在模型训练完毕后，我们通常希望将其部署到推理平台中，比如 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)
模型预览：
总结
导出 ONNX 模型的流程主要包括：
-
加载模型结构 + 权重
-
准备 dummy 输入张量
-
调用
torch.onnx.export()进行导出 -
设置
dynamic_axes可变尺寸以增强部署适配性
这套流程适用于大部分视觉模型（分类、去噪、分割等），也是后续进行 TensorRT 推理或移动端部署的基础。