当前位置：首页 > news >正文

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

news 来源：原创 2025/6/21 11:44:47

在模型训练完毕后，我们通常希望将其部署到推理平台中，比如 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 或 13do_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 npsess = 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)