2. 模型导出：PyTorch模型转ONNX、TensorFlow模型转ONNX、Keras模型转ONNX、ONNX模型验证

模型导出这一步，说白了就是给模型换个"身份证"。你训练好的模型在PyTorch里跑得欢，但到了部署环境，可能就得用ONNX这个通用格式。我最早接触ONNX时也觉得多此一举，直到有一次需要在边缘设备上跑模型，才发现没有ONNX真是寸步难行。

嗯，咱们直接进入正题。今天要讲的是四种主流框架的导出方法，以及导出后怎么验证它没坏。

2.1 PyTorch模型转ONNX

PyTorch转ONNX，我个人觉得是最顺手的。官方提供了torch.onnx.export这个接口，用起来很直接。

核心步骤：

  • 准备一个模型实例，设为eval模式
  • 构造一个 dummy input（形状要和真实输入一致）
  • 调用 export 函数，指定输出路径和 opset 版本
import torch
import torchvision.models as models

model = models.resnet18(pretrained=True)
model.eval()

dummy_input = torch.randn(1, 3, 224, 224)

torch.onnx.export(
    model,
    dummy_input,
    "resnet18.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'}}
)

这里有个坑，我踩过好几次。opset_version 别选太新，也别太旧。我个人习惯用11或12，兼容性最好。选15以上有些推理引擎可能不支持。

小技巧：dynamic_axes 参数很重要。如果你需要动态batch，一定要加上。否则导出的模型只能接受固定batch size。

2.2 TensorFlow模型转ONNX

TensorFlow转ONNX，我建议用tf2onnx这个工具。它支持SavedModel和Keras H5两种格式。

为什么会推荐这个工具？因为官方tf-onnx转换器有时候会出一些奇奇怪怪的算子兼容问题。tf2onnx相对来说更成熟一些。

import tf2onnx
import tensorflow as tf

# 假设你有一个SavedModel
model = tf.saved_model.load('path/to/saved_model')

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

# 转换
output_path = model.signatures['serving_default']
onnx_model, _ = tf2onnx.convert.from_function(
    output_path,
    input_signature=spec,
    opset=13,
    output_path="model.onnx"
)

我曾经在转换一个MobileNetV3时，遇到了一个"FusedBatchNormV3"算子不兼容的问题。后来发现是TensorFlow版本和tf2onnx版本不匹配。解决办法很简单：升级tf2onnx到最新版，或者降级TensorFlow到2.6左右。

注意：TensorFlow 2.x 的 SavedModel 导出时，一定要用 tf.function 装饰过的函数。否则转换出来的ONNX可能缺少输入输出名称。

2.3 Keras模型转ONNX

Keras模型转ONNX，其实和TensorFlow那套差不多。但如果你用的是纯Keras（不是tf.keras），那就要用keras2onnx这个库。

不过说实话，现在大部分人都用tf.keras了。所以直接用tf2onnx也能搞定。

import tf2onnx
import tensorflow as tf

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

# 转换
onnx_model, _ = tf2onnx.convert.from_keras(
    model,
    input_signature=None,
    opset=13,
    output_path="keras_model.onnx"
)

这里有个细节：如果你的Keras模型里有自定义层，转换前一定要确保自定义层已经注册了。否则转换时会报"Unknown layer"的错误。我建议在加载模型前，先把自定义层类导入进来。

避坑指南：我曾经在转换一个带Lambda层的Keras模型时，发现Lambda层里的函数无法序列化。后来改用自定义层继承tf.keras.layers.Layer，才解决了问题。

2.4 ONNX模型验证

模型导出后，千万别直接拿去用。一定要验证一下。我见过太多人导出后不检查，结果推理时发现精度不对。

验证分三步：

  1. 结构验证：用onnx.checker检查模型结构是否合法
  2. 数值验证：用onnxruntime跑一遍，对比原始框架的输出
  3. 形状验证：检查输入输出形状是否符合预期
import onnx
import onnxruntime as ort
import numpy as np

# 1. 结构验证
onnx_model = onnx.load("model.onnx")
onnx.checker.check_model(onnx_model)
print("模型结构验证通过")

# 2. 数值验证
ort_session = ort.InferenceSession("model.onnx")

# 构造输入
input_name = ort_session.get_inputs()[0].name
input_data = np.random.randn(1, 3, 224, 224).astype(np.float32)

# 推理
outputs = ort_session.run(None, {input_name: input_data})
print(f"ONNX输出形状: {outputs[0].shape}")

# 3. 对比原始框架输出（以PyTorch为例）
# 假设你还有原始PyTorch模型
# torch_output = torch_model(torch.from_numpy(input_data))
# np.testing.assert_allclose(outputs[0], torch_output.numpy(), rtol=1e-3, atol=1e-5)

你想想看，如果数值对不上，那部署上去就是灾难。我建议验证时用真实数据，别只用随机数。随机数有时候能过，但真实数据就不一定了。

个人经验：验证时最好用3-5组不同数据，包括边界情况（比如全0、全1、大数值）。这样能发现一些隐藏的数值溢出问题。

2.5 常见问题与解决方案

问题 原因 解决方案
算子不支持 ONNX opset版本太低 升级opset版本，或替换为支持的算子
动态轴失效 dynamic_axes未正确设置 检查输入输出名称是否匹配
精度下降 FP16转换或量化导致 使用FP32验证，再逐步降精度
转换失败 自定义算子未注册 使用onnx-simplifier简化模型

嗯，到这里模型导出这部分就差不多了。记住一个原则：导出后一定要验证，验证通过再部署。别嫌麻烦，这一步能省掉你后面90%的排查时间。

下一章我们会讲ONNX模型的优化和量化，到时候会用到今天导出的这些模型文件。

← 上一章 📖 返回目录 下一章 →