Converting log_softmax layer into ONNX format

I want to convert a network into ONNX format, and bumped into this problem.

The conversion of log_softmax layer is not supported. This code:

import torch
import torch.nn
import torch.nn.functional

img_size, num_inputs = 64, 3

class SimpleModel(torch.nn.Module):
    def __init__(self):
        super(SimpleModel, self).__init__()
        self.layer = torch.nn.AvgPool2d(int(64/4))
    def forward(self, x):
        return torch.nn.functional.log_softmax(x, dim=1)

if __name__ == "__main__":
    print("Torch version:", torch.__version__)
    model = SimpleModel()
    dummy_input = torch.randn(10, num_inputs, img_size, img_size, device='cpu')
    torch.onnx.export(model, dummy_input, "simple.onnx", verbose=True)

generates the following error output:

Torch version: 1.0.1.post2

Traceback (most recent call last):
  File "", line 18, in <module>
    torch.onnx.export(model, dummy_input, "simple.onnx", verbose=True)
  File "/usr3/ams/radio/rantonio/.local/lib/python3.5/site-packages/torch/onnx/", line 27, in export
    return utils.export(*args, **kwargs)
  File "/usr3/ams/radio/rantonio/.local/lib/python3.5/site-packages/torch/onnx/", line 104, in export
  File "/usr3/ams/radio/rantonio/.local/lib/python3.5/site-packages/torch/onnx/", line 287, in _export
    proto, export_map = graph._export_onnx(params, _onnx_opset_version, defer_weight_export, operator_export_type)
RuntimeError: ONNX export failed: Couldn't export operator aten::log_softmax

Is there a particular reason why porting log_softmax to ONNX is not supported by Pytorch?



Did you found a solution ?

The relevant bug report is probably this one: