4、模型导出实战（TensorFlow篇）：tf2onnx转换流程、Keras模型转ONNX、TF与ONNX的算子映射问题
好，咱们进入实战环节。前面聊了那么多ONNX的理论和工具链，今天就来点硬核的——把TensorFlow模型转成ONNX。
说实话，TF转ONNX是我个人觉得最折腾的一环。PyTorch那边有torch.onnx.export，基本一把梭。但TF这边，嗯，历史包袱有点重。你想想看，TF 1.x的静态图、TF 2.x的Eager模式、还有Keras这种高层API，每个流派导出方式都不一样。
我在项目中遇到过最离谱的一次，一个同事用TF 1.15训练了三个月的模型，导出时发现算子不支持，最后硬是回炉重训。所以今天咱们把这套流程捋清楚，避免你踩同样的坑。
4.1 tf2onnx：最主流的转换工具
先说结论：tf2onnx是目前TF转ONNX最靠谱的工具，没有之一。它由微软维护，社区活跃度很高。
安装很简单：
pip install tf2onnx
基本用法分两步走。第一步，把TF模型保存成SavedModel格式。第二步，用tf2onnx.convert命令转换。
我习惯用Python API来调，这样方便调试：
import tf2onnx
import tensorflow as tf
# 加载你的TF模型
model = tf.saved_model.load('path/to/saved_model')
# 指定输入输出的签名
spec = (tf.TensorSpec((None, 224, 224, 3), tf.float32, name="input"),)
# 转换
onnx_model, _ = tf2onnx.convert.from_function(
model.signatures['serving_default'],
input_signature=spec,
output_path='model.onnx'
)
这里有个细节要注意：input_signature里的name参数，决定了ONNX模型输入节点的名字。我个人习惯用"input"、"image"这种直观的名字，别搞什么"Placeholder_1:0"这种TF内部命名，部署时你会后悔的。
model.signatures['serving_default'].inputs 和 .outputs 打印出来看看。我每次转换前都会先看一眼，心里有底。
4.2 Keras模型转ONNX：更简洁的路径
如果你用的是Keras（不管是TF 2.x自带的tf.keras，还是独立的Keras），那转换路径更短。
直接上代码：
import tf2onnx
import tensorflow as tf
# 假设你有一个训练好的Keras模型
model = tf.keras.models.load_model('my_keras_model.h5')
# 一行命令搞定
onnx_model, _ = tf2onnx.convert.from_keras(model, output_path='model.onnx')
是不是很爽？from_keras这个API会自动帮你处理输入输出的签名映射。我曾经用这个方法一口气转了十几个Keras模型，基本没出过问题。
但注意，这里有个前提：你的Keras模型必须是用函数式API或Sequential API构建的。如果你用了自定义层或者子类化模型（即继承tf.keras.Model），那转换时可能会报错。
为什么会这样？因为子类化模型的forward逻辑是写在call()方法里的，tf2onnx没法静态地追踪计算图。说白了，它不知道你的模型长啥样。
4.3 TF与ONNX的算子映射问题
好，到了最头疼的部分——算子映射。说白了，就是TF里的某个op，ONNX里有没有对应的实现？
我整理了一个常见的映射对照表，你直接收藏就行：
| TensorFlow算子 | ONNX算子 | 映射难度 | 常见问题 |
|---|---|---|---|
| tf.nn.conv2d | Conv | 低 | 数据格式（NHWC vs NCHW） |
| tf.nn.relu | Relu | 低 | 基本无问题 |
| tf.nn.batch_normalization | BatchNormalization | 中 | 训练/推理模式切换 |
| tf.reshape | Reshape | 低 | 动态shape处理 |
| tf.split | Split | 中 | axis参数对齐 |
| tf.gather | Gather | 中 | 索引类型转换 |
| tf.image.resize | Resize | 高 | 插值模式、坐标对齐方式 |
| tf.while_loop | Loop | 极高 | 动态循环不支持 |
看到没？大部分常见算子映射难度是低到中。但有几个坑，我重点说一下。
第一个坑：数据格式。TF默认是NHWC（batch, height, width, channels），而ONNX默认是NCHW。tf2onnx会自动帮你做transpose，但有时候会多插一个Transpose节点，影响性能。我建议你在训练时就统一用NCHW格式，省得转换后多一堆没用的op。
第二个坑：动态shape。TF里很多操作支持动态shape（比如batch size不固定），但ONNX要求输入shape是静态的。你想想看，推理引擎需要提前分配内存，你搞个动态shape，它怎么知道该分配多大？
解决办法：在转换时指定input_shapes参数，把动态维度固定下来。比如：
spec = (tf.TensorSpec((1, 224, 224, 3), tf.float32, name="input"),)
这里batch size固定为1。如果你需要支持可变batch，可以用None，但ONNX Runtime里要配合动态shape的支持，比较麻烦。
第三个坑：自定义算子。这是最头疼的。如果你的模型里用了TF独有的算子（比如tf.raw_ops.*），那tf2onnx肯定不认识。怎么办？
我有两个建议：
- 方案一：用tf2onnx的custom_op机制，自己写一个映射函数。但说实话，这需要你懂ONNX的算子规范，门槛不低。
- 方案二：在TF模型里就把这些自定义算子替换成标准算子。比如用tf.nn.conv2d代替自定义的卷积实现。这是最省事的办法。
4.4 实战：一个完整的转换流程
最后，我给你一个完整的脚本，你直接拿去用：
import tf2onnx
import tensorflow as tf
import onnx
import onnxruntime as ort
import numpy as np
# 1. 加载Keras模型
print("[INFO] 加载Keras模型...")
model = tf.keras.models.load_model('model.h5')
# 2. 转换到ONNX
print("[INFO] 转换到ONNX...")
onnx_model, _ = tf2onnx.convert.from_keras(model, output_path='model.onnx')
# 3. 验证ONNX模型
print("[INFO] 验证ONNX模型结构...")
onnx_model = onnx.load('model.onnx')
onnx.checker.check_model(onnx_model)
print("[OK] ONNX模型结构验证通过")
# 4. 精度对比
print("[INFO] 精度对比...")
# 构造随机输入
dummy_input = np.random.randn(1, 224, 224, 3).astype(np.float32)
# TF推理
tf_output = model.predict(dummy_input)
# ONNX推理
ort_session = ort.InferenceSession('model.onnx')
ort_inputs = {ort_session.get_inputs()[0].name: dummy_input}
ort_output = ort_session.run(None, ort_inputs)
# 对比
diff = np.abs(tf_output - ort_output[0]).max()
print(f"[INFO] 最大绝对误差: {diff:.6f}")
if diff < 1e-5:
print("[OK] 精度验证通过！")
else:
print("[WARN] 精度偏差较大，请检查算子映射")
这个脚本我用了好几年，每次转模型都跑一遍。尤其是精度对比那一步，千万别省。我见过太多人转完模型直接部署，结果推理结果完全不对，最后排查半天发现是某个算子映射出了问题。
嗯，今天就先聊到这儿。下一章咱们讲PyTorch篇，那个相对简单一些，但也有不少细节要注意。