3、模型导出入门：PyTorch模型导出为ONNX（torch.onnx.export）
好，咱们今天来聊聊模型导出。说白了，就是把PyTorch训练好的模型，转成ONNX格式。这一步是模型部署的起点，也是很多坑的源头。我个人习惯在训练完模型后，第一时间就导出ONNX验证一下，而不是等到部署时才发现问题。
3.1 为什么需要导出ONNX？
你想想看，PyTorch模型虽然好用，但它是动态图。部署到生产环境时，我们需要一个静态的、标准化的中间表示。ONNX就是干这个的。它像一个通用翻译器，能把PyTorch模型翻译成TensorRT、OpenVINO、ONNX Runtime都能看懂的语言。
我在项目中遇到过，有同事直接把PyTorch模型放到生产环境推理，结果内存泄漏、推理速度慢得离谱。后来换成ONNX导出，再转成TensorRT，速度提升了3倍多。嗯，这就是标准化的力量。
3.2 torch.onnx.export 核心参数
导出函数就一个：torch.onnx.export()。但参数不少，我挑几个重点说。
| 参数名 | 说明 | 我的建议 |
|---|---|---|
| model | 要导出的PyTorch模型 | 记得先调成eval模式 |
| args | 模型的输入张量（或元组） | 形状要和实际推理时一致 |
| f | 输出ONNX文件路径 | 建议用 .onnx 后缀 |
| input_names | 输入节点的名称列表 | 起个有意义的名字，方便后续调试 |
| output_names | 输出节点的名称列表 | 同上 |
| dynamic_axes | 指定动态轴（如batch size） | 如果输入尺寸会变，这个必须配 |
| opset_version | ONNX算子集版本 | 我一般用13或14，兼容性好 |
3.3 最简单的导出示例
先来个最基础的。假设你有一个简单的CNN模型：
import torch
import torch.nn as nn
class SimpleCNN(nn.Module):
def __init__(self):
super().__init__()
self.conv = nn.Conv2d(3, 16, 3)
self.relu = nn.ReLU()
self.pool = nn.AdaptiveAvgPool2d((1, 1))
self.fc = nn.Linear(16, 10)
def forward(self, x):
x = self.conv(x)
x = self.relu(x)
x = self.pool(x)
x = x.view(x.size(0), -1)
x = self.fc(x)
return x
model = SimpleCNN()
model.eval() # 别忘了这步
# 创建一个示例输入
dummy_input = torch.randn(1, 3, 224, 224)
# 导出
torch.onnx.export(
model,
dummy_input,
"simple_cnn.onnx",
input_names=['input'],
output_names=['output'],
opset_version=13
)
print("导出成功！")
我的小技巧： 导出前一定要调用
model.eval()。我刚开始做的时候忘了这步，结果导出的模型里还带着Dropout和BatchNorm的训练行为，推理结果完全不对。折腾了一天才发现是这个问题。
3.4 处理动态输入尺寸
很多场景下，输入图片的尺寸不是固定的。比如目标检测模型，batch size可能变化，图片宽高也可能不同。这时候就需要配置 dynamic_axes。
# 假设输入是 [batch, 3, height, width]
# batch、height、width 都可能变化
torch.onnx.export(
model,
dummy_input,
"dynamic_model.onnx",
input_names=['input'],
output_names=['output'],
dynamic_axes={
'input': {0: 'batch_size', 2: 'height', 3: 'width'},
'output': {0: 'batch_size'}
},
opset_version=13
)这里要注意，dynamic_axes 是一个字典，key是节点名称，value是轴索引和轴名称的映射。轴名称可以随便起，但最好能让人一眼看懂。
我曾经踩过的坑： 动态轴配置不当，会导致ONNX模型在某些推理引擎上报错。比如TensorRT对动态shape支持有限，有些算子不支持动态尺寸。我的建议是：先固定尺寸导出，验证精度没问题后，再尝试动态尺寸。一步到位往往容易翻车。
3.5 导出后的验证
导出不是终点，验证才是。我每次导出后都会做两件事：
- 检查ONNX模型结构：用
onnx.load()加载，打印模型信息 - 对比输出精度：用相同的输入，对比PyTorch和ONNX的输出差异
import onnx
import onnxruntime as ort
import numpy as np
# 1. 检查模型结构
onnx_model = onnx.load("simple_cnn.onnx")
onnx.checker.check_model(onnx_model)
print("ONNX模型结构检查通过！")
# 2. 精度对比
# PyTorch推理
with torch.no_grad():
torch_output = model(dummy_input).numpy()
# ONNX Runtime推理
ort_session = ort.InferenceSession("simple_cnn.onnx")
ort_inputs = {ort_session.get_inputs()[0].name: dummy_input.numpy()}
ort_output = ort_session.run(None, ort_inputs)[0]
# 计算差异
diff = np.abs(torch_output - ort_output).max()
print(f"最大绝对误差：{diff:.6f}")
if diff < 1e-5:
print("精度对齐成功！")
else:
print("精度偏差较大，需要排查")
精度对齐的标准： 一般来说，最大绝对误差在1e-5以内就算对齐了。如果误差在1e-3级别，可能是某些算子精度问题。如果误差超过0.1，那基本就是导出配置有问题，需要仔细检查。
3.6 常见问题与避坑指南
- 算子不支持：有些PyTorch算子没有对应的ONNX算子。我遇到过
torch.einsum导出失败的情况。解决办法是改写模型，用ONNX支持的算子替代。 - 控制流问题：如果模型里有
if语句或for循环，导出时可能会报错。ONNX是静态图，不支持动态控制流。我的做法是：把控制流逻辑移到模型外部，或者用torch.where等向量化操作替代。 - 版本兼容性：PyTorch版本和ONNX版本要匹配。我建议用PyTorch 1.10以上版本，配合opset 13或14，兼容性最好。
嗯，模型导出这块，说白了就是多试、多验证。刚开始可能会遇到各种报错，但踩过几次坑之后，你就会发现其实套路很固定。我个人习惯把导出脚本写成一个函数，每次训练完新模型，直接调用就行，省时省力。
下一章咱们聊聊ONNX模型的优化和量化，那才是真正让模型跑得快的关键。