Keras Softmax axis=1 bug

apivovarov · August 7, 2019, 6:37am

I noticed that Keras Relay Frontend uses axis=1 by default

dmlc/tvm/blob/master/python/tvm/relay/frontend/keras.py#L82


        act_type = keras_layer.activation.__name__
if act_type == 'linear':
    if isinstance(keras_layer, str):
        return inexpr
    alpha = keras_layer.alpha if hasattr(keras_layer, 'alpha') else 1.
    beta = keras_layer.beta if hasattr(keras_layer, 'beta') else 0.
    alpha = _expr.const(alpha, dtype='float32')
    beta = _expr.const(beta, dtype='float32')
    return _op.add(_op.multiply(inexpr, alpha), beta)
if act_type == 'softmax':
    return _op.nn.softmax(inexpr, axis=1)
if act_type == 'sigmoid':
    return _op.sigmoid(inexpr)
if act_type == 'tanh':
    return _op.tanh(inexpr)
if act_type == 'relu':
    return _op.nn.relu(inexpr)
if act_type == 'softplus':
    return _op.log(_op.add(_op.exp(inexpr), _expr.const(1., dtype='float32')))
if act_type == 'elu':
    alpha = keras_layer.alpha if hasattr(keras_layer, 'alpha') else 1.

github.com

dmlc/tvm/blob/master/nnvm/python/nnvm/frontend/keras.py#L67


    else:
        act_type = keras_layer.activation.__name__
if act_type == 'linear':
    if isinstance(keras_layer, str):
        return insym
    alpha = keras_layer.alpha if hasattr(keras_layer, "alpha") else 1
    beta = keras_layer.beta if hasattr(keras_layer, "beta") else 0
    return _sym.__add_scalar__(_sym.__mul_scalar__(insym, \
        scalar=alpha), scalar=beta)
if act_type == 'softmax':
    return _sym.softmax(insym, axis=1)
if act_type == 'sigmoid':
    return _sym.sigmoid(insym)
if act_type == 'tanh':
    return _sym.tanh(insym)
if act_type == 'relu':
    return _sym.relu(insym)
if act_type == 'softplus':
    return _sym.log(_sym.__add_scalar__(_sym.exp(insym), scalar=1))
if act_type == 'elu':
    alpha = keras_layer.alpha if hasattr(keras_layer, "alpha") else 1

But I think it should be axis=-1
https://www.tensorflow.org/api_docs/python/tf/nn/softmax
https://keras.io/layers/advanced-activations/#softmax

github.com

dmlc/tvm/blob/ee74d00ea88ae20d0a3fb943eebb0f531fa51f34/topi/python/topi/nn/softmax.py#L23


# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
# KIND, either express or implied.  See the License for the
# specific language governing permissions and limitations
# under the License.
# pylint: disable=invalid-name
"""TVM operator for softmax and log_softmax compute."""
from __future__ import absolute_import
import tvm


@tvm.tag_scope(tag='softmax_output')
def softmax(x, axis=-1):
"""Perform softmax activation on the data


Parameters
----------
data : tvm.Tensor
    can be any dimension


axis : int
    channel axis