Hi,
I followed docs to export my self-defined operator to onnx, code like this:
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})
The error message is like this:
Traceback (most recent call last):
File "try.py", line 95, in <module>
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
How could I make it work please ?