如何使用 Python 运算符 (PyOp)

弃用说明：此功能已弃用，不再支持

Python 运算符提供了使用 ONNX Runtime 在 ONNX 图的单个节点内轻松调用任何自定义 Python 代码的能力。当模型需要 ONNX 和 ONNX Runtime 尚未官方支持的运算符时，这对于更快速的实验很有用，特别是如果所需功能已有 Python 实现。在生产场景中使用时应谨慎，并应事先考虑所有安全或其他风险。

设计概览

此功能位于 language_interop_ops 中。

以下是调用序列图

onnxruntime                        python capi                         script
     |                                  |                                 |
     | ------------------------------>  |                                 |
     |       call with tensor(s)        | ------------------------------> |
     |                                  |         call with numpy(s)      | 
     |                                  |                                 | compute
     |                                  | <------------------------------ |
     | <------------------------------  |           return numpys(s)      |
     |         return tensor(s)         |                                 |

如何使用

步骤 1

使用 --config Release --enable_language_interop_ops --build_wheel 构建 onnxruntime，并 pip 安装最新的 wheel 文件。

步骤 2

创建包含 Python 运算符节点的 onnx 模型

ad1_node = helper.make_node('Add', ['A','B'], ['S'])
mul_node = helper.make_node('Mul', ['C','D'], ['P'])
py1_node = helper.make_node(op_type = 'PyOp', #required, must be 'PyOp'
                            inputs = ['S','P'], #required
                            outputs = ['L','M','N'], #required
                            domain = 'pyopmulti_1', #required, must be unique
                            input_types = [TensorProto.FLOAT, TensorProto.FLOAT], #required
                            output_types = [TensorProto.FLOAT, TensorProto.FLOAT, TensorProto.FLOAT], #required
                            module = 'mymodule', #required
                            class_name = 'Multi_1', #required
                            compute = 'compute', #optional, 'compute' by default
                            W1 = '5', W2 = '7', W3 = '9') #optional, must all be strings
ad2_node = helper.make_node('Add', ['L','M'], ['H'])
py2_node = helper.make_node('PyOp',['H','N','E'],['O','W'], domain = 'pyopmulti_2',
                            input_types = [TensorProto.FLOAT, TensorProto.FLOAT, TensorProto.FLOAT],
                            output_types = [TensorProto.FLOAT, TensorProto.FLOAT],
                            module = 'mymodule', class_name = 'Multi_2')
sub_node = helper.make_node('Sub', ['O','W'], ['F'])
graph = helper.make_graph([ad1_node,mul_node,py1_node,ad2_node,py2_node,sub_node], 'multi_pyop_graph', [A,B,C,D,E], [F])
model = helper.make_model(graph, producer_name = 'pyop_model')
onnx.save(model, './model.onnx')

步骤 3

实现 mymodule.py

class Multi_1:
    def __init__(self, W1, W2, W3):
        self.W1 = int(W1)
        self.W2 = int(W2)
        self.W3 = int(W3)
    def compute(self, S, P):
        ret = S + P
        return ret + self.W1, ret + self.W2, ret + self.W3
class Multi_2:
    def compute(self, *kwargs):
        return sum(kwargs[0:-1]), sum(kwargs[1:])

步骤 4

将 mymodule.py 复制到 Python sys.path 中，然后使用 onnxruntime python API 运行模型。在 Windows 上，请事先设置 PYTHONHOME。它应指向 Python 安装目录，例如 C:\Python37 或如果使用 conda，则指向 C:\ProgramData\Anaconda3\envs\myconda1。

支持的数据类型

TensorProto.BOOL
TensorProto.UINT8
TensorProto.UINT16
TensorProto.UINT32
TensorProto.INT16
TensorProto.INT32
TensorProto.FLOAT
TensorProto.DOUBLE

限制

推理和编译环境必须安装相同版本的 Python。
在 Windows 上，使用 --config Debug 已知存在问题。如果需要调试符号，请使用 --config RelWithDebInfo 进行构建。
由于 Python C API 的限制，多线程被禁用，因此 Python 运算符将按顺序运行。

测试覆盖率

此运算符已在多个平台上进行测试，包括使用或不使用 conda 的情况

平台	Python 3.5	Python 3.6	Python 3.7
Windows	(conda) 通过	(conda) 通过	通过
Linux	(conda) 通过	(conda) 通过	通过
Mac	(conda) 通过	(conda) 通过	(conda) 通过

示例

开发者可以在模型转换期间使用 PyOp 来处理缺失的运算符

import os
import numpy as np
from onnx import *
from skl2onnx import convert_sklearn
from skl2onnx.common.data_types import FloatTensorType
from skl2onnx.common.utils import check_input_and_output_numbers

X = np.array([[1, 1], [2, 1], [3, 1.2], [4, 1], [5, 0.8], [6, 1]],dtype=np.single)
nmf = NMF(n_components=2, init='random', random_state=0)
W = np.array(nmf.fit_transform(X), dtype=np.single)

def calculate_sklearn_nmf_output_shapes(operator):
    check_input_and_output_numbers(operator, output_count_range=1, input_count_range=1)
    operator.outputs[0].type.shape = operator.inputs[0].type.shape

def convert_nmf(scope, operator, container):
    ws = [str(w) for w in W.flatten()]
    attrs = {'W':'|'.join(ws)}
    container.add_node(op_type='PyOp', name='nmf', inputs=['X'], outputs=['variable'],
                       op_version=10, op_domain='MyDomain', module='mymodule', class_name='MyNmf',
                       input_types=[TensorProto.FLOAT], output_types=[TensorProto.FLOAT], **attrs)

custom_shape_calculators = {type(nmf): calculate_sklearn_nmf_output_shapes}
custom_conversion_functions = {type(nmf): convert_nmf}
initial_types = [('X', FloatTensorType([6,2]))]
onx = convert_sklearn(nmf, '', initial_types, '', None, custom_conversion_functions, custom_shape_calculators)
with th open("model.onnx", "wb") as f:
    f.write(onx.SerializeToString())

mymodule.py

import numpy as np
class MyNmf:
    def __init__(self,W):
        A = []
        for w in W.split('|'):
            A.append(float(w))
        self.__W = np.array(A,dtype=np.single).reshape(6,2)
    def compute(self,X):
        return self.__W