评价此页

★ ★ ★ ★ ★

beginner/onnx/onnx_registry_tutorial

在 Google Colab 中运行

注意

跳转至页面底部下载完整示例代码。

ONNX 简介 || 将 PyTorch 模型导出为 ONNX || 扩展 ONNX 导出器算子支持 || 将带有控制流的模型导出为 ONNX

扩展 ONNX 导出器算子支持#

创建日期：2023 年 10 月 6 日 | 最后更新：2025 年 3 月 5 日 | 最后验证：2024 年 11 月 5 日

作者： Ti-Tai Wang, Justin Chu

概述#

本教程介绍如何为不受支持的 PyTorch 算子创建 ONNX 实现，或使用您自己的实现替换现有实现。

我们将涵盖需要扩展 ONNX 导出器算子支持的三种场景：

覆盖现有 PyTorch 算子的实现
使用自定义 ONNX 算子
支持自定义 PyTorch 算子

您将学到什么

如何覆盖 PyTorch 算子或为其添加 ONNX 支持。
如何为专用运行时集成自定义 ONNX 算子。
如何实现自定义 PyTorch 算子并将其转换为 ONNX。

先决条件#

在开始本教程之前，请确保您已完成以下准备工作：

torch >= 2.6
目标 PyTorch 算子
在继续之前已完成 ONNX Script 教程
使用 ONNX Script 实现算子

覆盖现有 PyTorch 算子的实现#

尽管 ONNX 导出器团队尽最大努力支持所有 PyTorch 算子，但其中一些可能尚未得到支持。在本节中，我们将演示如何将不受支持的 PyTorch 算子添加到 ONNX 注册表中。

注意

实现不受支持的 PyTorch 算子的步骤，与使用自定义实现替换现有 PyTorch 算子的实现步骤相同。由于在本教程中我们实际上并没有一个不受支持的 PyTorch 算子，我们将利用这一点，并以类似处理未实现算子的方式，用自定义实现替换 torch.ops.aten.add.Tensor 的实现。

当模型由于算子不受支持而无法导出为 ONNX 时，ONNX 导出器将显示类似以下的错误信息：

No decompositions registered for [...]

错误信息表明不受支持的 PyTorch 算子是 torch.ops.aten.add.Tensor。该算子属于 <class 'torch._ops.OpOverload'> 类型，我们将使用该算子作为目标来注册我们的自定义实现。

import torch
import onnxscript

# Opset 18 is the standard supported version as of PyTorch 2.6
from onnxscript import opset18 as op


# Create a model that uses the operator torch.ops.aten.add.Tensor
class Model(torch.nn.Module):
    def forward(self, input_x, input_y):
        return torch.ops.aten.add.Tensor(input_x, input_y)


# NOTE: The function signature (including parameter names) must match the signature of the unsupported PyTorch operator.
# https://github.com/pytorch/pytorch/blob/main/aten/src/ATen/native/native_functions.yaml
# All attributes must be annotated with type hints.
def custom_aten_add(self, other, alpha: float = 1.0):
    if alpha != 1.0:
        alpha = op.CastLike(alpha, other)
        other = op.Mul(other, alpha)
    # To distinguish the custom implementation from the builtin one, we switch the order of the inputs
    return op.Add(other, self)


x = torch.tensor([1.0])
y = torch.tensor([2.0])

# Then we provide the custom implementation to the ONNX exporter as a ``custom_translation_table``.
onnx_program = torch.onnx.export(
    Model().eval(),
    (x, y),
    dynamo=True,
    custom_translation_table={
        torch.ops.aten.add.Tensor: custom_aten_add,
    },
)
# Optimize the ONNX graph to remove redundant nodes
onnx_program.optimize()

[torch.onnx] Obtain model graph for `Model()` with `torch.export.export(..., strict=False)`...
[torch.onnx] Obtain model graph for `Model()` with `torch.export.export(..., strict=False)`... ✅
[torch.onnx] Run decompositions...
/usr/lib/python3.10/copyreg.py:101: FutureWarning: `isinstance(treespec, LeafSpec)` is deprecated, use `isinstance(treespec, TreeSpec) and treespec.is_leaf()` instead.
  return cls.__new__(cls, *args)
[torch.onnx] Run decompositions... ✅
[torch.onnx] Translate the graph into ONNX...
[torch.onnx] Translate the graph into ONNX... ✅
[torch.onnx] Optimize the ONNX graph...
[torch.onnx] Optimize the ONNX graph... ✅

现在让我们检查该模型，并验证模型是否正在使用自定义实现。

print(onnx_program.model)

<
    ir_version=10,
    opset_imports={'': 20},
    producer_name='pytorch',
    producer_version='2.12.0+cu130',
    domain=None,
    model_version=None,
>
graph(
    name=main_graph,
    inputs=(
        %"input_x"<FLOAT,[1]>,
        %"input_y"<FLOAT,[1]>
    ),
    outputs=(
        %"add"<FLOAT,[1]>
    ),
) {
    0 |  # node_add
         %"add"<FLOAT,[1]> ⬅️ ::Add(%"input_y", %"input_x")
    return %"add"<FLOAT,[1]>
}

转换过程使用了我们的自定义实现：在节点 node_Add_0 中，input_y 现在排在第一位，input_x 排在第二位。

我们可以使用 ONNX Runtime 来运行模型，并通过直接在输入张量上调用 torch.onnx.ONNXProgram 来验证结果。

result = onnx_program(x, y)[0]
torch.testing.assert_close(result, torch.tensor([3.0]))

使用自定义 ONNX 算子#

在这种情况下，我们使用标准的 PyTorch 算子创建模型，但运行时（例如微软的 ONNX Runtime）可以为该内核提供自定义实现，从而有效地替换现有实现。

在下面的示例中，我们使用 ONNX Runtime 提供的 com.microsoft.Gelu 算子，它与 ONNX 规范中的 Gelu 不同。

class GeluModel(torch.nn.Module):
    def forward(self, input_x):
        return torch.ops.aten.gelu(input_x)


# Create a namespace for the custom operator using ONNX Script
# ``com.microsoft`` is an official ONNX Runtime namespace
microsoft_op = onnxscript.values.Opset(domain="com.microsoft", version=1)

# NOTE: The function signature (including parameter names) must match the signature of the unsupported PyTorch operator.
# https://github.com/pytorch/pytorch/blob/main/aten/src/ATen/native/native_functions.yaml
# NOTE: All attributes must be annotated with type hints.
# The function must be scripted using the ``@onnxscript.script()`` decorator when
# using operators from custom domains. This may be improved in future versions.
from onnxscript import FLOAT


@onnxscript.script(microsoft_op)
def custom_aten_gelu(self: FLOAT, approximate: str = "none") -> FLOAT:
    return microsoft_op.Gelu(self)


onnx_program = torch.onnx.export(
    GeluModel().eval(),
    (x,),
    dynamo=True,
    custom_translation_table={
        torch.ops.aten.gelu.default: custom_aten_gelu,
    },
)

# Optimize the ONNX graph to remove redundant nodes
onnx_program.optimize()

[torch.onnx] Obtain model graph for `GeluModel()` with `torch.export.export(..., strict=False)`...
[torch.onnx] Obtain model graph for `GeluModel()` with `torch.export.export(..., strict=False)`... ✅
[torch.onnx] Run decompositions...
/usr/lib/python3.10/copyreg.py:101: FutureWarning: `isinstance(treespec, LeafSpec)` is deprecated, use `isinstance(treespec, TreeSpec) and treespec.is_leaf()` instead.
  return cls.__new__(cls, *args)
[torch.onnx] Run decompositions... ✅
[torch.onnx] Translate the graph into ONNX...
[torch.onnx] Translate the graph into ONNX... ✅
[torch.onnx] Optimize the ONNX graph...
[torch.onnx] Optimize the ONNX graph... ✅

让我们检查模型，并验证模型使用了来自 com.microsoft 命名空间的 Gelu 算子类型 (op_type)。

print(onnx_program.model)

<
    ir_version=10,
    opset_imports={'com.microsoft': 1, '': 20},
    producer_name='pytorch',
    producer_version='2.12.0+cu130',
    domain=None,
    model_version=None,
>
graph(
    name=main_graph,
    inputs=(
        %"input_x"<FLOAT,[1]>
    ),
    outputs=(
        %"gelu"<FLOAT,[1]>
    ),
) {
    0 |  # n0
         %"gelu"<FLOAT,[1]> ⬅️ com.microsoft::Gelu(%"input_x")
    return %"gelu"<FLOAT,[1]>
}

与上一个示例类似，我们可以使用 ONNX Runtime 运行模型并验证结果。

result = onnx_program(x)[0]
torch.testing.assert_close(result, torch.ops.aten.gelu(x))

支持自定义 PyTorch 算子#

在这种情况下，算子是用户实现并注册到 PyTorch 的算子。

在下面的示例中，我们希望使用一个自定义算子，它接受一个张量输入并返回一个输出。该算子将输入与自身相加，并返回舍入后的结果。

首先，我们假设自定义算子已经实现并使用 torch.library.custom_op() 进行了注册。您可以参考在 Python 中创建新的自定义算子获取有关如何创建自定义算子的详细指南。

# Define and use the operator in PyTorch
@torch.library.custom_op("mylibrary::add_and_round_op", mutates_args=())
def add_and_round_op(input: torch.Tensor) -> torch.Tensor:
    return torch.round(input + input)


@add_and_round_op.register_fake
def _add_and_round_op_fake(tensor_x):
    return torch.empty_like(tensor_x)


class AddAndRoundModel(torch.nn.Module):
    def forward(self, input):
        return add_and_round_op(input)


# Implement the custom operator in ONNX using ONNX Script
def onnx_add_and_round(input):
    return op.Round(op.Add(input, input))


onnx_program = torch.onnx.export(
    AddAndRoundModel().eval(),
    (x,),
    dynamo=True,
    custom_translation_table={
        torch.ops.mylibrary.add_and_round_op.default: onnx_add_and_round,
    },
)

# Optimize the ONNX graph to remove redundant nodes
onnx_program.optimize()
print(onnx_program)

[torch.onnx] Obtain model graph for `AddAndRoundModel()` with `torch.export.export(..., strict=False)`...
[torch.onnx] Obtain model graph for `AddAndRoundModel()` with `torch.export.export(..., strict=False)`... ✅
[torch.onnx] Run decompositions...
/usr/lib/python3.10/copyreg.py:101: FutureWarning: `isinstance(treespec, LeafSpec)` is deprecated, use `isinstance(treespec, TreeSpec) and treespec.is_leaf()` instead.
  return cls.__new__(cls, *args)
[torch.onnx] Run decompositions... ✅
[torch.onnx] Translate the graph into ONNX...
[torch.onnx] Translate the graph into ONNX... ✅
[torch.onnx] Optimize the ONNX graph...
[torch.onnx] Optimize the ONNX graph... ✅
ONNXProgram(
    model=
        <
            ir_version=10,
            opset_imports={'': 20},
            producer_name='pytorch',
            producer_version='2.12.0+cu130',
            domain=None,
            model_version=None,
        >
        graph(
            name=main_graph,
            inputs=(
                %"input"<FLOAT,[1]>
            ),
            outputs=(
                %"add_and_round_op"<FLOAT,[1]>
            ),
        ) {
            0 |  # node_Add_0
                 %"val_0"<FLOAT,[1]> ⬅️ ::Add(%"input", %"input")
            1 |  # node_add_and_round_op
                 %"add_and_round_op"<FLOAT,[1]> ⬅️ ::Round(%"val_0")
            return %"add_and_round_op"<FLOAT,[1]>
        }


    ,
    exported_program=
        ExportedProgram:
            class GraphModule(torch.nn.Module):
                def forward(self, input: "f32[1]"):
                    input_1 = input

                    # File: /var/lib/workspace/beginner_source/onnx/onnx_registry_tutorial.py:215 in forward, code: return add_and_round_op(input)
                    add_and_round_op: "f32[1]" = torch.ops.mylibrary.add_and_round_op.default(input_1);  input_1 = None
                    return (add_and_round_op,)

        Graph signature:
            # inputs
            input: USER_INPUT

            # outputs
            add_and_round_op: USER_OUTPUT

        Range constraints: {}

)

转换过程使用我们的自定义实现，将 torch.export.ExportedProgram 中的 torch.ops.mylibrary.add_and_round_op.default 算子转换为 ONNX 算子 Add 和 Round。

最后，我们验证结果。

result = onnx_program(x)[0]
torch.testing.assert_close(result, add_and_round_op(x))

结论#

恭喜！在本教程中，我们探索了 custom_translation_table 选项，并发现了如何使用 ONNX Script 为不受支持或现有的 PyTorch 算子创建自定义实现。

最后，我们利用 ONNX Runtime 执行模型并将结果与 PyTorch 进行对比，这让我们全面了解了如何在 ONNX 生态系统中处理不受支持的算子。

延伸阅读#

下面的列表引用了从基本示例到高级场景的教程，不一定按列出的顺序。您可以随时跳转到您感兴趣的特定主题，或者坐下来，享受学习 ONNX 导出器所有知识的乐趣。

将 PyTorch 模型导出到 ONNX
扩展 ONNX 导出器算子支持
将带控制流的模型导出到 ONNX

脚本运行总耗时：（0 分钟 3.349 秒）