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 3.
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'