When using torch.onnx.export(), errors occurs: AssertionError: Only output_size=[1, 1] is supported

1

Environments:
ubuntu 16.04
cuda8.0
python: 3.5.2
pytorch: 1.0.1
torchvision: 0.2.2

Hi guys, when I use torch.onnx.export, I get some problems. The code is just from the " Example: End-to-end AlexNet from PyTorch to Caffe" (https://pytorch.org/docs/stable/onnx.html?highlight=export#torch.onnx.export)

import torch
import torchvision

dummy_input = torch.randn(10, 3, 224, 224, device='cuda')
model = torchvision.models.alexnet(pretrained=True).cuda()
input_names = [ "actual_input_1" ] + [ "learned_%d" % i for i in range(16) ]
output_names = [ "output1" ]

torch.onnx.export(model, dummy_input, "alexnet.onnx", verbose=True, input_names=input_names, output_names=output_names)

Here is the error:

Traceback (most recent call last):
File “pytorch_to_onnx.py”, line 10, in
torch.onnx.export(model, dummy_input, “alexnet.onnx”, verbose=True)#, input_names=input_names, output_names=output_names)
File “/usr/local/lib/python3.5/dist-packages/torch/onnx/init.py”, line 27, in export
return utils.export(*args, **kwargs)
File “/usr/local/lib/python3.5/dist-packages/torch/onnx/utils.py”, line 104, in export
operator_export_type=operator_export_type)
File “/usr/local/lib/python3.5/dist-packages/torch/onnx/utils.py”, line 281, in _export
example_outputs, propagate)
File “/usr/local/lib/python3.5/dist-packages/torch/onnx/utils.py”, line 227, in _model_to_graph
graph = _optimize_graph(graph, operator_export_type)
File “/usr/local/lib/python3.5/dist-packages/torch/onnx/utils.py”, line 155, in _optimize_graph
graph = torch._C._jit_pass_onnx(graph, operator_export_type)
File “/usr/local/lib/python3.5/dist-packages/torch/onnx/init.py”, line 52, in _run_symbolic_function
return utils._run_symbolic_function(*args, **kwargs)
File “/usr/local/lib/python3.5/dist-packages/torch/onnx/utils.py”, line 504, in _run_symbolic_function
return fn(g, *inputs, **attrs)
File “/usr/local/lib/python3.5/dist-packages/torch/onnx/symbolic.py”, line 89, in wrapper
return fn(g, *args)
File “/usr/local/lib/python3.5/dist-packages/torch/onnx/symbolic.py”, line 600, in adaptive_avg_pool2d
assert output_size == [1, 1], “Only output_size=[1, 1] is supported”
AssertionError: Only output_size=[1, 1] is supported

By the way, the documents about torch.onnx.export() are vague. I cannot understand what the arguments input_names and output_names are for?

Any cues would be appreciated!

2

Finally, I fix this problem by downgrading the torchvision from 0.2.2 to 0.2.1.
Cannot believe that there are such huge differences

4 Likes
3

Thans a lot !!!

4

Thank you very much!

5

The original AlexNet definition in v0.2.1 does not have AdaptiveAvgPool2d, but v0.2.2 does. As stated in the error log, the error is mainly caused by AdaptiveAvgPool2d having output size other than [1,1].

I had the same error as you, solved by changing it to fixed size pooling (nn.AvgPool2d).

2 Likes
6

I tried it ， it works. in"alexnet.py" row 34:self.avgpool = nn.AdaptiveAvgPool2d((6, 6)) to self.avgpool = nn.AvgPool2d((1, 1))