Cannot export custom operator to onnx

username: coincheung
cooked: <p>Hi,</p>
<p>I followed docs to export my self-defined operator to onnx, code like this:</p>
<pre><code class="lang-auto">
import torch
import torchvision
from torch.onnx.symbolic_helper import parse_args
from torch.onnx import register_custom_op_symbolic

torch.set_grad_enabled(False)

@parse_args('v', 'v', 'f', 'i')
def symbolic_taylor_softmax(g, inten, dim, n, use_log):
    return g.op("TaylorSoftmaxV3", inten, dim, n, use_log)

register_custom_op_symbolic('custom_ops::taylor_softmax', symbolic_taylor_softmax, 9)

class Wrapper(torch.nn.Module):

    def __init__(self):
        super(Wrapper, self).__init__()
        net = torchvision.models.resnet50(pretrained=False)
        net.eval()
        self.net = net

    def forward(self, x):
        scores = self.net(x)[0]
        scores = torch.ops.custom_ops.taylor_softmax(scores, 1, 4, True)
        return scores

net = Wrapper()
net.eval()
dummy_input = torch.randn(1, 3, 224, 224)
input_names = ['input_image']
output_names = ['preds',]

torch.onnx.export(net, dummy_input, 'model.onnx',
    input_names=input_names, output_names=output_names,
    verbose=False, opset_version=11,
    custom_opsets={'custom_domain': 2})
</code></pre>
<p>The error message is like this:</p>
<pre><code class="lang-auto">Traceback (most recent call last):
  File "try.py", line 95, in &lt;module&gt;
    custom_opsets={'custom_domain': 2})
  File "/root/miniconda3/lib/python3.7/site-packages/torch/onnx/__init__.py", line 276, in export
    custom_opsets, enable_onnx_checker, use_external_data_format)
  File "/root/miniconda3/lib/python3.7/site-packages/torch/onnx/utils.py", line 94, in export
    use_external_data_format=use_external_data_format)
  File "/root/miniconda3/lib/python3.7/site-packages/torch/onnx/utils.py", line 698, in _export
    dynamic_axes=dynamic_axes)
  File "/root/miniconda3/lib/python3.7/site-packages/torch/onnx/utils.py", line 456, in _model_to_graph
    use_new_jit_passes)
  File "/root/miniconda3/lib/python3.7/site-packages/torch/onnx/utils.py", line 417, in _create_jit_graph
    graph, torch_out = _trace_and_get_graph_from_model(model, args)
  File "/root/miniconda3/lib/python3.7/site-packages/torch/onnx/utils.py", line 377, in _trace_and_get_graph_from_model
    torch.jit._get_trace_graph(model, args, strict=False, _force_outplace=False, _return_inputs_states=True)
  File "/root/miniconda3/lib/python3.7/site-packages/torch/jit/_trace.py", line 1139, in _get_trace_graph
    outs = ONNXTracedModule(f, strict, _force_outplace, return_inputs, _return_inputs_states)(*args, **kwargs)
  File "/root/miniconda3/lib/python3.7/site-packages/torch/nn/modules/module.py", line 889, in _call_impl
    result = self.forward(*input, **kwargs)
  File "/root/miniconda3/lib/python3.7/site-packages/torch/jit/_trace.py", line 130, in forward
    self._force_outplace,
  File "/root/miniconda3/lib/python3.7/site-packages/torch/jit/_trace.py", line 116, in wrapper
    outs.append(self.inner(*trace_inputs))
  File "/root/miniconda3/lib/python3.7/site-packages/torch/nn/modules/module.py", line 887, in _call_impl
    result = self._slow_forward(*input, **kwargs)
  File "/root/miniconda3/lib/python3.7/site-packages/torch/nn/modules/module.py", line 860, in _slow_forward
    result = self.forward(*input, **kwargs)
  File "try.py", line 83, in forward
    scores = torch.ops.custom_ops.taylor_softmax(scores, 1, 4, True)
  File "/root/miniconda3/lib/python3.7/site-packages/torch/_ops.py", line 60, in __getattr__
    op = torch._C._jit_get_operation(qualified_op_name)
RuntimeError: No such operator custom_ops::taylor_softmax
</code></pre>
<p>How could I make it work please ?</p>