使用 Python 函数创建自定义算子

自定义算子是 ONNX Runtime 中的一个强大功能，它允许用户通过实现自己的算子来扩展运行时的功能，以执行标准 ONNX 算子集中未提供的特定操作。

在本文档中，我们将介绍如何使用 Python 函数创建自定义算子并将其集成到 ONNX Runtime 中进行推理。

步骤 1：定义自定义算子的 Python 函数

首先定义将作为自定义算子实现的 Python 函数。确保该函数与您期望的自定义算子的输入和输出张量形状兼容。Python 装饰器 @onnx_op 会将函数转换为自定义算子实现。以下是一个我们为 tokenizer 创建函数的示例

@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

因为 ONNXRuntimme 在加载模型时需要自定义算子 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 model 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 如下所示

  • ONNXRuntime 自定义 API 文档
  • 集成在 ONNXRuntime Extensions 中的第三方库 API 文档，可在 C++ 代码中使用
    • OpenCV API 文档 https://docs.opencv.ac.cn/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。