目錄
  • 教程 >
  • 擴充套件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政策。