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