onnx2torch 是一个 ONNX 到 PyTorch 的转换器。 我们的转换器：

• 使用简单——只需调用 convert 函数即可转换 ONNX 模型；
• 易于扩展——您可以在 PyTorch 中编写自定义层，并使用 @add_converter 注册它；
• 可以转换回 ONNX——您可以使用 torch.onnx.export 函数将模型转换回 ONNX 格式。

如果您发现任何问题，请告知我们！也欢迎您提交合并请求。

请注意，此转换器仅支持有限数量的 PyTorch 和 ONNX 模型及操作。请在此处告诉我们您正在使用或希望从 ONNX 转换为 PyTorch 的模型：这里。

安装

pip install onnx2torch

或者

conda install -c conda-forge onnx2torch

使用

以下是一些使用示例。

转换

import onnx
import torch
from onnx2torch import convert

# ONNX 模型路径
onnx_model_path = "/some/path/mobile_net_v2.onnx"
# 您可以传递 ONNX 模型路径来直接转换，或者...
torch_model_1 = convert(onnx_model_path)

# 或者您也可以加载一个普通的 ONNX 模型并将其传递给转换器
onnx_model = onnx.load(onnx_model_path)
torch_model_2 = convert(onnx_model)

执行

我们可以像执行原始的 PyTorch 模型一样执行返回的 PyTorch 模型。

import onnxruntime as ort

# 创建示例数据
x = torch.ones((1, 2, 224, 224)).cuda()

out_torch = torch_model_1(x)

ort_sess = ort.InferenceSession(onnx_model_path)
outputs_ort = ort_sess.run(None, {"input": x.numpy()})

# 将 ONNX 输出与 PyTorch 输出进行比较
print(torch.max(torch.abs(outputs_ort - out_torch.detach().numpy())))
print(np.allclose(outputs_ort, out_torch.detach().numpy(), atol=1.0e-7))

模型

我们已经测试了以下模型：

分割模型：

• DeepLabV3+
• DeepLabV3 ResNet-50 (TorchVision)
• HRNet
• UNet (TorchVision)
• FCN ResNet-50 (TorchVision)
• LRASPP MobileNetV3 (TorchVision)

来自 MMdetection 的检测模型：

• SSDLite，带有 MobileNetV2 主干网络
• RetinaNet R50
• SSD300，带有 VGG 主干网络
• YOLOv3 d53
• YOLOv5

来自 TorchVision 的分类模型：

• ResNet-18
• ResNet-50
• MobileNetV2
• MobileNetV3 Large
• EfficientNet-B{0, 1, 2, 3}
• WideResNet-50
• ResNext-50
• VGG-16
• GoogLeNet
• MnasNet
• RegNet

Transformer 模型：

• ViT
• Swin
• GPT-J

:page_facing_up: 当前支持的操作列表可以在这里找到 [operators.md]。

如何向转换器添加新操作

下面我们展示如何通过新增 ONNX 操作来扩展 onnx2torch，这些操作同时被 PyTorch 和 ONNX 支持。

@add_converter(operation_type="Relu", version=6)
@add_converter(operation_type="Relu", version=13)
@add_converter(operation_type="Relu", version=14)
def _(node: OnnxNode, graph: OnnxGraph) -> OperationConverterResult:
    return OperationConverterResult(
        torch_module=nn.ReLU(),
        onnx_mapping=onnx_mapping_from_node(node=node),
    )

# 建议对字符串类型的 ONNX 属性使用 Enum。
class ReductionOnnxAttr(Enum):
    NONE = "none"
    ADD = "add"
    MUL = "mul"


class OnnxScatterND(nn.Module, OnnxToTorchModuleWithCustomExport):
    def __init__(self, reduction: ReductionOnnxAttr):
        super().__init__()
        self._reduction = reduction

    # 下面的方法应该以字典形式返回 ONNX 属性及其值。
    # 属性的数量、名称和值取决于算子集版本；
    # 该方法应返回正确的属性集合。
    # 注意：为每个键添加类型后缀：reduction -> reduction_s，其中 s 表示“字符串”。
    def _onnx_attrs(self, opset_version: int) -> Dict[str, Any]:
        onnx_attrs: Dict[str, Any] = {}

        # 在这里处理 opset 版本 < 16 的情况，此时没有 "reduction" 属性。
        if opset_version < 16:
            if self._reduction != ReductionOnnxAttr.NONE:
                raise ValueError(
                    "opset < 16 的 ScatterND 不支持"
                    f"reduction 属性不等于 {ReductionOnnxAttr.NONE.value}，"
                    f"但实际得到的是 {self._reduction.value}"
                )
            return onnx_attrs

        onnx_attrs["reduction_s"] = self._reduction.value
        return onnx_attrs

    def forward(
        self,
        data: torch.Tensor,
        indices: torch.Tensor,
        updates: torch.Tensor,
    ) -> torch.Tensor:
        def _forward():
            # ScatterND 的前向实现...
            return output

        if torch.onnx.is_in_onnx_export():
            # 请遵循我们的约定，参数包括：
            # 前向函数、操作类型、操作输入、操作属性。
            onnx_attrs = self._onnx_attrs(opset_version=get_onnx_version())
            return DefaultExportToOnnx.export(
                _forward, "ScatterND", data, indices, updates, onnx_attrs
            )

        return _forward()


@add_converter(operation_type="ScatterND", version=11)
@add_converter(operation_type="ScatterND", version=13)
@add_converter(operation_type="ScatterND", version=16)
def _(node: OnnxNode, graph: OnnxGraph) -> OperationConverterResult:
    node_attributes = node.attributes
    reduction = ReductionOnnxAttr(node_attributes.get("reduction", "none"))
    return OperationConverterResult(
        torch_module=OnnxScatterND(reduction=reduction),
        onnx_mapping=onnx_mapping_from_node(node=node),
    )

算子集版本的 workaround

如果您正在使用较旧算子集版本的模型，请尝试以下 workaround：

ONNX 版本转换 - 官方文档

import onnx
from onnx import version_converter
import torch
from onnx2torch import convert

# 加载 ONNX 模型。
model = onnx.load("model.onnx")
# 将模型转换为目标版本。
target_version = 13
converted_model = version_converter.convert_version(model, target_version)
# 转换为 PyTorch。
torch_model = convert(converted_model)
torch.save(torch_model, "model.pt")

注意：仅在现有算子集版本无法将模型转换为 PyTorch 时才使用此方法。结果可能会有所不同。

引用

要引用 onnx2torch，请使用“引用此仓库”按钮，或：

@misc{onnx2torch,
  title={onnx2torch},
  author={ENOT 开发者以及 Kalgin, Igor、Yanchenko, Arseny、Ivanov, Pyoter 和 Goncharenko, Alexander},
  year={2021},
  howpublished={\url{https://enot.ai/}},
  note={版本：x.y.z}
}

致谢

感谢 Dmitry Chudakov @cakeofwar42 的贡献。
特别感谢 Andrey Denisov @denisovap2013 的 logo 设计。

onnx2torch 快速上手指南

onnx2torch 是一个将 ONNX 模型转换为 PyTorch 模型的开源工具。它支持一键转换、自定义算子扩展，并允许将转换后的模型重新导出为 ONNX 格式。

环境准备

在开始之前，请确保您的开发环境满足以下要求：

• 操作系统：Linux, macOS 或 Windows
• Python 版本：建议 Python 3.7+
• 前置依赖：
  • PyTorch (torch)
  • ONNX (onnx)
  • NumPy (numpy)

提示：如果您尚未安装 PyTorch，请访问 PyTorch 官网 获取适合您环境的安装命令。国内用户可使用清华源加速安装：

pip install torch torchvision torchaudio --index-url https://pypi.tuna.tsinghua.edu.cn/simple

安装步骤

您可以通过 pip 或 conda 进行安装。推荐国内用户使用镜像源以提升下载速度。

方式一：使用 pip 安装（推荐）

pip install onnx2torch -i https://pypi.tuna.tsinghua.edu.cn/simple

方式二：使用 conda 安装

conda install -c conda-forge onnx2torch

基本使用

以下是将 ONNX 模型转换为 PyTorch 模型并进行推理验证的最简流程。

1. 转换模型

您可以直接传入 ONNX 文件路径，或者先加载 ONNX 模型对象再传入。

import onnx
import torch
from onnx2torch import convert

# 路径示例
onnx_model_path = "mobile_net_v2.onnx"

# 方法 A：直接传入文件路径
torch_model = convert(onnx_model_path)

# 方法 B：传入已加载的 ONNX 模型对象
# onnx_model = onnx.load(onnx_model_path)
# torch_model = convert(onnx_model)

# 将模型设置为评估模式并移至 GPU (可选)
torch_model.eval()
torch_model = torch_model.cuda()

2. 执行推理与验证

转换后的 torch_model 是一个标准的 PyTorch nn.Module，可以直接用于推理。以下示例对比了 ONNX Runtime 和转换后 PyTorch 模型的输出结果。

import numpy as np
import onnxruntime as ort

# 创建测试数据 (假设输入形状为 1x2x224x224)
x = torch.ones((1, 2, 224, 224)).cuda()

# 1. 运行 PyTorch 模型
with torch.no_grad():
    out_torch = torch_model(x)

# 2. 运行原始 ONNX 模型 (用于验证一致性)
ort_sess = ort.InferenceSession(onnx_model_path)
outputs_ort = ort_sess.run(None, {"input": x.cpu().numpy()})

# 3. 对比输出结果
# 计算最大绝对误差
max_diff = torch.max(torch.abs(outputs_ort[0] - out_torch.detach().cpu().numpy()))
print(f"Max difference: {max_diff}")

# 检查数值是否接近 (容差 1e-7)
is_close = np.allclose(outputs_ort[0], out_torch.detach().cpu().numpy(), atol=1.0e-7)
print(f"Outputs match: {is_close}")

3. 重新导出为 ONNX (可选)

如果需要将转换后的 PyTorch 模型再次导出为 ONNX 格式，可以使用原生 PyTorch 接口：

dummy_input = torch.randn(1, 2, 224, 224).cuda()
torch.onnx.export(
    torch_model, 
    dummy_input, 
    "exported_model.onnx",
    opset_version=13,
    input_names=['input'],
    output_names=['output']
)