从 PyTorch 到 ONNX：深度学习模型导出全解析

news2025/5/25 22:15:58

在模型训练完毕后，我们通常希望将其部署到推理平台中，比如 TensorRT、ONNX Runtime 或移动端框架。而 ONNX（Open Neural Network Exchange）正是 PyTorch 与这些平台之间的桥梁。

本文将以一个图像去噪模型 SimpleDenoiser 为例，手把手带你完成 PyTorch 模型导出为 ONNX 格式的全过程，并解析每一行代码背后的逻辑。

准备工作

我们假设你已经训练好一个图像去噪模型并保存为 .pth 文件，模型结构自编码器实现如下（略）：

class SimpleDenoiser(nn.Module):
    def __init__(self):
        super(SimpleDenoiser, self).__init__()
        self.encoder = nn.Sequential(
            nn.Conv2d(3, 64, 3, padding=1), nn.ReLU(),
            nn.Conv2d(64, 64, 3, padding=1), nn.ReLU()
        )
        self.decoder = nn.Sequential(
            nn.Conv2d(64, 64, 3, padding=1), nn.ReLU(),
            nn.Conv2d(64, 3, 3, padding=1)
        )

    def forward(self, x):
        x = self.encoder(x)
        x = self.decoder(x)
        return x

导出代码分解

我们现在来看导出脚本的核心逻辑，并分块解释它的每一部分。

1. 导入模块 & 设置路径

//torch：核心框架

//train.SimpleDenoiser：从训练脚本复用模型结构

//os：用于创建输出目录

import torch
from train import SimpleDenoiser  # 模型结构
import os

2. 导出函数定义

//这个函数接收三个参数：

//pth_path: 训练得到的模型参数文件路径

//onnx_path: 导出的 ONNX 文件保存路径

//input_size: 模拟推理输入的尺寸（默认 1×3×256×256）
def export_model_to_onnx(pth_path, onnx_path, input_size=(1, 3, 256, 256)):

3. 加载模型和权重

//自动检测 CUDA 可用性，加载模型到对应设备；

//使用 load_state_dict() 加载训练好的参数；

//model.eval() 让模型切换到推理模式（关闭 Dropout/BatchNorm 更新）；
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

model = SimpleDenoiser().to(device)
model.load_state_dict(torch.load(pth_path, map_location=device))
model.eval()

4. 构造假输入（Dummy Input）

//ONNX 导出需要一个具体的输入样本，我们这里用 torch.randn 生成一个形状为 (1, 3, 256, 256) 的随机图//像；

//输入必须放在同一个设备上（GPU 或 CPU）；
dummy_input = torch.randn(*input_size).to(device)

5. 导出为 ONNX

torch.onnx.export(
    model,  //要导出的模型
    dummy_input,  //示例输入张量
    onnx_path, //	导出路径
    export_params=True,  //是否导出权重
    opset_version=11,  //ONNX 的算子集版本，通常推荐 11 或 13
    do_constant_folding=True,  //优化常量表达式，减小模型体积
    input_names=['input'],  //自定义输入输出张量的名称
    output_names=['output'],  //声明哪些维度可以变动，比如 batch size、图像大小等（部署时更灵活）
    dynamic_axes={
        'input': {0: 'batch_size', 2: 'height', 3: 'width'},
        'output': {0: 'batch_size', 2: 'height', 3: 'width'}
    }
)

6. 创建目录并调用函数

//确保输出文件夹存在，并调用导出函数生成最终模型。
if __name__ == "__main__":
    os.makedirs("onnx", exist_ok=True)
    export_model_to_onnx("weights/denoiser.pth", "onnx/denoiser.onnx")

导出后如何验证？

pip install onnxruntime

import onnxruntime
import numpy as np

sess = onnxruntime.InferenceSession("onnx/denoiser.onnx")
input = np.random.randn(1, 3, 256, 256).astype(np.float32)
output = sess.run(None, {"input": input})
print("输出 shape:", output[0].shape)