onnx自定义op注册

Question

我正在寻找与ONNX/ONNX运行时中的keras自定义层类似的特性。我理解解决这个问题的方法是在onnx中实现一个自定义操作符以便进行实验。文档似乎将C++中的实现作为一个共享库，并在python中使用它。https://onnxruntime.ai/docs/reference/operators/add-custom-op.html

是否有一种方法可以在python中为onnx定义自定义op，仅用于实验目的，并用于推断？我试过这样做，但给出了‘错误: PyOp不是一个注册函数/op’https://onnxruntime.ai/docs/reference/operators/custom-python-operator.html

Python代码：

import onnx
import onnxruntime as ort

A = onnx.helper.make_tensor_value_info('A', onnx.TensorProto.FLOAT, [4])
B = onnx.helper.make_tensor_value_info('B', onnx.TensorProto.FLOAT, [4])
C = onnx.helper.make_tensor_value_info('C', onnx.TensorProto.FLOAT, [4])
D = onnx.helper.make_tensor_value_info('D', onnx.TensorProto.FLOAT, [4])
E = onnx.helper.make_tensor_value_info('E', onnx.TensorProto.FLOAT, [4])
F = onnx.helper.make_tensor_value_info('F', onnx.TensorProto.FLOAT, [4])

ad1_node = onnx.helper.make_node('Add', ['A', 'B'], ['S'])

mul_node = onnx.helper.make_node('Mul', ['C','D'], ['P'])

ad2_node = onnx.helper.make_node('Add', ['S', 'P'], ['H'])

py1_node = onnx.helper.make_node(op_type = 'PyOp', #required, must be 'PyOp'
                            inputs = ['H'], #required
                            outputs = ['F'], #required
                            domain = 'pyopadd_2', #required, must be unique
                            input_types = [onnx.TensorProto.FLOAT], #required
                            output_types = [onnx.TensorProto.FLOAT], #required
                            module = 'mymodule', #required
                            class_name = 'Add_2', #required
                            compute = 'compute') #optional, 'compute' by default

graph = onnx.helper.make_graph([ad1_node,mul_node,ad2_node, py1_node], 'multi_pyop_graph', [A,B,C,D], [F])
model = onnx.helper.make_model(graph, opset_imports=[onnx.helper.make_opsetid('pyopadd_2', 1)], producer_name = 'pyop_model')
onnx.save(model, './modeltemp.onnx')

ort_session = ort.InferenceSession('./modeltemp.onnx')
ort_output = ort_session.run(["F"], {'A':[1,2,3,4], 'B':[1,1,1,1], 'C':[2,2,2,2], 'D':[3,3,3,3]})
print(ort_output)

mymodule.py

class Add_2:
    def compute(self, S):
        return S+2

Answer 1

这不是您使用PyOp的方式。首先:您需要实现在python中尝试使用的操作符。第二:您需要注册在ONNXRuntime会话中实现的操作符。第三:运行包含自定义操作的模型的推断。

这里可以找到一个例子，tutorial.ipynb，您必须查看名为“在python中实现op”的部分。警告！这个笔记本中的PyOp叫做PyCustomOpDef，如果您直接从ONNX模型开始，您不需要担心这个笔记本中的任何其他东西。只需要将"ai.onnx.contrib"包含在模型的选项集中，并在节点域中设置相同的域"ai.onnx.contrib"。

# for the model
DOMAIN = "ai.onnx.contrib"
VERSION = 1 # try 2 or 3, I had some issues with the versioning
new_opset = onnx.helper.make_opsetid(DOMAIN, VERSION)
loaded_model.opset_import.append(new_opset)

# for the node, like in your code 
domain = 'ai.onnx.contrib', #required, must be unique

Answer 2

正如自定义python-运算符明确指出的那样，您必须自己构建with运行时：

--config Release --enable_language_interop_ops --build_wheel

该功能不随of运行时的预构建版本而来。