3、模型导出入门：使用PyTorch导出ONNX模型（torch.onnx.export），理解导出参数含义。

好，咱们今天来聊聊模型导出。说白了，就是把PyTorch训练好的模型，转成ONNX格式。这一步是部署的起点，也是很多坑的起点。我个人习惯，在写训练代码时就把导出逻辑想好，不然后面改起来很痛苦。

3.1 为什么需要导出ONNX？

你想想看，PyTorch模型跑在Python环境里，依赖一大堆库。到了生产环境，比如用C++推理、在手机上跑、或者用TensorRT加速，PyTorch就不太方便了。ONNX就像一个中间语言，把模型的计算图标准化，让不同框架和硬件都能读懂。

我在项目中遇到过，客户要求用C++部署，但模型是PyTorch训练的。当时我直接导出ONNX，然后用ONNX Runtime加载，问题就解决了。嗯，这里要注意，导出时有些细节没处理好，推理结果可能对不上。

3.2 torch.onnx.export 核心参数

这个函数是导出的核心。我刚开始用的时候，参数记不全，经常翻文档。后来总结了一下，其实就几个关键点。

参数名	类型	说明	我的建议
model	torch.nn.Module	要导出的模型	记得设成eval模式
args	tuple	模型的输入张量	形状和类型要跟实际一致
f	str	输出ONNX文件路径	用.onnx后缀
export_params	bool	是否导出模型参数	默认True，别改
opset_version	int	ONNX算子集版本	建议用14或以上
do_constant_folding	bool	是否做常量折叠优化	推理时建议True
input_names	list[str]	输入节点名称	方便后续调试
output_names	list[str]	输出节点名称	同上
dynamic_axes	dict	动态轴定义	处理变长输入时必用

3.3 一个完整的导出示例

咱们直接看代码。假设我有一个简单的分类模型，输入是3x224x224的图片。

import torch
import torch.onnx

# 1. 准备模型
model = MyClassifier()
model.eval()  # 重要！一定要设成eval模式
checkpoint = torch.load('model.pth')
model.load_state_dict(checkpoint['state_dict'])

# 2. 准备输入
dummy_input = torch.randn(1, 3, 224, 224)

# 3. 导出
torch.onnx.export(
    model,
    dummy_input,
    'model.onnx',
    export_params=True,
    opset_version=14,
    do_constant_folding=True,
    input_names=['input'],
    output_names=['output'],
    dynamic_axes={
        'input': {0: 'batch_size'},
        'output': {0: 'batch_size'}
    }
)

print("导出成功！")

这段代码我用了很多次。注意model.eval()，如果不加，BatchNorm和Dropout的行为会不一样，导出的模型推理结果可能不对。我曾经因为这个坑，排查了一整天。

3.4 参数详解与避坑指南

3.4.1 opset_version 怎么选？

这个参数决定了ONNX算子集的版本。版本越高，支持的算子越多，但兼容性可能差一些。我个人习惯用14，因为大部分推理框架都支持。如果你用了很新的PyTorch算子，可能需要更高版本。

3.4.2 dynamic_axes 处理变长输入

很多场景下，输入batch size是变化的，或者序列长度不固定。这时候就要用dynamic_axes。说白了，就是告诉ONNX哪些维度是动态的。

dynamic_axes = {
    'input': {0: 'batch_size', 2: 'height', 3: 'width'},
    'output': {0: 'batch_size'}
}

这样导出的模型，输入可以是任意batch size，图片尺寸也可以变。但要注意，模型本身要支持这种动态输入，比如全连接层需要处理不同大小的特征。

3.4.3 do_constant_folding 优化

这个参数默认是False，但我建议设成True。它会提前计算一些常量表达式，比如权重和偏置的运算，减少推理时的计算量。嗯，这里要注意，如果模型里有动态控制流，比如if语句，常量折叠可能会出问题。

3.5 导出后的验证

导出完别急着用，先验证一下。我一般会做两件事：

1. 检查模型结构：用Netron打开ONNX文件，看看计算图对不对。
2. 对比推理结果：用ONNX Runtime加载模型，输入同样的数据，对比输出差异。

import onnxruntime as ort

# 加载ONNX模型
session = ort.InferenceSession('model.onnx')

# 准备输入
input_name = session.get_inputs()[0].name
output_name = session.get_outputs()[0].name
input_data = torch.randn(1, 3, 224, 224).numpy()

# 推理
result = session.run([output_name], {input_name: input_data})

# 对比PyTorch结果
with torch.no_grad():
    torch_result = model(torch.from_numpy(input_data)).numpy()

# 检查差异
diff = np.abs(result[0] - torch_result).max()
print(f"最大差异：{diff}")

如果差异在1e-5以内，基本没问题。如果差异很大，就要检查导出参数或者模型本身了。

3.6 常见问题与解决方案

• 导出报错：Unsupported operator：说明某个PyTorch算子没有对应的ONNX实现。可以尝试升级opset_version，或者用torch.onnx.register_custom_op注册自定义算子。
• 推理结果对不上：检查模型是否在eval模式，输入数据是否预处理一致。我曾经遇到过，训练时用了数据增强，导出时忘了关，结果对不上。
• 动态轴不生效：确认dynamic_axes的key跟input_names一致。另外，模型内部如果有reshape操作，要确保它支持动态形状。

好了，这一章就到这里。记住，导出ONNX不是终点，而是部署的起点。多验证、多测试，才能保证线上不出问题。下一章咱们聊聊ONNX模型的结构，看看里面到底长什么样。