Hi all,
I’ve run into some behavior that I can’t explain. Below is a snippet of code that was lifted from this discussion: https://github.com/pytorch/pytorch/issues/24235.
What I’ve found is that if bias=False in testSubMod then the compilation runs with no problem and outputs test.onnx as a file but if bias=True then the compilation fails with the stack trace shown below.
I’m pretty new to working with ONNX and TorchScript so any insights would be really appreciated!
Thanks,
Zach
import torch
import onnx
class testSubMod(torch.nn.Module):
def __init__(self, rnn_dims=32):
super().__init__()
self.lin = torch.nn.Linear(32, 32, bias=True) # < This, this bit here!
def forward(self, out):
for _ in torch.arange(8):
out = self.lin(out)
return out
class test(torch.nn.Module):
def __init__(self, rnn_dims=32):
super().__init__()
self.submod = torch.jit.script(testSubMod())
def forward(self, x):
out = torch.ones(
[
x.size(0),
x.size(1)
],
dtype=x.dtype,
device=x.device
)
return self.submod(out)
if __name__=='__main__':
model = test()
model = torch.jit.script(model)
input_data = torch.ones((32, 32, 32)).float()
output = model(input_data)
torch.onnx.export(model, input_data,
'test.onnx', example_outputs=output)
onnx_model = onnx.load("test.onnx")
print(onnx.helper.printable_graph(onnx_model.graph))
/home/USERNAME/anaconda3/envs/rain_man/lib/python3.7/site-packages/torch/onnx/utils.py:617: UserWarning: ONNX export failed on ATen operator dim because torch.onnx.symbolic_opset9.dim does not exist
.format(op_name, opset_version, op_name))
Traceback (most recent call last):
File "compile.py", line 38, in <module>
'test.onnx', example_outputs=output)
File "/home/USERNAME/anaconda3/envs/rain_man/lib/python3.7/site-packages/torch/onnx/__init__.py", line 143, in export
strip_doc_string, dynamic_axes, keep_initializers_as_inputs)
File "/home/USERNAME/anaconda3/envs/rain_man/lib/python3.7/site-packages/torch/onnx/utils.py", line 66, in export
dynamic_axes=dynamic_axes, keep_initializers_as_inputs=keep_initializers_as_inputs)
File "/home/USERNAME/anaconda3/envs/rain_man/lib/python3.7/site-packages/torch/onnx/utils.py", line 382, in _export
fixed_batch_size=fixed_batch_size)
File "/home/USERNAME/anaconda3/envs/rain_man/lib/python3.7/site-packages/torch/onnx/utils.py", line 262, in _model_to_graph
fixed_batch_size=fixed_batch_size)
File "/home/USERNAME/anaconda3/envs/rain_man/lib/python3.7/site-packages/torch/onnx/utils.py", line 132, in _optimize_graph
graph = torch._C._jit_pass_onnx(graph, operator_export_type)
File "/home/USERNAME/anaconda3/envs/rain_man/lib/python3.7/site-packages/torch/onnx/__init__.py", line 174, in _run_symbolic_function
return utils._run_symbolic_function(*args, **kwargs)
File "/home/USERNAME/anaconda3/envs/rain_man/lib/python3.7/site-packages/torch/onnx/utils.py", line 644, in _run_symbolic_function
torch._C._jit_pass_onnx_block(b, new_block, operator_export_type, env)
File "/home/USERNAME/anaconda3/envs/rain_man/lib/python3.7/site-packages/torch/onnx/__init__.py", line 174, in _run_symbolic_function
return utils._run_symbolic_function(*args, **kwargs)
File "/home/USERNAME/anaconda3/envs/rain_man/lib/python3.7/site-packages/torch/onnx/utils.py", line 618, in _run_symbolic_function
op_fn = sym_registry.get_registered_op(op_name, '', opset_version)
File "/home/USERNAME/anaconda3/envs/rain_man/lib/python3.7/site-packages/torch/onnx/symbolic_registry.py", line 91, in get_registered_op
return _registry[(domain, version)][opname]
KeyError: 'dim'