如何使用 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