4、TensorFlow模型导出ONNX：tf2onnx工具使用、Keras模型转换、SavedModel转换、TF与ONNX的op映射差异

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

说实话，TF转ONNX比PyTorch要麻烦一些。为什么？因为TF的图结构太灵活了，各种控制流、自定义op满天飞。我早期做这个的时候，经常被各种奇奇怪怪的报错搞得头大。

但别怕，有套路。核心工具就是tf2onnx。咱们一步步来。

4.1 tf2onnx工具：你的主力军

tf2onnx是微软开源的工具。我个人习惯用它的命令行版本，简单直接。当然，Python API也行，看个人喜好。

安装很简单：

pip install tf2onnx

基本用法就一行命令：

python -m tf2onnx.convert \
    --saved-model ./saved_model_dir \
    --output model.onnx \
    --opset 13

这里--opset我建议至少用13。为什么？因为opset 13支持很多新算子，兼容性更好。我踩过坑，用opset 11转出来的模型，在ONNX Runtime上跑起来性能差一截。

4.2 Keras模型转换：H5格式的处理

很多人喜欢用Keras的.h5格式保存模型。嗯，这个格式转ONNX也简单。

假设你有一个训练好的model.h5：

python -m tf2onnx.convert \
    --keras model.h5 \
    --output model.onnx

但这里有个坑。Keras模型有时候会包含自定义层。比如我自己写过一个AttentionLayer，直接转就会报错。

怎么办？两种办法：

• 方法一： 把自定义层用标准Keras层重写。这是最稳妥的。
• 方法二： 在转换时指定--custom-ops参数，告诉tf2onnx怎么处理。

我个人推荐方法一。因为自定义op在ONNX里很难找到对应的实现，后面部署时也会出问题。

4.3 SavedModel转换：生产环境的标准姿势

SavedModel是TensorFlow官方推荐的格式。我建议所有生产环境的TF模型，都先导出成SavedModel，再转ONNX。

为什么？因为SavedModel包含了完整的图结构、变量、签名信息。转ONNX时信息最完整。

转换命令：

python -m tf2onnx.convert \
    --saved-model ./saved_model \
    --output model.onnx \
    --tag serve \
    --signature_def serving_default

这里--tag serve和--signature_def serving_default是啥意思？

说白了，SavedModel里可以存多个图（比如训练图、推理图）。--tag serve告诉工具我要用推理的那个图。--signature_def指定输入输出的签名。

我遇到过有人忘了指定签名，结果转出来的ONNX输入输出名字全是乱的。嗯，那排查起来可费劲了。

4.4 TF与ONNX的op映射差异：避坑指南

这是最核心的部分。TF和ONNX的算子不是一一对应的。有些TF的op，ONNX里根本没有。

我整理了一个常见差异表：

TF算子	ONNX映射	注意事项
tf.nn.conv2d	Conv	数据格式问题：TF默认NHWC，ONNX默认NCHW。tf2onnx会自动处理，但你要确认一下。
tf.nn.batch_normalization	BatchNormalization	训练模式和推理模式要区分。转ONNX时一定要用推理模式。
tf.split	Split	轴（axis）的索引方式不同。TF是0-based，ONNX也是0-based，但有些版本有差异。
tf.gather	Gather	这个映射比较稳定，但注意indices的数据类型必须是int64。
tf.reshape	Reshape	允许-1维度，ONNX也支持。但TF的reshape更灵活，有些动态形状会出问题。
tf.cond / tf.while_loop	If / Loop	这是大坑。TF的控制流在ONNX里支持有限。我建议尽量避免。

我曾经接手过一个项目，模型里用了大量的tf.while_loop做动态序列处理。转ONNX时直接报错，说Loop算子不支持。最后没办法，只能把模型结构改了，用固定长度的序列替代。

所以，你想想看，如果一开始就规划好，能省多少事？

4.5 实战：一个完整的转换流程

咱们来个完整的例子。假设你有一个Keras模型：

import tensorflow as tf
import tf2onnx

# 1. 创建或加载模型
model = tf.keras.Sequential([
    tf.keras.layers.Conv2D(32, (3,3), activation='relu', input_shape=(224,224,3)),
    tf.keras.layers.MaxPooling2D((2,2)),
    tf.keras.layers.Flatten(),
    tf.keras.layers.Dense(10, activation='softmax')
])

# 2. 先保存为SavedModel
model.save('./my_model', save_format='tf')

# 3. 转ONNX
python -m tf2onnx.convert \
    --saved-model ./my_model \
    --output my_model.onnx \
    --opset 13

转完之后，记得验证一下：

import onnx
onnx_model = onnx.load('my_model.onnx')
onnx.checker.check_model(onnx_model)
print("模型验证通过！")

这一步不能省。我见过有人转完直接部署，结果推理结果全是错的。一查，原来是某个op映射错了。

好了，这一章就到这里。TF转ONNX的核心就是：用SavedModel、选对opset、避开控制流。记住这三点，能少踩90%的坑。

下一章咱们聊PyTorch转ONNX，那个相对简单一些，但也有不少细节要注意。