4、模型导出实战：从TensorFlow导出ONNX模型，tf2onnx工具使用

好，咱们进入实战环节。

前面讲了那么多ONNX的底层结构和原理，说实话，不亲手导一次模型，总觉得心里没底。我个人习惯是，学一个新工具，先跑通一个最简单的例子，然后再去啃那些复杂的模型。

这一章，我们就拿TensorFlow开刀。毕竟TF在工业界存量太大了，很多老项目都是TF 1.x或者2.x写的。怎么把这些宝贝模型转成ONNX，让它们能在别的框架里跑起来？这就是我们今天要解决的问题。

4.1 为什么选TensorFlow？

你想想看，PyTorch转ONNX基本是原生支持，一条torch.onnx.export就搞定了。但TensorFlow这边，情况就复杂一些。

TF的图是静态图（尤其是1.x版本），而且有各种Saver、Checkpoint、SavedModel格式。我在项目中遇到过，一个同事训练了三天的模型，最后部署的时候发现转不了ONNX，急得直跺脚。说白了，就是没搞清楚TF的模型保存格式和ONNX之间的映射关系。

所以，掌握tf2onnx这个工具，是每个TF用户的必修课。

4.2 环境准备：装好你的工具箱

嗯，这里要注意。版本问题是个大坑。我曾经因为TensorFlow和tf2onnx版本不匹配，折腾了一整个下午。

我建议你直接用以下组合，目前最稳：

# Python 3.8 - 3.10 都行
pip install tensorflow==2.12.0
pip install tf2onnx==1.15.1
pip install onnx==1.14.0
pip install onnxruntime==1.15.1  # 用来验证转换结果
⚠️ 避坑指南： 我曾经用tf2onnx 1.14版本去转TF 2.13的模型，结果报了一堆Op不支持的错。后来发现是tf2onnx还没跟上TF的更新节奏。所以，尽量保持tf2onnx和TF的版本差距不要太大。

4.3 实战一：导出Keras模型（最简单的方式）

如果你用的是TF 2.x的Keras API，那恭喜你，这是最省心的路径。

说白了，Keras模型就是一个h5文件或者SavedModel目录。我们直接拿它开刀。

4.3.1 先训练一个玩具模型

为了演示，我们搞个简单的图像分类模型。别嫌简单，核心流程是一样的。

import tensorflow as tf
import tf2onnx
import onnx

# 1. 构建一个简单的Keras模型
model = tf.keras.Sequential([
    tf.keras.layers.Conv2D(32, (3,3), activation='relu', input_shape=(28, 28, 1)),
    tf.keras.layers.MaxPooling2D((2,2)),
    tf.keras.layers.Flatten(),
    tf.keras.layers.Dense(10, activation='softmax')
])

model.compile(optimizer='adam', loss='sparse_categorical_crossentropy')

# 2. 保存模型（SavedModel格式）
model.save('my_keras_model', save_format='tf')
print("模型已保存为 SavedModel 格式")

4.3.2 用tf2onnx转换

这里有两种方式，我推荐你用命令行，因为方便集成到CI/CD流水线里。

方式一：命令行（推荐）

python -m tf2onnx.convert \
    --saved-model my_keras_model \
    --output model.onnx \
    --opset 16

方式二：Python API

import tf2onnx

# 指定输入输出的名字（Keras模型会自动推断）
spec = (tf.TensorSpec((None, 28, 28, 1), tf.float32, name="input"),)
output_path = "model.onnx"

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

# 保存
onnx.save(model_proto, output_path)
print(f"ONNX模型已保存到: {output_path}")
💡 个人经验： opset参数我一般设成16或17。opset版本越高，支持的新算子越多，但兼容性可能差一点。如果你要部署到老设备上，建议用opset 13或14。

4.4 实战二：导出TF 1.x的冻结图（Freeze Graph）

这才是真正的硬骨头。很多老项目还是TF 1.x的.pb文件。我当年接手过一个工业检测项目，模型就是这种老古董。

TF 1.x的模型通常是一个.pb文件（冻结图）加一个.pbtxt（文本描述）。转换步骤稍微多一点。

4.4.1 准备冻结图

假设你有一个frozen_graph.pb文件。首先，你得知道输入输出的节点名字。

怎么查？用saved_model_cli或者写个小脚本打印所有节点。

import tensorflow as tf

# 加载冻结图
with tf.io.gfile.GFile('frozen_graph.pb', 'rb') as f:
    graph_def = tf.compat.v1.GraphDef()
    graph_def.ParseFromString(f.read())

# 打印所有节点（只看最后几个输出节点）
with tf.Graph().as_default() as graph:
    tf.import_graph_def(graph_def, name='')
    for op in graph.get_operations():
        print(op.name, op.type)
    # 通常输出节点名字类似: "output/BiasAdd" 或 "softmax"

4.4.2 转换冻结图

python -m tf2onnx.convert \
    --input frozen_graph.pb \
    --inputs "input:0" \
    --outputs "output/BiasAdd:0" \
    --output model.onnx \
    --opset 16
⚠️ 注意： 这里的--inputs和--outputs必须和你的图里实际的名字完全一致。多一个冒号、少一个斜杠都不行。我曾经因为输出节点名字写成了output/BiasAdd（少了:0），结果转换出来的模型推理结果全是错的。

4.5 验证转换结果

模型转出来了，怎么知道对不对？

我一般做两件事：

  1. 用ONNX Runtime跑一遍推理，看结果和TF原模型是否一致。
  2. 用Netron可视化，看计算图结构是否完整。
import onnxruntime as ort
import numpy as np

# 加载ONNX模型
sess = ort.InferenceSession("model.onnx")

# 构造随机输入
input_data = np.random.randn(1, 28, 28, 1).astype(np.float32)

# 推理
outputs = sess.run(None, {"input": input_data})
print("ONNX推理结果:", outputs[0].shape)

# 对比TF结果（略，但逻辑就是比较两个输出的数值误差）
# 一般误差在1e-5以内就算成功
💡 小技巧： 如果发现数值对不上，别慌。先检查opset版本，再检查输入输出的dtype。TF默认是float32，但有些模型内部用了float16，转换时可能会丢失精度。

4.6 常见问题与避坑

我整理了几个高频问题，都是我在实际项目中踩过的坑。

问题 原因 解决方案
转换时报错：Unsupported Op tf2onnx不支持某个TF算子 升级tf2onnx版本；或者用--custom-ops自定义映射
推理结果全是NaN 输入输出名字写错，或者opset版本太低 检查节点名字；尝试opset 15以上
模型太大，转换超时 图中有循环或动态shape 用--fold_const折叠常量；或者手动指定--input_shape
转换成功但推理速度慢 ONNX Runtime没有启用优化 设置SessionOptions的graph_optimization_level

4.7 总结

好了，这一章的内容就这些。

说白了，从TensorFlow导出ONNX模型，核心就三步：

  1. 搞清楚你的模型格式（Keras h5？SavedModel？还是冻结图？）
  2. 用tf2onnx转换（命令行或Python API）
  3. 验证结果（数值对比 + 可视化）

我个人觉得，最难的不是转换本身，而是排查那些莫名其妙的错误。但只要你掌握了上面这些套路，大部分问题都能迎刃而解。

下一章，我们会讲PyTorch模型的导出。那个更简单，但也有自己的小陷阱。到时候见。

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