3
$\begingroup$

We've been looking at softmax results produced by different frameworks (TF, PyTorch, Caffe2, Glow and ONNX runtime) and were surprised to find that the results differ between the frameworks.

From the documents for each framework it is clear that they do handle softmax differently. PyTorch and Tensorflow produce similar results that fall in line with what I would expect. ONNX and Caffe2 results are very different in terms of the actual probabilities while the order of the numerically sorted probabilities appear to be consistent.

Hopefully it isn't just poor search skills but I have been unsuccessful in finding any reference that explains why Caffe2 and ONNX define softmax the way they do.

The ONNX softmax operator is defined as follows (which appears to come from Caffe2):

The operator computes the softmax (normalized exponential) values for each layer in the batch of the given input.

The input does not need to explicitly be a 2D vector; rather, it will be coerced into one. For an arbitrary n-dimensional tensor input \in [a_0, a_1, ..., a_{k-1}, a_k, ..., a_{n-1}] and k is the axis provided, then input will be coerced into a 2-dimensional tensor with dimensions [a_0 * ... * a_{k-1}, a_k * ... * a_{n-1}]. For the default case where axis=1, this means the input tensor will be coerced into a 2D tensor of dimensions [a_0, a_1 * ... * a_{n-1}], where a_0 is often the batch size. In this situation, we must have a_0 = N and a_1 * ... * a_{n-1} = D. Each of these dimensions must be matched correctly, or else the operator will throw errors. The output tensor has the same shape and contains the softmax values of the corresponding input.

While the PyTorch softmax operator defines it as I would expect:

enter image description here

The results we see are as follows:

Input tensor to softmax:

  tensor([[[1., 1.],[2., 2.]],[[3., 3.],[4., 4.]]])

TensorFlow:

    2: [[[0.5 0.5] [0.5 0.5]] [[0.5 0.5] [0.5 0.5]]]
    1: [[[0.26894143 0.26894143] [0.7310586  0.7310586 ]] [[0.26894143 0.26894143] [0.7310586  0.7310586 ]]]
    0: [[[0.11920291 0.11920291] [0.11920291 0.11920291]] [[0.880797   0.880797  ] [0.880797   0.880797  ]]]

PyTorch:

    2: [[[0.5000, 0.5000], [0.5000, 0.5000]], [[0.5000, 0.5000], [0.5000, 0.5000]]]
    1: [[[0.2689, 0.2689], [0.7311, 0.7311]], [[0.2689, 0.2689], [0.7311, 0.7311]]]
    0: [[[0.1192, 0.1192], [0.1192, 0.1192]], [[0.8808, 0.8808], [0.8808, 0.8808]]]

Caffe2:

    2: [[[0.5 0.5]  [0.5 0.5]] [[0.5 0.5] [0.5 0.5]]]
    1: [[[0.13447072 0.13447072] [0.3655293  0.3655293 ]] [[0.13447072 0.13447072] [0.3655293  0.3655293 ]]]
    0: [[[0.0160293  0.0160293 ] [0.04357216 0.04357216]] [[0.1184414  0.1184414 ] [0.3219571  0.3219571 ]]]

ONNX Runtime:

     2: [[[0.5, 0.5], [0.5, 0.5]], [[0.5, 0.5], [0.5, 0.5]]]
     1: [[[0.13447072, 0.13447072], [0.3655293 , 0.3655293 ]], [[0.13447072, 0.13447072], [0.3655293 , 0.3655293 ]]]
     0: [[[0.0160293 , 0.0160293 ], [0.04357216, 0.04357216]], [[0.1184414 , 0.1184414 ], [0.3219571 , 0.3219571 ]]]

Can someone explain the rationale behind the difference?

$\endgroup$
5
$\begingroup$

The ONNX documentation you wrote describes the reshaping that is done by their softmax implementation: an input tensor is always reshaped to 2 dimensions before applying softmax along the second axis (I can't explain why they choose to do it this way). It is roughly equivalent to the following TensorFlow code:

import tensorflow as tf


def onnx_softmax(tensor, axis: int) -> tf.Tensor:
    tensor = tf.convert_to_tensor(tensor)
    shape = tf.shape(tensor)
    new_shape_0 = tf.math.reduce_prod(shape[:axis])
    new_shape_1 = tf.math.reduce_prod(shape[axis:])
    new_shape = tf.stack([new_shape_0, new_shape_1])
    tensor_2d = tf.reshape(tensor, shape=new_shape)
    softmax = tf.nn.softmax(tensor_2d, axis=-1)
    return tf.reshape(softmax, shape=shape)

Examples (using TensorFlow 2 for eager execution):

  • Softmax along axis 2:

    tensor = [[[1., 1.], [2., 2.]], [[3., 3.], [4., 4.]]]
    print(onnx_softmax(tensor, axis=2))
    

    Output:

    tf.Tensor(
    [[[0.5 0.5]
      [0.5 0.5]]
    
     [[0.5 0.5]
      [0.5 0.5]]], shape=(2, 2, 2), dtype=float32)
    
  • Softmax along axis 1:

    tensor = [[[1., 1.], [2., 2.]], [[3., 3.], [4., 4.]]]
    print(onnx_softmax(tensor, axis=1))
    

    Output:

    tf.Tensor(
    [[[0.13447072 0.13447072]
      [0.3655293  0.3655293 ]]
    
     [[0.13447072 0.13447072]
      [0.3655293  0.3655293 ]]], shape=(2, 2, 2), dtype=float32)
    
  • Softmax along axis 0:

    tensor = [[[1., 1.], [2., 2.]], [[3., 3.], [4., 4.]]]
    print(onnx_softmax(tensor, axis=0))
    

    Output:

    tf.Tensor(
    [[[0.0160293  0.0160293 ]
      [0.04357216 0.04357216]]
    
     [[0.1184414  0.1184414 ]
      [0.3219571  0.3219571 ]]], shape=(2, 2, 2), dtype=float32)
    

As you can see, this agrees with your ONNX/Caffe2 printouts.

$\endgroup$
  • 1
    $\begingroup$ FWIW, there are threads about the softmax implementation on the ONNX Github and there is some discussion around this design choice. After skimming it, I'm still not sure why they made this decision. github.com/onnx/onnx/issues/2289 $\endgroup$ – Sycorax says Reinstate Monica Jan 29 at 19:18
  • $\begingroup$ Thank you. I was hoping someone would have some insight into the reason behind the reshaping but I appreciate your response. $\endgroup$ – Mike N. Feb 7 at 21:37

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service, privacy policy and cookie policy

Not the answer you're looking for? Browse other questions tagged or ask your own question.