I wonder if it is possible to register a custom Op of ONNX from tvm.
- I’ve created a opset for this op.
- So the onnx file is valid.
The testing code is like:
import onnx
from onnx import helper, TensorProto
from tvm.relax.frontend.onnx import from_onnx
from tvm import relax
op_name = "my_custom_op"
op_set_name = "ai.onnx.contrib"
def create_custom_opset():
opset = helper.make_opsetid(op_set_name, 2)
return opset
def create_custom_model():
shape = [32, 32]
node = helper.make_node(op_name, ["x"], ["y"], domain="ai.onnx.contrib")
graph = helper.make_graph(
[node],
"custom_op_graph",
inputs=[helper.make_tensor_value_info("x", TensorProto.FLOAT, shape)],
outputs=[helper.make_tensor_value_info("y", TensorProto.FLOAT, shape)],
)
model = helper.make_model(graph, producer_name="custom_op_model")
model.opset_import.append(create_custom_opset())
return model
if __name__ == "__main__":
model = create_custom_model()
# try to check it using onnx.check, it should success
try:
onnx.checker.check_model(model)
print("Model is valid.")
except Exception as exception:
print("Model is not valid: ", str(exception))
ir_mod = from_onnx(model, keep_params_in_input=True)
ir_mod, params = relax.frontend.detach_params(ir_mod)
ir_mod.show()
It has a new operator called my_custom_op, when I use relax.frontend.onnx.from_onnx to import this model, it won’t work.
The error is like:
Model is valid.
Traceback (most recent call last):
File "/home/xxx/xxx/sandbox/tvm_tunning/sdk/frontend/onnx_frontend.py", line 41, in <module>
ir_mod = from_onnx(model, keep_params_in_input=True)
File "/home/xxx/xxx/sandbox/tvm/python/tvm/relax/frontend/onnx/onnx_frontend.py", line 2173, in from_onnx
return g.from_onnx(graph, opset)
File "/home/xxx/xxx/sandbox/tvm/python/tvm/relax/frontend/onnx/onnx_frontend.py", line 1829, in from_onnx
self._check_for_unsupported_ops(graph)
File "/home/xxx/xxx/sandbox/tvm/python/tvm/relax/frontend/onnx/onnx_frontend.py", line 1949, in _check_for_unsupported_ops
raise tvm.error.OpNotImplemented(msg)
tvm.error.OpNotImplemented: The following operators are not supported for frontend ONNX: my_custom_op
Currently, I know that ONNX Runtime supports registering new op, Is there a way to register this op in TVM? Any help will be appreciate. Thanks advance!