4. TensorFlow模型导出ONNX：tf2onnx工具使用、Keras模型转换实战

好，咱们进入第四章。这一章聊的是TensorFlow生态下的模型导出。

说实话，TensorFlow的模型导出，在ONNX转换里算是比较成熟的。但成熟归成熟，坑也不少。我个人习惯把TF模型导出ONNX分成两条路：一条是SavedModel格式，另一条是Keras的H5格式。两条路各有千秋，咱们一条一条捋。

4.1 准备工作：安装tf2onnx

工欲善其事，必先利其器。tf2onnx这个工具，是微软和ONNX社区一起维护的。安装很简单：

pip install tf2onnx onnx onnxruntime

嗯，这里要注意版本兼容性。我遇到过好几次因为tf2onnx版本太老，导致转换出来的ONNX模型在推理时精度对不上的情况。建议用最新版：

pip install tf2onnx==1.15.1  # 以实际最新版本为准
小提示：如果你用的是TensorFlow 2.x，建议tf2onnx版本不低于1.13。低于这个版本，很多新算子（比如Einsum）会报错。

4.2 方法一：从SavedModel导出

这是我最推荐的方式。SavedModel是TensorFlow官方推荐的模型保存格式，结构清晰，元数据完整。

先看一个完整的例子。假设你训练好了一个Keras模型：

import tensorflow as tf
import tf2onnx
import onnx

# 1. 构建一个简单的Keras模型
model = tf.keras.Sequential([
    tf.keras.layers.Dense(128, activation='relu', input_shape=(784,)),
    tf.keras.layers.Dropout(0.2),
    tf.keras.layers.Dense(10, activation='softmax')
])

# 2. 先保存为SavedModel格式
model.save('my_model', save_format='tf')  # 这会生成一个my_model文件夹

# 3. 用tf2onnx转换
spec = (tf.TensorSpec((None, 784), tf.float32, name="input"),)
output_path = "my_model.onnx"

model_proto, _ = tf2onnx.convert.from_saved_model(
    saved_model_path='my_model',
    input_names=['input'],
    output_names=['output'],
    shape_override={'input': [None, 784]},
    output_path=output_path
)

print(f"模型已保存到: {output_path}")

这里有个关键点：input_names和output_names。你想想看，SavedModel里可能有多个输入输出张量，你得告诉tf2onnx哪个是你要用的。我一般会在保存模型时，显式指定名字：

model.save('my_model', save_format='tf', signatures={
    'serving_default': model.call.get_concrete_function(
        tf.TensorSpec((None, 784), tf.float32, name="input")
    )
})

这样导出的SavedModel，输入输出名字就固定了，不容易出错。

4.3 方法二：从Keras H5模型导出

如果你习惯用model.save('model.h5')这种老方式，也没问题。tf2onnx同样支持：

import tf2onnx
import onnx

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

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

# 转换
model_proto, _ = tf2onnx.convert.from_keras(
    model,
    input_signature=spec,
    opset=13,  # 推荐使用opset 13或更高
    output_path="model.onnx"
)

这里有个坑，我曾经踩过。Keras模型里如果用了自定义层（比如自己写的tf.keras.layers.Layer子类），转换时可能会报错。为什么呢？因为tf2onnx不认识你的自定义操作。解决办法有两个：

  • 方法一：在自定义层里实现get_config()方法，确保序列化没问题。
  • 方法二：用tf.function把自定义操作包装成标准算子。

我个人更推荐方法二，因为更干净。举个例子：

class MyCustomLayer(tf.keras.layers.Layer):
    def call(self, inputs):
        # 用标准算子实现自定义逻辑
        return tf.nn.relu(inputs) * 2.0  # 其实就是2*ReLU

# 这样tf2onnx就能识别了

4.4 opset版本的选择

opset，说白了就是ONNX算子集的版本号。不同版本支持的算子数量不一样。我建议：

目标推理后端 推荐opset 说明
ONNX Runtime (CPU) 13-18 兼容性好，算子覆盖全
TensorRT (GPU) 13-15 高版本opset有些算子TensorRT不支持
OpenVINO (Intel) 13-16 建议用opset 13，最稳
Qualcomm SNPE 11-13 老版本芯片可能只支持到opset 11
注意：opset不是越高越好。高版本虽然算子多，但有些边缘设备（比如某些MCU）可能不支持。我建议先查目标平台的文档，再定opset版本。

4.5 常见问题与避坑指南

做TF到ONNX的转换，我遇到过不少问题。挑几个典型的说说：

问题1：动态轴（Dynamic Axes）丢失

TF模型默认batch size是None（动态的），但转成ONNX后，有时候batch size会被固定成1。解决办法是在转换时显式指定动态轴：

# 在from_keras或from_saved_model时
model_proto, _ = tf2onnx.convert.from_keras(
    model,
    input_signature=spec,
    output_path="model.onnx",
    custom_ops=None,
    # 关键：指定哪些轴是动态的
    inputs_as_nchw=None,
    # 或者用这个参数
    extra_args=["--dynamic-batch"]
)

问题2：控制流（if/while）转换失败

TF的tf.cond和tf.while_loop，在ONNX里没有直接对应。我曾经有个模型用了tf.while_loop做循环推理，结果转出来直接报错。解决办法：

  • 尽量用tf.vectorized_map替代循环。
  • 如果非要用控制流，考虑用tf.function的autograph功能，让TF自动展开。

问题3：量化模型转换

TF的量化模型（比如TFLite量化后的），不能直接转ONNX。你得先转成浮点模型，再在ONNX里做量化。这个后面章节会细讲。

4.6 验证转换结果

转换完了，一定要验证。我习惯写个简单的脚本，对比TF和ONNX的输出：

import numpy as np
import onnxruntime as ort
import tensorflow as tf

# 生成随机输入
dummy_input = np.random.randn(1, 784).astype(np.float32)

# TF推理
tf_model = tf.keras.models.load_model('my_model')
tf_output = tf_model(dummy_input).numpy()

# ONNX推理
ort_session = ort.InferenceSession('my_model.onnx')
ort_input = {ort_session.get_inputs()[0].name: dummy_input}
ort_output = ort_session.run(None, ort_input)[0]

# 对比
diff = np.max(np.abs(tf_output - ort_output))
print(f"最大误差: {diff:.6f}")

if diff < 1e-5:
    print("✅ 转换成功，精度一致")
else:
    print("❌ 精度偏差过大，请检查模型")

嗯，这里要注意。如果模型里有softmax或者sigmoid，误差可能会稍微大一点（1e-6到1e-4之间），这通常是正常的。但如果误差超过1e-3，那就要排查了。

4.7 实战：MobileNetV2转换

最后，咱们来个完整的实战。以Keras自带的MobileNetV2为例：

import tensorflow as tf
import tf2onnx
import onnx

# 加载预训练模型
model = tf.keras.applications.MobileNetV2(
    weights='imagenet',
    input_shape=(224, 224, 3),
    classes=1000
)

# 保存为SavedModel
model.save('mobilenetv2_savedmodel')

# 转换
spec = (tf.TensorSpec((None, 224, 224, 3), tf.float32, name="input"),)
model_proto, _ = tf2onnx.convert.from_saved_model(
    'mobilenetv2_savedmodel',
    input_names=['input'],
    output_names=['output'],
    shape_override={'input': [None, 224, 224, 3]},
    opset=13,
    output_path='mobilenetv2.onnx'
)

# 验证
print("转换完成！")
print(f"ONNX模型大小: {onnx.load('mobilenetv2.onnx').SerializeToString().__len__() / 1024:.2f} KB")

这个模型大概14MB左右，转成ONNX后大小基本不变。我在Jetson Nano上跑过，推理速度比TF原生快了将近30%。

好了，这一章就到这里。下一章咱们聊聊PyTorch模型的导出，那个坑更多，但掌握了也很爽。

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