import torch
pretrained = "pretrained/states_pt_places2.pth"
generator_state_dict = torch.load(pretrained)['G']
if 'stage1.conv1.conv.weight' in generator_state_dict.keys():
from model.networks import Generator
else:
from model.networks_tf import Generator
# set up network
generator = Generator(cnum_in=5, cnum=48, return_flow=False)
generator.load_state_dict(generator_state_dict, strict=True)
img = torch.rand([5, 512, 512]).unsqueeze(0).cpu()
mask = torch.empty(img.shape[0], 1, img.shape[2], img.shape[3]).cpu()
data = (img, mask)
generator.cpu().eval()
dynamic_axes = None
dynamic_export = True
if dynamic_export:
dynamic_axes = {
'input': {
0: 'batch',
2: 'height',
3: 'width'
},
'output': {
0: 'batch',
2: 'height',
3: 'width'
}
}
with torch.no_grad():
torch.onnx.export(
generator,
data,
"output.onnx",
input_names=['input'],
output_names=['output'],
export_params=True,
keep_initializers_as_inputs=False,
verbose=False,
opset_version=12,
dynamic_axes=dynamic_axes)
print(f'Successfully exported ONNX model')
Traceback (most recent call last):
File "deepfillv2-pytorch/export2onnx.py", line 65, in <module>
torch.onnx.export(
File "/miniconda3/envs/py39/lib/python3.9/site-packages/torch/onnx/__init__.py", line 316, in export
return utils.export(model, args, f, export_params, verbose, training,
File "/miniconda3/envs/py39/lib/python3.9/site-packages/torch/onnx/utils.py", line 107, in export
_export(model, args, f, export_params, verbose, training, input_names, output_names,
File "/miniconda3/envs/py39/lib/python3.9/site-packages/torch/onnx/utils.py", line 724, in _export
_model_to_graph(model, args, verbose, input_names,
File "/miniconda3/envs/py39/lib/python3.9/site-packages/torch/onnx/utils.py", line 497, in _model_to_graph
graph = _optimize_graph(graph, operator_export_type,
File "/miniconda3/envs/py39/lib/python3.9/site-packages/torch/onnx/utils.py", line 216, in _optimize_graph
graph = torch._C._jit_pass_onnx(graph, operator_export_type)
File "/miniconda3/envs/py39/lib/python3.9/site-packages/torch/onnx/__init__.py", line 373, in _run_symbolic_function
return utils._run_symbolic_function(*args, **kwargs)
File "/miniconda3/envs/py39/lib/python3.9/site-packages/torch/onnx/utils.py", line 1032, in _run_symbolic_function
return symbolic_fn(g, *inputs, **attrs)
File "/miniconda3/envs/py39/lib/python3.9/site-packages/torch/onnx/symbolic_helper.py", line 172, in wrapper
return fn(g, *args, **kwargs)
File "/miniconda3/envs/py39/lib/python3.9/site-packages/torch/onnx/symbolic_opset9.py", line 1281, in _convolution
raise RuntimeError("Unsupported: ONNX export of convolution for kernel "
RuntimeError: Unsupported: ONNX export of convolution for kernel of unknown shape.
I've google it but they said that because there is an operation in pytorch that onnx doesn't support. But I could not investigate what operation.
Could you help me please? Thank you