Unstable softmax trace when exporting to onnx

[Enolerobotti]

🐛 Bug

The ONNX model with softmax which is exported with torch.onnx module predicts NaNs instead of expected numerical values. The actual output does not depend on mantissa. The actual result is unstable. Meaning that expected value 5.0390004e-32 actually may be converted either to NaN or to 5.0390004e-32. If for example 5.0390004e-32 was correctly converted to 5.0390004e-32 then it may happen that other element of the tensor, e.g. 1.0000000e+00, might be converted to NaN.

To Reproduce

Please consider the following code to reproduse

import torch.onnx
import torch.nn as nn
import onnxruntime as rt
import torch.nn.functional as F


class TraceCheck(nn.Module):
    def __init__(self):
        super(TraceCheck, self).__init__()

    def forward(self, x):
        return F.softmax(-100*x, -2)


if __name__ == '__main__':
    net = TraceCheck()
    x = torch.Tensor(torch.randn(5,5))
    pred = net.forward(x).detach().cpu().numpy()
    tmp_filename = 'test.onnx'
    torch.onnx.export(net, x, tmp_filename)
    sess = rt.InferenceSession(tmp_filename)
    pred_onnx = sess.run([sess.get_outputs()[-1].name], {sess.get_inputs()[-1].name: x.detach().cpu().numpy()})[0]
    print(pred)
    print(pred_onnx)

Example outputs are :
pred:

[[1.0000000e+00 0.0000000e+00 0.0000000e+00 1.0000000e+00 0.0000000e+00]
[0.0000000e+00 0.0000000e+00 0.0000000e+00 0.0000000e+00 0.0000000e+00]
[0.0000000e+00 6.1350977e-01 0.0000000e+00 0.0000000e+00 0.0000000e+00]
[0.0000000e+00 3.8649023e-01 8.0113404e-19 0.0000000e+00 1.0000000e+00]
[1.1149160e-09 5.0298510e-32 1.0000000e+00 0.0000000e+00 0.0000000e+00]]

pred_onnx:

[[ nan 0.0000000e+00 0.0000000e+00 nan 0.0000000e+00]
[0.0000000e+00 0.0000000e+00 0.0000000e+00 0.0000000e+00 0.0000000e+00]
[0.0000000e+00 6.1350977e-01 0.0000000e+00 0.0000000e+00 0.0000000e+00]
[0.0000000e+00 3.8649023e-01 nan 0.0000000e+00 nan]
[ nan 5.0298498e-32 nan 0.0000000e+00 0.0000000e+00]]

We see that pred != pred_onnx

Expected behavior

pred == pred_onnx

Environment

• PyTorch Version : 1.4.0
• OS: Windows 10 Pro 64 bit
• How you installed PyTorch (conda, pip, source): current repo (git clone ...)
• Build command you used (if compiling from source): setup.py build/install
• Python version: 3.6
• CUDA/cuDNN version: None
• GPU models and configuration: None
• Any other relevant information: opset 11

Additional context

cc @houseroad @spandantiwari @lara-hdr @BowenBao @neginraoof

[spandantiwari]

@Enolerobotti - Do you see anything wrong with the exported ONNX model? Let us know if you do. If not, then most likely this could be an implementation issue in ORT. Could you please report this in onnxruntime repo?

[Enolerobotti]

@spandantiwari yes, I see something wrong with the exported graph.

graph(%0 : Float(5, 5)):
  %1 : Float() = onnx::Constant[value={-100}]()
  %2 : Float(5, 5) = onnx::Mul(%0, %1)
  %3 : Tensor = onnx::Exp(%2)
  %4 : Tensor = onnx::ReduceSum[axes=[0]](%3)
  %5 : Float(5, 5) = onnx::Div(%3, %4) # c:\anaconda3\envs\python36b\lib\site-packages\torch\nn\functional.py:1231:0
  return (%5)

I guess I get NaN during the division Div(%3, %4). So, this bug is not come from the onnxruntime.

I've found a workaround. The workaround is replace negative axis -2 in the softmax with x.dim()-1. After the replacement I get correct results. Now the graph is

graph(%0 : Float(5, 5)):
  %1 : Float() = onnx::Constant[value={-100}]()
  %2 : Float(5, 5) = onnx::Mul(%0, %1)
  %3 : Float(5, 5) = onnx::Softmax[axis=1](%2) # c:\anaconda3\envs\python36b\lib\site-packages\torch\nn\functional.py:1231:0
  return (%3)

Unfortunately, this workaround is not useful for me because actually I try to export a third-party net

[p12tic]

The problem here is that the semantics of softmax are different in Pytorch and ONNX and Pytorch can use ONNX softmax only in certain situations. #37326 improves the situation.

cc @spandantiwari