4、TensorFlow 模型导出为 ONNX：tf2onnx 工具使用、Keras 模型转换、SavedModel 转换

好，咱们进入第四章。这一章聊的是 TensorFlow 模型怎么转成 ONNX。

说实话，TensorFlow 的生态有点「大而全」。从最早的 Session 模式，到后来的 Keras 高层 API，再到 SavedModel 标准格式，版本迭代快，坑也不少。我个人在项目里踩过不少 TensorFlow 转 ONNX 的雷，今天把这些经验都抖出来。

4.1 为什么需要 tf2onnx？

你想想看，TensorFlow 模型训练好了，但部署环境不一定支持 TF。比如你要在嵌入式设备上跑，或者用 ONNX Runtime 做推理加速。这时候就需要一个桥梁——tf2onnx。

说白了，tf2onnx 就是一个转换器。它把 TensorFlow 的计算图翻译成 ONNX 的计算图。我刚开始用的时候，总觉得这玩意儿就是个黑盒，后来才发现，理解它的工作原理能省很多调试时间。

4.2 环境准备

先装好必要的库。嗯，这里要注意版本兼容性。

pip install tensorflow==2.12.0
pip install onnx==1.14.0
pip install tf2onnx==1.15.1
pip install onnxruntime==1.15.1  # 用于验证转换结果

我曾经因为版本不匹配，卡了整整一个下午。TensorFlow 2.x 的版本迭代很快，tf2onnx 的适配有时候会滞后。我的建议是：先查一下 tf2onnx 的 GitHub Release 页面，看看它支持哪个 TF 版本。

4.3 方法一：Keras 模型转换

这是最直观的方式。你训练完一个 Keras 模型，保存成 .h5 文件，然后一行命令搞定。

import tensorflow as tf
import tf2onnx

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

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

# 转换
onnx_model, _ = tf2onnx.convert.from_keras(model, input_signature=spec, opset=13)

# 保存
with open("model.onnx", "wb") as f:
    f.write(onnx_model.SerializeToString())

这里有个坑：input_signature 必须指定。如果不写，tf2onnx 会尝试自动推断，但经常推断错。我遇到过 batch size 被固定成 1 的情况，推理时只能一次处理一张图。

4.4 方法二：SavedModel 转换

SavedModel 是 TensorFlow 官方推荐的导出格式。它包含了模型的计算图、权重、以及签名信息。我个人更推荐这种方式，因为它更「干净」。

# 假设你已经有了一个 SavedModel 目录
# 目录结构类似：./saved_model/  ├── assets/  ├── variables/  └── saved_model.pb

import tf2onnx

# 命令行方式（最常用）
!python -m tf2onnx.convert \
    --saved-model ./saved_model \
    --output model.onnx \
    --opset 13 \
    --inputs "input:0" \
    --outputs "output:0"

嗯，这里要注意输入输出张量的名字。你可以用 saved_model_cli show --dir ./saved_model --all 来查看。我刚开始不知道这个命令，每次都要写代码打印张量名，后来发现这个 CLI 工具太好用了。

4.5 方法三：从具体 GraphDef 转换

有些老项目用的是 frozen graph（.pb 文件）。这时候可以用 from_graphdef 方法。

import tf2onnx

# 加载 frozen graph
with tf.io.gfile.GFile("frozen_graph.pb", "rb") as f:
    graph_def = tf.compat.v1.GraphDef()
    graph_def.ParseFromString(f.read())

# 转换
onnx_graph = tf2onnx.tfonnx.process_tf_graph(
    graph_def,
    input_names=["input:0"],
    output_names=["output:0"],
    opset=13
)

# 保存
model_proto = onnx_graph.make_model("my_model")
with open("model.onnx", "wb") as f:
    f.write(model_proto.SerializeToString())

说实话，这种方法我很少用了。因为 frozen graph 里经常有一些 TF 内部算子（比如 Switch、Merge），转换起来很麻烦。如果你还在用这种老格式，我建议先升级到 SavedModel。

4.6 验证转换结果

转换完了，一定要验证。不然部署时才发现问题，那就晚了。

import onnx
import onnxruntime as ort
import numpy as np

# 检查 ONNX 模型结构
onnx_model = onnx.load("model.onnx")
onnx.checker.check_model(onnx_model)
print("ONNX 模型检查通过！")

# 用 ONNX Runtime 跑一次推理
session = ort.InferenceSession("model.onnx")
input_name = session.get_inputs()[0].name
output_name = session.get_outputs()[0].name

# 构造随机输入
dummy_input = np.random.randn(1, 224, 224, 3).astype(np.float32)

# 推理
result = session.run([output_name], {input_name: dummy_input})
print("推理结果 shape:", result[0].shape)

我一般会对比 TensorFlow 和 ONNX Runtime 的输出。如果误差在 1e-5 以内，就算通过。如果误差很大，那就要检查是不是某些算子转换出了问题。

4.7 常见问题与避坑指南

问题	原因	解决办法
转换时报错「Unsupported op」	某些 TF 算子没有对应的 ONNX 算子	升级 opset 版本，或者用自定义算子映射
推理结果不一致	可能是 Batch Normalization 的融合问题	检查 TF 和 ONNX 的 BN 层参数顺序
输入输出名字找不到	SavedModel 的签名信息不完整	用 saved_model_cli 查看具体签名
模型太大，转换超时	tf2onnx 默认有内存限制	设置环境变量 TF2ONNX_MAX_MEMORY=8G

4.8 总结

TensorFlow 转 ONNX，说白了就是三步：准备模型、选择转换方法、验证结果。

我个人最推荐 SavedModel 方式，因为它最规范，坑最少。Keras 模型转换虽然方便，但 input_signature 容易出错。至于 frozen graph，能不用就别用了。

下一章我们会聊 PyTorch 模型怎么转 ONNX。嗯，那个坑更多，但我会把经验都告诉你。