使用 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 在加载模型时需要自定义运算符模式，请通过 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 都在下面列出

  • ONNXRuntime 自定义 API 文档
  • 集成在 ONNXRuntime Extensions 中的第三方库 API 文档，可以在 C++ 代码中使用
    • OpenCV API 文档 https://docs.opencv.ac.cn/4.x/
    • Google SentencePiece Library 文档 https://github.com/google/sentencepiece/blob/master/doc/api.md
    • dlib（矩阵和 ML 库）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 文件夹
  • 查看 build-package，了解如何构建用于生产的不同语言包。

请查看 贡献，以了解是否有可能将自定义运算符贡献给 onnxruntime-extensions。