模型导出与验证：从PyTorch/TensorFlow导出ONNX、使用onnxruntime验证模型正确性、常见导出错误排查

这一章，我们来聊聊模型导出。说白了，就是把你在PyTorch或TensorFlow里训好的模型，转成ONNX这种中间格式。为什么要这么干？因为ONNX是跨平台的，能跑在各种各样的推理引擎上。我个人习惯是，模型训完第一件事就是导出ONNX，然后立刻验证。别等到部署时才发现问题，那可就晚了。

从PyTorch导出ONNX

PyTorch导出ONNX，主要靠 torch.onnx.export 这个函数。嗯，这里要注意，导出时模型得设成 eval 模式，不然BatchNorm和Dropout这些层会捣乱。

import torch
import torch.onnx

# 假设你有一个训练好的模型
model = MyModel()
model.eval()

# 创建一个 dummy input，形状要和实际输入一致
dummy_input = torch.randn(1, 3, 224, 224)

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

这里有个关键点：dynamic_axes。如果你希望模型支持动态batch size，就得配这个参数。我在项目中遇到过，一开始没配，结果部署时发现只能跑固定batch size，又得重新导出，挺折腾的。

从TensorFlow导出ONNX

TensorFlow这边，得用 tf2onnx 这个工具。流程稍微绕一点，但也不复杂。

import tensorflow as tf
import tf2onnx

# 加载你的TensorFlow模型
model = tf.keras.models.load_model('my_model.h5')

# 指定输入签名
spec = (tf.TensorSpec((None, 224, 224, 3), tf.float32, name="input"),)

# 导出
output_path = "model.onnx"
model_proto, _ = tf2onnx.convert.from_keras(
    model,
    input_signature=spec,
    opset=13,
    output_path=output_path
)

你想想看，TensorFlow的图结构比PyTorch复杂，所以导出时更容易出问题。我建议导出后立刻用 onnx.checker 检查一下。

使用onnxruntime验证模型正确性

模型导出来了，怎么知道它对不对？跑一下推理，跟原始框架的结果对比一下。这就是验证的核心思路。

import onnxruntime as ort
import numpy as np

# 创建推理会话
ort_session = ort.InferenceSession("model.onnx")

# 准备输入数据
input_data = np.random.randn(1, 3, 224, 224).astype(np.float32)

# 运行推理
ort_inputs = {ort_session.get_inputs()[0].name: input_data}
ort_outputs = ort_session.run(None, ort_inputs)

# 对比结果
# 假设你在PyTorch里也跑了一次
torch_output = model(torch.from_numpy(input_data)).detach().numpy()

# 计算误差
diff = np.abs(ort_outputs[0] - torch_output).max()
print(f"最大误差: {diff}")

if diff < 1e-5:
    print("验证通过！")
else:
    print("验证失败，需要排查问题")

常见导出错误排查

导出过程中，你大概率会遇到各种报错。我把常见的几种列出来，方便你对照排查。

错误类型	典型报错信息	原因与解决方案
算子不支持	Unsupported operator: aten::xxx	PyTorch的某些算子ONNX不支持。可以换用等价算子，或者升级opset版本。我遇到过 aten::meshgrid 不支持，换成 torch.stack 加广播就解决了。
动态shape问题	RuntimeError: Only tensors with static shape are supported	模型里有动态shape的操作，比如 torch.where 或 torch.nonzero。建议用 torch.ones_like 或 torch.zeros_like 替代。
输入输出命名冲突	ValueError: Duplicate name in graph	多个输入或输出用了相同的名字。检查 input_names 和 output_names，确保唯一。
TensorFlow自定义层	AttributeError: 'CustomLayer' object has no attribute 'get_config'	自定义层没有实现 get_config 方法。继承 tf.keras.layers.Layer 并实现序列化接口即可。

验证通过后的检查清单

验证通过不代表万事大吉。我建议你对照这个清单再检查一遍：

• 数值精度：最大误差是否在可接受范围内？通常1e-5以内算正常。
• 动态batch：用不同batch size跑一下，看结果是否一致。
• 多输入多输出：如果你的模型有多个输入输出，每个都要单独验证。
• 量化模型：如果是量化后的模型，误差会大一些，但趋势应该一致。
• 性能基线：记录一下ONNX的推理耗时，跟原始框架对比，确保没有性能退化。

嗯，这一章的内容就这些。导出和验证是模型部署的第一步，也是最容易踩坑的一步。你想想看，如果模型导出来就是错的，后面所有的优化工作都白费了。所以，花点时间把这一步做扎实，绝对值得。