PyTorch教程：将简单模型导出为ONNX格式的完整指南

原创于 2025-06-05 09:01:49 发布 · 481 阅读

CC 4.0 BY-SA版权

PyTorch教程：将简单模型导出为ONNX格式的完整指南

前言

在深度学习领域，PyTorch因其灵活性和易用性而广受欢迎。然而，在实际生产环境中，我们经常需要将训练好的模型部署到不同的平台和设备上。ONNX（Open Neural Network Exchange）作为一种开放的神经网络交换格式，能够帮助我们实现这一目标。本文将详细介绍如何使用PyTorch将模型导出为ONNX格式。

ONNX简介

ONNX是一种开放的机器学习模型格式标准，它允许模型在不同的框架和硬件平台之间进行转换和运行。通过ONNX，我们可以：

实现跨框架的模型转换（如PyTorch到TensorFlow）
优化模型性能
部署到各种硬件平台（包括云端和边缘设备）

准备工作

在开始之前，我们需要安装必要的依赖包：

pip install --upgrade onnx onnxscript onnxruntime

构建示例模型

我们将使用一个简单的图像分类器模型作为示例，这个模型结构与PyTorch官方教程中的"60分钟入门"示例类似：

import torch
import torch.nn as nn
import torch.nn.functional as F

class ImageClassifierModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv1 = nn.Conv2d(1, 6, 5)  # 输入通道1，输出通道6，卷积核5x5
        self.conv2 = nn.Conv2d(6, 16, 5) # 输入通道6，输出通道16，卷积核5x5
        self.fc1 = nn.Linear(16*5*5, 120) # 全连接层
        self.fc2 = nn.Linear(120, 84)    # 全连接层
        self.fc3 = nn.Linear(84, 10)     # 输出层，10个类别

    def forward(self, x):
        x = F.max_pool2d(F.relu(self.conv1(x)), (2, 2))
        x = F.max_pool2d(F.relu(self.conv2(x)), 2)
        x = torch.flatten(x, 1)
        x = F.relu(self.fc1(x))
        x = F.relu(self.fc2(x))
        x = self.fc3(x)
        return x

模型导出步骤

1. 实例化模型并准备输入数据

torch_model = ImageClassifierModel()
example_inputs = (torch.randn(1, 1, 32, 32),)  # 批大小1，1个通道，32x32图像

2. 导出模型到ONNX格式

PyTorch 2.5引入了新的ONNX导出器，使用dynamo=True参数启用：

onnx_program = torch.onnx.export(torch_model, example_inputs, dynamo=True)

3. 优化ONNX模型（可选）

onnx_program.optimize()  # 执行常量折叠和冗余节点消除

4. 保存ONNX模型

onnx_program.save("image_classifier_model.onnx")

5. 验证ONNX模型

import onnx
onnx_model = onnx.load("image_classifier_model.onnx")
onnx.checker.check_model(onnx_model)  # 检查模型是否有效

可视化模型

我们可以使用Netron工具来可视化导出的ONNX模型。Netron是一个开源的神经网络可视化工具，支持多种框架的模型格式。

访问Netron的网页版
上传我们保存的image_classifier_model.onnx文件
查看模型结构和各层参数

使用ONNX Runtime运行模型

1. 准备输入数据

onnx_inputs = [tensor.numpy(force=True) for tensor in example_inputs]

2. 创建ONNX Runtime会话

ort_session = onnxruntime.InferenceSession(
    "./image_classifier_model.onnx", 
    providers=["CPUExecutionProvider"]
)

3. 运行模型

onnxruntime_input = {input_arg.name: input_value 
                    for input_arg, input_value 
                    in zip(ort_session.get_inputs(), onnx_inputs)}
onnxruntime_outputs = ort_session.run(None, onnxruntime_input)[0]

验证结果一致性

为了确保导出的ONNX模型与原始PyTorch模型行为一致，我们需要比较它们的输出：

torch_outputs = torch_model(*example_inputs)

# 确保输出数量一致
assert len(torch_outputs) == len(onnxruntime_outputs)

# 比较每个输出是否接近
for torch_output, onnxruntime_output in zip(torch_outputs, onnxruntime_outputs):
    torch.testing.assert_close(torch_output, torch.tensor(onnxruntime_output))

print("PyTorch和ONNX Runtime输出匹配成功！")