3、模型导出前的准备工作：模型eval模式、输入数据格式要求、设备一致性检查

好，咱们直接进入正题。

很多人拿到一个训练好的模型，二话不说就开始转ONNX。结果呢？报错、精度对不上、推理结果全是NaN。我见过太多这样的案例了。其实，导出前的准备工作做得好，后面能省掉80%的麻烦。

今天我就把这三件必须做的事，掰开揉碎了讲清楚。

3.1 模型eval模式：别让Dropout和BN坑了你

先说第一个，也是最容易被忽略的——模型模式切换。

PyTorch模型有两种模式：train()和eval()。训练时我们用model.train()，但导出ONNX时，必须用model.eval()。

为什么？

因为Dropout和Batch Normalization（BN）在两种模式下的行为完全不同。

• Dropout：训练时会随机丢弃神经元，导出时如果没关掉，每次推理结果都不一样。你想想看，ONNX图里还带着随机性，这能行吗？
• Batch Normalization：训练时用的是当前batch的均值和方差，而推理时用的是全局统计量。如果不切到eval模式，BN层会继续用batch统计量，结果就是——同一个输入，每次输出都不一样。

我个人的习惯是，在导出脚本里这样写：

model = MyModel()
model.load_state_dict(torch.load('checkpoint.pth'))
model.eval()  # 必须！必须！必须！

我曾经遇到过一个项目，模型在PyTorch里跑得好好的，转成ONNX后精度直接崩了。查了两天，最后发现是忘了切eval模式。嗯，这种低级错误，犯过一次就再也不会忘了。

3.2 输入数据格式要求：Tensor的形状、类型和设备

第二个准备工作，是搞清楚ONNX期望的输入格式。

ONNX是一个静态图，它要求输入的形状、数据类型、设备位置都是确定的。你不能像PyTorch那样灵活地传各种形状的Tensor。

3.2.1 形状（Shape）

ONNX导出时，你需要提供一个dummy input（虚拟输入）。这个dummy input的形状决定了ONNX模型的输入形状。

举个例子：

dummy_input = torch.randn(1, 3, 224, 224)  # batch_size=1, 3通道, 224x224
torch.onnx.export(model, dummy_input, "model.onnx")

这里有个坑：batch size。如果你用1作为batch size导出，那ONNX模型就只接受batch size为1的输入。如果你想支持动态batch size，需要设置dynamic_axes参数。

3.2.2 数据类型（Dtype）

ONNX支持的数据类型有限，主要是float32、float16、int32、int64等。PyTorch默认是float32，这没问题。

但如果你用了float64（双精度），或者int8、uint8，导出时可能会报错或自动转换。我个人建议：统一用float32，兼容性最好。

检查一下你的dummy input：

print(dummy_input.dtype)  # 应该是 torch.float32

3.2.3 设备一致性（Device）

这是第三个要点，也是很多人踩坑的地方。

你的模型可能在GPU上训练，但导出ONNX时，dummy input和设备必须在同一个设备上。

什么意思？

如果你的模型在GPU上：

model = model.cuda()
dummy_input = torch.randn(1, 3, 224, 224).cuda()  # 必须也在GPU上

如果你的模型在CPU上：

model = model.cpu()
dummy_input = torch.randn(1, 3, 224, 224)  # 默认在CPU上

我曾经见过一个同事，模型在GPU上，但dummy input忘了加.cuda()。结果导出成功了，但推理时ONNX模型的行为完全不对。查了半天，发现是设备不一致导致的隐式类型转换问题。

3.3 完整的导出前检查清单

好了，说了这么多，我总结一个检查清单。每次导出前，按这个顺序过一遍：

检查项	操作	常见错误
模型模式	model.eval()	忘了切模式，导致BN/Dropout行为异常
输入形状	确认dummy_input的shape	batch size不匹配，或动态轴未设置
数据类型	确认dummy_input的dtype	用了float64或int8等不常用类型
设备一致性	模型和输入在同一设备	模型在GPU，输入在CPU（反之亦然）

3.4 一个完整的示例

最后，给你看一个我常用的导出模板：

import torch
import torch.onnx

# 1. 加载模型
model = MyModel()
model.load_state_dict(torch.load('model.pth'))

# 2. 切到eval模式
model.eval()

# 3. 创建dummy input
dummy_input = torch.randn(1, 3, 224, 224)

# 4. 设备一致性检查
device = next(model.parameters()).device
dummy_input = dummy_input.to(device)
print(f"模型设备: {device}, 输入设备: {dummy_input.device}")

# 5. 导出
torch.onnx.export(
    model,
    dummy_input,
    "model.onnx",
    input_names=['input'],
    output_names=['output'],
    dynamic_axes={'input': {0: 'batch_size'}, 'output': {0: 'batch_size'}}
)

print("导出成功！")

嗯，这个模板我用了两年，基本没出过问题。你拿去用，把里面的模型名和输入形状改成你自己的就行。

记住：准备工作做得好，导出过程就是一行命令的事。准备工作没做好，后面调试能让你怀疑人生。