详解PyTorch模型转换：从.pth到.onnx的超详细步骤

详解PyTorch模型转换：从.pth到.onnx的超详细步骤

引用

CSDN

1.

https://blog.csdn.net/zhouzongxin94/article/details/144060690

随着深度学习技术的不断发展，模型的跨平台部署和优化变得越来越重要。PyTorch作为主流的深度学习框架之一，其模型转换为ONNX格式可以实现更广泛的部署和优化。本文将详细介绍如何将PyTorch模型转换为ONNX格式，并通过具体示例帮助读者掌握这一技能。

PyTorch与ONNX简介

PyTorch

PyTorch是由Facebook AI Research团队开发的开源深度学习框架，以其动态图机制和简洁的API设计，广受研究社区和工业界的欢迎。PyTorch支持多种神经网络结构，具有强大的自动微分功能，适用于复杂的模型设计与实验。

ONNX

ONNX（Open Neural Network Exchange）是由微软和Facebook共同发起的开源项目，旨在实现深度学习模型的跨平台互操作性。ONNX定义了一种通用的模型表示格式，支持多种深度学习框架（如PyTorch、TensorFlow等），使得模型可以在不同的平台和设备上无缝迁移和部署。

为什么需要将PyTorch模型转换为ONNX

将PyTorch模型转换为ONNX格式具有以下优势：

1. 跨平台部署：ONNX支持在多种平台（如Windows、Linux、macOS）和设备（如CPU、GPU、移动端）上部署，提升模型的适用范围。
2. 框架互操作性：通过ONNX，能够实现不同深度学习框架之间的模型转换，方便在已有生态系统中集成和应用。
3. 优化与加速：ONNX Runtime等工具提供了对ONNX模型的优化和加速，提升推理性能，降低资源消耗。
4. 标准化管理：ONNX作为开放标准，促进模型的标准化管理和共享，便于团队协作与知识传播。

转换步骤详解

本文将通过具体的步骤和代码示例，详细介绍如何将PyTorch模型转换为ONNX格式，并验证转换后的模型。

环境准备

在开始转换之前，确保系统中已安装以下必要的软件和库：

1. Python：建议使用Python 3.6及以上版本。
2. PyTorch：确保安装了最新版本的PyTorch。
3. ONNX：用于处理ONNX模型。
4. ONNX Runtime（可选）：用于验证ONNX模型的推理正确性。

可以使用以下命令安装所需库：

pip install torch onnx onnxruntime  

模型准备

以一个简单的PyTorch模型为例，本文将展示模型的定义、训练（或加载预训练模型），并进行转换。

import torch
import torch.nn as nn
# 定义一个简单的神经网络
class SimpleNet(nn.Module):
    def __init__(self, input_size=784, hidden_size=500, num_classes=10):
        super(SimpleNet, self).__init__()
        self.fc1 = nn.Linear(input_size, hidden_size) 
        self.relu = nn.ReLU()
        self.fc2 = nn.Linear(hidden_size, num_classes)  

    def forward(self, x):
        out = self.fc1(x)
        out = self.relu(out)
        out = self.fc2(out)
        return out
# 初始化模型
model = SimpleNet()
# 假设已经训练好并保存为.pth文件
model.load_state_dict(torch.load('model.pth'))
model.eval()

转换代码示例

以下是将PyTorch模型转换为ONNX格式的详细步骤：

1. 定义输入示例：ONNX转换需要输入一个示例输入，以确定模型的输入输出形状和类型。
2. 执行转换：使用torch.onnx.export函数进行模型转换。
3. 保存ONNX模型：将转换后的模型保存为.onnx文件。

import torch.onnx
# 定义输入示例（批量大小为1，784个特征）
dummy_input = torch.randn(1, 784)
# 指定输出ONNX模型的文件名
onnx_model_path = "model.onnx"
# 导出模型
torch.onnx.export(
    model,                      # 要转换的PyTorch模型
    dummy_input,                # 模型的输入示例
    onnx_model_path,            # ONNX模型的存储路径
    export_params=True,         # 是否导出训练好的参数
    opset_version=11,           # ONNX的操作集版本
    do_constant_folding=True,   # 是否优化常量折叠
    input_names = ['input'],    # 输入节点的名称
    output_names = ['output'],  # 输出节点的名称
    dynamic_axes={'input' : {0 : 'batch_size'},    # 动态轴设置
                  'output' : {0 : 'batch_size'}}
)
print(f"模型已成功转换并保存为 {onnx_model_path}")

验证ONNX模型

为了确保转换的ONNX模型与原始PyTorch模型在推理结果上保持一致，可以使用ONNX Runtime进行验证。

import onnx
import onnxruntime
import numpy as np
# 加载ONNX模型
onnx_model = onnx.load(onnx_model_path)
onnx.checker.check_model(onnx_model)
print("ONNX模型检查通过！")
# 创建ONNX Runtime会话
ort_session = onnxruntime.InferenceSession(onnx_model_path)
# 准备输入数据
input_data = dummy_input.numpy()
# 使用PyTorch模型进行推理
with torch.no_grad():
    torch_output = model(dummy_input).numpy()
# 使用ONNX Runtime进行推理
ort_inputs = {ort_session.get_inputs()[0].name: input_data}
ort_output = ort_session.run(None, ort_inputs)[0]
# 比较两者的输出
np.testing.assert_allclose(torch_output, ort_output, rtol=1e-03, atol=1e-05)
print("ONNX模型的推理结果与PyTorch模型一致！")

常见问题与解决方案

1. 操作集（opset）版本不兼容

如果在转换过程中遇到Unsupported operator错误，可能是因为使用的opset版本过低或过高。尝试调整opset_version参数，例如使用11或12版本。

2. 动态轴设置问题

在某些模型中，批量大小或序列长度可能是动态变化的。确保在torch.onnx.export中正确设置dynamic_axes参数，以支持动态输入。

3. 自定义层或操作不支持

如果模型中包含自定义的层或不被ONNX支持的操作，可能需要手动添加自定义的ONNX算子实现，或将其简化为ONNX支持的操作。

4. 权重参数未正确导出

确保在torch.onnx.export中设置export_params=True，以导出模型的所有权重参数。

实践案例

将训练好的ResNet模型转换为ONNX

以下是将预训练的ResNet模型转换为ONNX格式的示例：

import torch
import torchvision.models as models
import torch.onnx
# 加载预训练的ResNet18模型
resnet18 = models.resnet18(pretrained=True)
resnet18.eval()
# 定义输入示例（批量大小为1，3个通道，224x224图像）
dummy_input = torch.randn(1, 3, 224, 224)
# 导出ONNX模型
torch.onnx.export(
    resnet18,
    dummy_input,
    "resnet18.onnx",
    export_params=True,
    opset_version=11,
    do_constant_folding=True,
    input_names = ['input'],
    output_names = ['output'],
    dynamic_axes={'input' : {0 : 'batch_size'},
                  'output' : {0 : 'batch_size'}}
)
print("ResNet18模型已成功转换为ONNX格式！")

验证转换后的ResNet18模型

import onnx
import onnxruntime
import numpy as np
from PIL import Image
from torchvision import transforms
# 加载ONNX模型
onnx_model = onnx.load("resnet18.onnx")
onnx.checker.check_model(onnx_model)
print("ONNX模型检查通过！")
# 创建ONNX Runtime会话
ort_session = onnxruntime.InferenceSession("resnet18.onnx")
# 准备输入图像
preprocess = transforms.Compose([
    transforms.Resize(256),
    transforms.CenterCrop(224),
    transforms.ToTensor(),
])
img = Image.open("example.jpg")
input_tensor = preprocess(img)
input_batch = input_tensor.unsqueeze(0).numpy()
# 使用PyTorch模型进行推理
with torch.no_grad():
    torch_output = resnet18(torch.from_numpy(input_batch)).numpy()
# 使用ONNX Runtime进行推理
ort_inputs = {ort_session.get_inputs()[0].name: input_batch}
ort_output = ort_session.run(None, ort_inputs)[0]
# 计算两者的差异
difference = np.abs(torch_output - ort_output)
print(f"输出差异最大值: {difference.max()}")

总结

将PyTorch模型转换为ONNX格式，是实现模型跨平台部署和优化的重要步骤。通过本文详细的步骤和代码示例，相信你已经掌握了这一技巧。从环境准备、模型定义、转换执行到模型验证，每一步都至关重要。遇到转换过程中的常见问题时，参考本文提供的解决方案，可以帮助你迅速排查并解决问题。掌握这一技能，将极大地提升你在深度学习项目中的灵活性和效率，助力你的模型在多样化的应用场景中大放异彩。