onnx2torch - ONNX模型转换至PyTorch的转换器

项目介绍：onnx2torch

项目概述

onnx2torch是一个将ONNX模型转换为PyTorch模型的转换器。它的设计目标是简单易用，同时又便于扩展。用户可以通过简单的函数调用来完成ONNX模型的转换，并且可以通过自定义PyTorch层来扩展转换功能。

项目特点

易于使用：只需通过函数调用convert即可将ONNX模型转换为PyTorch模型。
易于扩展：用户可以通过编写自定义的PyTorch层，并使用@add_converter注解来注册该层，轻松扩展转换器的功能。
双向转换：不仅可以将ONNX模型转换为PyTorch模型，还可以使用torch.onnx.export函数将模型转换回ONNX格式。

使用说明

安装方法

用户可以通过pip或conda命令轻松安装onnx2torch转换器：

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模型文件路径传递给convert函数进行转换
torch_model_1 = convert(onnx_model_path)

# 或者可以加载一个ONNX模型对象并传递给convert函数
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))

支持的模型

onnx2torch目前已经成功测试了多种模型，包括：

分割模型：DeepLabV3+、DeepLabV3 ResNet-50、HRNet、UNet、FCN ResNet-50、LRASPP MobileNetV3等。
检测模型：如SSDLite、RetinaNet、SSD300、YOLOv3和YOLOv5等。
分类模型：包括ResNet、MobileNet、EfficientNet、VGG、GoogLeNet等。
Transformer模型：ViT、Swin、GPT-J等。

如何添加新操作

用户可以通过编写和注册自定义转换器来为onnx2torch添加新操作。以下是示例代码片段，展示了如何添加对一个简单运算的支持：

@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官方文档进行版本转换。

致谢

特别感谢Dmitry Chudakov和Andrey Denisov在项目开发和设计中的贡献。项目由ENOT开发者团队，包括Igor Kalgin、Arseny Yanchenko、Pyoter Ivanov和Alexander Goncharenko于2021年推出。

onnx2torch是一个开源项目，社区欢迎任何形式的贡献和反馈。只需使用GitHub的功能提交问题或创建合并请求，便可参与到项目的改进中。