评价此页
★ ★ ★ ★ ★
beginner/onnx/onnx_registry_tutorial

注意

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

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 算子

您将学到什么

  • 如何重写或添加对 ONNX 中 PyTorch 算子的支持。

  • 如何为专用运行时集成自定义 ONNX 算子。

  • 如何实现自定义 PyTorch 算子并将其转换为 ONNX。

先决条件#

在开始本教程之前，请确保您已满足以下先决条件：

重写现有 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.13.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 算子类型。 

print(onnx_program.model)

<
    ir_version=10,
    opset_imports={'com.microsoft': 1, '': 20},
    producer_name='pytorch',
    producer_version='2.13.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.13.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 导出器所有知识的乐趣。

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

脚本运行总时间：（0 分 3.280 秒）

在此页面上
PyTorch 库