4. 模型训练与导出

好，咱们进入第四章。前面几章我们把 ONNX 的基础概念和工具链搭好了，现在该动真格的了——把 TensorFlow 训练出来的模型转成 ONNX。

说实话，TF 转 ONNX 在几个框架里算是最折腾的。PyTorch 有 torch.onnx.export，一步到位。TF 这边呢？版本多、API 乱、还有 SavedModel、Keras H5、Checkpoint 好几种格式。我刚开始做这个的时候，光踩版本坑就花了两天。

不过别怕。今天我就把 tf2onnx 这个工具给你讲透。你跟着走一遍，以后遇到 TF 模型导出，心里就有底了。

4.1 tf2onnx 是什么？为什么选它？

tf2onnx 是微软开源的一个转换工具。它的核心作用就一个：把 TensorFlow 的计算图翻译成 ONNX 的计算图。

你可能会问：为什么不直接用 ONNX 官方工具？原因很简单——TF 的算子太“野”了。很多自定义算子、控制流、动态 shape，ONNX 原生不认。tf2onnx 相当于一个翻译官，把这些 TF 特有的东西映射到 ONNX 标准算子上去。

核心原理一句话： tf2onnx 读取 TF 的 GraphDef 或 SavedModel，逐节点匹配算子映射表，输出 ONNX 模型。

我个人习惯用 tf2onnx 的 Python API 来做转换，而不是命令行。为什么？因为可以在代码里做预处理、后处理，调试起来方便。当然，命令行也有它的好处——快速验证。

4.2 环境准备：版本匹配是头等大事

嗯，这里我要重点强调一下。TF 和 ONNX 的版本匹配，是转换成功的第一道门槛。我在项目中遇到过好几次，明明代码没问题，就是转不过去，最后发现是版本不兼容。

下面这个表格是我实测过的稳定组合，你直接照着来就行：

TensorFlow 版本 tf2onnx 版本 ONNX 版本 备注
2.4.x 1.9.x 1.9.x 稳定组合
2.5.x - 2.8.x 1.10.x - 1.12.x 1.10.x - 1.12.x 推荐
2.9.x - 2.12.x 1.13.x - 1.15.x 1.13.x - 1.15.x 最新稳定

避坑指南： 我曾经因为 TF 2.6 配了 tf2onnx 1.8，结果转换时卡在 FusedBatchNorm 算子上，报错信息完全看不懂。后来发现是 tf2onnx 版本太老，不支持 TF 2.6 新增的算子变体。升级到 1.10 就好了。

安装命令很简单：

pip install tf2onnx onnx onnxruntime

如果你用的是 TF 2.x，tf2onnx 会自动检测并适配。但建议你装完后跑一下版本检查：

python -c "import tf2onnx; print(tf2onnx.__version__)"

4.3 方法一：Keras H5 模型转换

这是最常见的情况。你用 Keras 的 Sequential 或 Functional API 训练好模型，保存成 .h5 文件。现在要转成 ONNX。

流程其实就三步：

  1. 加载 Keras 模型：用 tf.keras.models.load_model()
  2. 指定输入签名：告诉 tf2onnx 输入张量的名字、形状和数据类型
  3. 调用转换函数：tf2onnx.convert.from_keras()

看代码：

import tensorflow as tf
import tf2onnx
import onnx

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

# 2. 定义输入签名
# 注意：batch_size 可以设为 None，表示动态 batch
spec = (tf.TensorSpec((None, 224, 224, 3), tf.float32, name="input"),)

# 3. 转换
output_path = "my_model.onnx"
model_proto, _ = tf2onnx.convert.from_keras(
    model,
    input_signature=spec,
    output_path=output_path
)

print(f"转换完成，模型保存至 {output_path}")

小技巧： input_signature 里的 name 参数很重要。这个 name 会成为 ONNX 模型的输入节点名。你在嵌入式端部署时，推理代码里要用这个名字来绑定输入数据。我一般习惯用 "input" 或 "data"，简单明了。

你想想看，为什么这里要显式指定 TensorSpec？因为 Keras 模型本身不保存输入张量的名字和 dtype 信息。如果不指定，tf2onnx 会自己猜，但猜错的情况我遇到过好几次。所以，手动指定是最稳妥的。

4.4 方法二：SavedModel 转换

SavedModel 是 TF 2.x 官方推荐的保存格式。它比 .h5 更完整——包含了模型的计算图、权重、以及 serving 相关的签名信息。

转换 SavedModel 有两种方式：命令行和 Python API。

4.4.1 命令行方式（快速验证）

python -m tf2onnx.convert \
    --saved-model ./saved_model_dir \
    --output model.onnx \
    --inputs input:0 \
    --outputs output:0

这里 --inputs 和 --outputs 指定的是张量的名字。怎么知道名字？用 saved_model_cli 工具查看：

saved_model_cli show --dir ./saved_model_dir --all

输出里会显示 SignatureDef，里面就有输入输出的 tensor_name。

4.4.2 Python API 方式（推荐）

import tf2onnx

# 直接指定 SavedModel 目录
output_path = "model_from_saved_model.onnx"
model_proto, _ = tf2onnx.convert.from_saved_model(
    saved_model_path="./saved_model_dir",
    input_names=["input:0"],
    output_names=["output:0"],
    output_path=output_path
)

关键区别： SavedModel 转换不需要加载模型到内存，tf2onnx 直接读取磁盘上的 protobuf 文件。对于超大模型（比如 1GB 以上），这种方式更省内存。

我记得有一次，客户给了一个 2.3GB 的 SavedModel，用 Keras 加载直接 OOM 了。换成 from_saved_model 方式，内存占用不到 500MB，顺利转完。

4.5 方法三：Checkpoint + GraphDef 转换

这个方法现在用得少了，但如果你维护的是 TF 1.x 的老项目，可能会遇到。流程是：从 checkpoint 恢复图结构，冻结成 frozen graph，再转 ONNX。

import tensorflow as tf
from tensorflow.python.framework import graph_io

# 1. 恢复图
with tf.Session() as sess:
    saver = tf.train.import_meta_graph('model.ckpt.meta')
    saver.restore(sess, 'model.ckpt')
    
    # 2. 冻结图
    frozen_graph = tf.graph_util.convert_variables_to_constants(
        sess,
        sess.graph_def,
        output_node_names=['output']
    )
    
    # 3. 保存 frozen graph
    graph_io.write_graph(frozen_graph, '.', 'frozen_graph.pb', as_text=False)

# 4. 用 tf2onnx 转换 frozen graph
import tf2onnx
tf2onnx.convert.from_graph_def(
    frozen_graph,
    input_names=["input:0"],
    output_names=["output:0"],
    output_path="model_from_frozen.onnx"
)

注意： TF 1.x 的图里经常有控制流操作（如 tf.cond、tf.while_loop），这些在 ONNX 里没有直接对应。tf2onnx 会尝试展开，但有时候会失败。我曾经遇到一个循环展开后节点数暴增 10 倍，推理速度反而变慢了。这种情况建议在 TF 侧先重写模型，避免复杂控制流。

4.6 三种方法对比：什么时候用哪个？

方法 适用场景 优点 缺点
from_keras Keras .h5 模型 API 简洁，代码量少 需要加载模型到内存
from_saved_model TF 2.x 标准导出格式 省内存，支持签名 需要知道 tensor name
from_graph_def TF 1.x 老项目、frozen graph 兼容旧版 流程繁琐，控制流难处理

我个人建议：新项目优先用 SavedModel。它是 TF 官方推荐格式，tf2onnx 对它的支持也最完善。Keras 转换适合快速原型验证。至于 frozen graph，除非你维护的是遗留系统，否则尽量别碰。

4.7 转换后的验证：别急着部署

模型转完了，别急着往嵌入式板子上烧。先做两件事：

  1. ONNX 模型结构检查：用 onnx.checker 验证模型是否合法
  2. 数值一致性对比：用同一份输入数据，对比 TF 和 ONNX 的输出差异

验证代码：

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

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

# 2. 准备测试数据
test_input = np.random.randn(1, 224, 224, 3).astype(np.float32)

# 3. TF 推理
tf_model = tf.keras.models.load_model('my_model.h5')
tf_output = tf_model.predict(test_input)

# 4. ONNX 推理
ort_session = ort.InferenceSession("my_model.onnx")
ort_inputs = {ort_session.get_inputs()[0].name: test_input}
ort_output = ort_session.run(None, ort_inputs)[0]

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

if diff < 1e-4:
    print("✅ 转换成功，数值一致")
else:
    print("⚠️ 差异较大，请检查模型")

经验之谈： 我一般要求差异小于 1e-4。如果差异在 1e-3 量级，可能是算子精度问题，可以接受。但如果差异超过 0.1，那肯定有问题——要么是输入输出名字搞错了，要么是某个算子映射不对。

4.8 常见问题与避坑

  • 算子不支持：tf2onnx 会报 "Unsupported Op"。解决办法是升级 tf2onnx 版本，或者用 --opset 参数指定更高的 opset 版本。
  • 动态 shape 问题：ONNX 默认支持动态 batch，但有些算子要求静态 shape。我建议在转换时把 batch 维度设为 None，其他维度固定。
  • 控制流炸了：TF 的 tf.while_loop 转 ONNX 后可能变成一堆 If 节点，推理效率极低。能避免就避免。
  • 输入输出名字：名字搞错是最常见的低级错误。转换前用 print(model.inputs) 确认一下。

4.9 本章知识体系

下面这张图总结了 TF 模型导出 ONNX 的完整流程和三种路径：

TensorFlow 模型导出 ONNX 流程 TensorFlow 模型 Keras .h5 模型 SavedModel 格式 Frozen Graph from_keras() from_saved_model() from_graph_def() ONNX 模型 验证步骤 ① onnx.checker 结构检查 → ② 数值一致性对比

从图上可以看得很清楚：不管你的 TF 模型是什么格式，最终都汇聚到 ONNX。关键就是选对转换方法，然后做好验证。

好了，第四章就到这里。你只要把 tf2onnx 的三种方法都跑一遍，以后遇到 TF 模型导出就不会慌了。记住：版本匹配、输入签名、数值验证，这三步做到位，转换成功率能到 95% 以上。

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