注意
请 跳转到末尾 下载完整的示例代码。
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。
先决条件#
在开始本教程之前，请确保您已完成以下先决条件：
torch >= 2.6目标 PyTorch 运算符
已完成 ONNX Script 教程，然后继续。
ONNX Script 的运算符实现。
重写现有 PyTorch 运算符的实现#
尽管 ONNX 导出器团队尽最大努力支持所有 PyTorch 运算符，但其中一些可能尚未得到支持。在本节中，我们将演示如何将不受支持的 PyTorch 运算符添加到 ONNX Registry。
注意
实现不受支持的 PyTorch 运算符的步骤与使用自定义实现替换现有 PyTorch 运算符的实现步骤相同。因为本教程中实际上没有不受支持的 PyTorch 运算符，所以我们将利用这一点，以同样的方式替换 `torch.ops.aten.add.Tensor` 的实现，就像该运算符未被 ONNX 导出器实现一样。
当模型因不受支持的运算符而无法导出到 ONNX 时，ONNX 导出器将显示类似以下的错误消息：
No decompositions registered for [...]
错误消息表明不受支持的 PyTorch 运算符是 `torch.ops.aten.add.Tensor`。该运算符的类型是 `
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 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={'': 20},
producer_name='pytorch',
producer_version='2.9.0+cu128',
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 运算符的模型，但运行时（如 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()
[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` 命名空间的 `Gelu` op_type。
print(onnx_program.model)
<
ir_version=10,
opset_imports={'com.microsoft': 1, '': 20},
producer_name='pytorch',
producer_version='2.9.0+cu128',
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={'': 20},
producer_name='pytorch',
producer_version='2.9.0+cu128',
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`。
最后，我们验证结果。
结论#
恭喜！在本教程中，我们探索了 `custom_translation_table` 选项，并学习了如何使用 ONNX Script 为不受支持或现有的 PyTorch 运算符创建自定义实现。
最后，我们利用 ONNX Runtime 执行模型并与 PyTorch 的结果进行比较，从而全面了解了在 ONNX 生态系统中处理不受支持的运算符。
延伸阅读#
下面的列表引用了从基本示例到高级场景的教程，不一定按列出的顺序。您可以随时跳转到您感兴趣的特定主题，或者坐下来，享受学习 ONNX 导出器所有知识的乐趣。
脚本总运行时间： (0 分钟 2.891 秒)
