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 "test.py", 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/__init__.py", line 27, in export
return utils.export(*args, **kwargs)
File "/usr3/ams/radio/rantonio/.local/lib/python3.5/site-packages/torch/onnx/utils.py", line 104, in export
operator_export_type=operator_export_type)
File "/usr3/ams/radio/rantonio/.local/lib/python3.5/site-packages/torch/onnx/utils.py", 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?
Thanks,