目录
  • 教程 >
  • 扩展ONNX导出器操作符支持
快捷方式
beginner/onnx/onnx_registry_tutorial

注意

点击 此处 下载完整示例代码

ONNX入门 || 将PyTorch模型导出到ONNX || 扩展ONNX导出器操作符支持 || 将包含控制流的模型导出到ONNX

扩展ONNX导出器操作符支持

创建日期：2023年10月06日 | 最后更新：2025年03月05日 | 最后验证：2024年11月05日

作者： 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时，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()

/usr/local/lib/python3.10/dist-packages/onnxscript/converter.py:823: FutureWarning:

'onnxscript.values.Op.param_schemas' is deprecated in version 0.1 and will be removed in the future. Please use '.op_signature' instead.

/usr/local/lib/python3.10/dist-packages/onnxscript/converter.py:823: FutureWarning:

'onnxscript.values.OnnxFunction.param_schemas' is deprecated in version 0.1 and will be removed in the future. Please use '.op_signature' instead.

[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 decomposition...
[torch.onnx] Run decomposition... ✅
[torch.onnx] Translate the graph into ONNX...
[torch.onnx] Translate the graph into ONNX... ✅

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

print(onnx_program.model)

<
    ir_version=10,
    opset_imports={'pkg.onnxscript.torch_lib.common': 1, '': 18},
    producer_name='pytorch',
    producer_version='2.7.0+cu126',
    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_0
         %"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操作符的模型，但运行时（例如Microsoft的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()

'Gelu' is not a known op in 'com.microsoft'
[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 decomposition...
[torch.onnx] Run decomposition... ✅
[torch.onnx] Translate the graph into ONNX...
[torch.onnx] Translate the graph into ONNX... ✅

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

print(onnx_program.model)

<
    ir_version=10,
    opset_imports={'pkg.onnxscript.torch_lib.common': 1, 'com.microsoft': 1, '': 18},
    producer_name='pytorch',
    producer_version='2.7.0+cu126',
    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 decomposition...
[torch.onnx] Run decomposition... ✅
[torch.onnx] Translate the graph into ONNX...
[torch.onnx] Translate the graph into ONNX... ✅
ONNXProgram(
    model=
        <
            ir_version=10,
            opset_imports={'pkg.onnxscript.torch_lib.common': 1, '': 18},
            producer_name='pytorch',
            producer_version='2.7.0+cu126',
            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_Round_1
                 %"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: ExportGraphSignature(input_specs=[InputSpec(kind=<InputKind.USER_INPUT: 1>, arg=TensorArgument(name='input'), target=None, persistent=None)], output_specs=[OutputSpec(kind=<OutputKind.USER_OUTPUT: 1>, arg=TensorArgument(name='add_and_round_op'), target=None)])
        Range constraints: {}

)

翻译正在使用我们的自定义实现，将 torch.ops.mylibrary.add_and_round_op.default 操作符在 torch.export.ExportedProgram` 中的实现转换为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 分钟 2.409 秒)

由Sphinx-Gallery生成

文档

查阅PyTorch的全面开发者文档

查看文档

教程

获取针对初学者和高级开发者的深入教程

查看教程

资源

查找开发资源并获得问题的解答

查看资源

为了分析流量并优化您的体验，本网站使用cookies。点击或导航即表示您同意我们使用cookies。作为本网站的当前维护者，Facebook的Cookies政策适用。了解更多，包括可用的控制选项： Cookies政策。