使用 Python 函式建立自定義運算元

自定義運算元是 ONNX Runtime 中的一項強大功能,它允許使用者透過實現自己的運算元來擴充套件執行時的功能,以執行標準 ONNX 運算元集中不提供的特定操作。

本文件將介紹如何使用 Python 函式建立自定義運算元,並將其整合到 ONNX Runtime 中進行推理。

步驟 1:為自定義運算元定義 Python 函式

首先定義將作為自定義運算元實現的 Python 函式。確保該函式與您自定義運算元所需的輸入和輸出張量形狀相容。Python 裝飾器 @onnx_op 會將該函式轉換為自定義運算元實現。以下是我們為分詞器建立函式示例:

@onnx_op(op_type="GPT2Tokenizer",
            inputs=[PyCustomOpDef.dt_string],
            outputs=[PyCustomOpDef.dt_int64, PyCustomOpDef.dt_int64],
            attrs={"padding_length": PyCustomOpDef.dt_int64})
def bpe_tokenizer(s, **kwargs):
    padding_length = kwargs["padding_length"]
    input_ids, attention_mask = cls.tokenizer.tokenizer_sentence([s[0]], padding_length)
    return input_ids, attention_mask

因為 ONNX Runtime 在載入模型時需要自定義運算元的 schema(模式),請透過 onnx_op 引數指定它們。如果 ONNX 節點有屬性,也需要‘attrs’引數,它可以是一個將名稱對映到型別的字典,如果所有型別都只是字串,也可以是一個列表。

步驟 2:建立包含自定義運算元的 ONNX 模型

自定義運算元註冊到 ONNX Runtime 後,您可以建立利用它的 ONNX 模型。您可以修改現有 ONNX 模型以包含自定義運算元,或從頭開始建立一個新模型。

要建立包含自定義運算元的新 ONNX 模型,可以使用 ONNX Python API。這裡有一個示例:test_pyops.py

從頭開始在 C++ 中建立自定義運算元

在實現自定義運算元之前,您需要一個包含一個或多個 ORT 自定義運算元的 ONNX 模型,這些運算元由 ONNX 轉換器(例如 ONNX-ScriptONNX 模型 API 等)建立。

1. 使用 PythonOp 快速驗證(可選)

在您真正為自己的用例開發自定義運算元之前,如果您想使用 Python 快速驗證 ONNX 模型,可以按照上述方式使用 Python 函式封裝自定義運算元。

import numpy
from onnxruntime_extensions import PyOp, onnx_op

# Implement the CustomOp by decorating a function with onnx_op
@onnx_op(op_type="Inverse", inputs=[PyOp.dt_float])
def inverse(x):
    # the user custom op implementation here:
    return numpy.linalg.inv(x)

# Run the model with this custom op
# model_func = PyOrtFunction(model_path)
# outputs = model_func(inputs)
# ...

2. 從 ONNX 模型生成自定義運算元的 C++ 模板程式碼(可選)

python -m onnxruntime-extensions.cmd --cpp-gen <model_path> <repository_dir>` If you are familiar with the ONNX model detail, you create the custom operator C++ classes directly.

3. 在生成的 C++ 檔案中實現 CustomOp Kernel Compute 方法。

自定義運算元核心 C++ 程式碼示例可以在 operators 資料夾中找到,例如 gaussian_blur。所有可在核心實現中使用的 C++ API 如下所示:

  • ONNX Runtime 自定義 API 文件
  • 整合在 ONNX Runtime Extensions 中可用於 C++ 程式碼的第三方庫 API 文件
    • OpenCV API 文件 https://docs.opencv.tw/4.x/
    • Google SentencePiece 庫文件 https://github.com/google/sentencepiece/blob/master/doc/api.md
    • dlib(矩陣和機器學習庫)C++ API 文件 https://dlib.net/algorithms.html
    • BlingFire 庫 https://github.com/microsoft/BlingFire
    • Google RE2 庫 https://github.com/google/re2/wiki/CplusplusAPI
    • JSON 庫 https://json.nlohmann.me/api/basic_json/

3. 構建和測試

  • 單元測試可以用 Python 或 C++ 實現,請檢視 test 資料夾以獲取更多示例
  • 請檢視 構建包,瞭解如何構建用於生產的不同語言包。

請檢視貢獻指南,瞭解是否可以將自定義運算元貢獻給 onnxruntime-extensions。