使用 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-Script、ONNX 模型 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. 構建和測試
請檢視貢獻指南,瞭解是否可以將自定義運算元貢獻給 onnxruntime-extensions。