4、模型导出实战：从TensorFlow导出ONNX模型（tf2onnx），踩坑与兼容性处理

好，咱们今天来聊聊从TensorFlow导出ONNX模型这件事。说实话，这个环节是我在实际项目中踩坑最多的一个阶段。你想想看，TensorFlow的生态和ONNX的生态，虽然都是搞深度学习的，但底层逻辑差异其实挺大的。

我个人习惯把模型导出比作「翻译」。你有一本中文书（TensorFlow模型），想翻译成英文（ONNX模型）。翻译得好，读者能看懂；翻译得不好，那就成了「中式英语」，跑起来全是bug。

4.1 准备工作：环境与工具链

在开始之前，先把工具链搭好。我建议你直接用pip安装，干净利落。

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的版本和TensorFlow版本需要匹配。我记得有一次项目紧急，我随手装了个最新版的tf2onnx，结果跟公司的TensorFlow 2.8死活对不上，折腾了半天才发现是版本兼容性问题。

4.2 基础导出流程：从SavedModel开始

最稳妥的方式，是从TensorFlow的SavedModel格式导出。为什么呢？因为SavedModel是TensorFlow官方推荐的模型保存格式，它把模型结构、权重、签名信息都打包在一起了，tf2onnx解析起来最省事。

来看一个最简单的例子：

import tensorflow as tf
import tf2onnx

# 假设你有一个训练好的模型
model = tf.keras.models.load_model('my_model.h5')

# 先保存为SavedModel格式
tf.saved_model.save(model, 'saved_model_dir')

# 然后用tf2onnx转换
import subprocess
subprocess.run([
    'python', '-m', 'tf2onnx.convert',
    '--saved-model', 'saved_model_dir',
    '--output', 'model.onnx',
    '--opset', '15'
])

嗯，这里要注意--opset参数。opset是ONNX的操作集版本，我一般用15或16。版本太低的话，有些新算子不支持；版本太高，又怕推理引擎不兼容。说白了，选一个中间值最稳妥。

4.3 踩坑实录：那些年我遇到的兼容性问题

好，接下来是重头戏。我把自己踩过的坑整理了一下，大概分这么几类：

4.3.1 算子不兼容

这是最常见的坑。TensorFlow有些算子，ONNX里没有直接对应的。比如tf.nn.fused_batch_norm，在ONNX里可能被拆成多个基础算子。

我曾经遇到过一个模型，里面用了tf.raw_ops.ResizeNearestNeighbor，结果转换时报错说「Unsupported op」。后来查了半天，发现得用tf.image.resize替代才行。

4.3.2 动态形状问题

TensorFlow默认支持动态形状，但ONNX对动态形状的支持有限。说白了，ONNX更喜欢你把输入形状固定下来。

我记得有一次，我的模型输入是(None, 224, 224, 3)，batch size是动态的。导出时没指定形状，结果推理时各种报错。后来我加了个参数：

subprocess.run([
    'python', '-m', 'tf2onnx.convert',
    '--saved-model', 'saved_model_dir',
    '--output', 'model.onnx',
    '--opset', '15',
    '--inputs', 'input:0[1,224,224,3]'  # 固定batch size为1
])

这样就把batch size固定为1了。如果你需要动态batch，可以用--dynamic-batch参数，但说实话，我建议你尽量固定形状，省得后面推理时出幺蛾子。

4.3.3 控制流操作

TensorFlow里的tf.while_loop、tf.cond这些控制流操作，在ONNX里支持得不太好。我建议你在导出前，尽量把这些操作去掉或者替换掉。

举个例子，如果你的模型里有循环，可以考虑用tf.map_fn替代，或者干脆在预处理阶段把循环展开。

4.4 实战：一个完整的导出案例

好，咱们来一个完整的案例。假设我有一个图像分类模型，用的是MobileNetV2。

import tensorflow as tf
import numpy as np
import tf2onnx
import onnx

# 1. 加载模型
model = tf.keras.applications.MobileNetV2(weights='imagenet')

# 2. 定义输入签名（这一步很关键）
spec = (tf.TensorSpec((None, 224, 224, 3), tf.float32, name="input"),)

# 3. 转换为ONNX
output_path = "mobilenetv2.onnx"
model_proto, _ = tf2onnx.convert.from_keras(
    model, 
    input_signature=spec,
    opset=15,
    output_path=output_path
)

# 4. 验证ONNX模型
onnx_model = onnx.load(output_path)
onnx.checker.check_model(onnx_model)
print("模型验证通过！")

# 5. 用ONNX Runtime跑一下推理
import onnxruntime as ort
session = ort.InferenceSession(output_path)
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(f"推理结果shape: {result[0].shape}")

4.5 高级技巧：自定义算子映射

有些时候，你遇到的算子确实没有现成的映射。这时候就需要自定义映射了。说实话，这个操作有点复杂，我一般只在万不得已时才用。

基本思路是这样的：

import tf2onnx
from tf2onnx import utils

# 自定义一个映射函数
def my_custom_op_mapping(ctx, node, name, args):
    # 这里写你的映射逻辑
    # 把TensorFlow的op映射到ONNX的op
    pass

# 注册自定义映射
tf2onnx.utils.register_onnx_op_mapping("MyCustomOp", my_custom_op_mapping)

不过说实话，我建议你尽量别走这条路。能换op就换op，能改模型就改模型。自定义映射的维护成本太高了，而且容易出bug。

4.6 导出后的验证清单

每次导出完模型，我都会按这个清单检查一遍：

检查项	检查方法	常见问题
模型结构完整性	用Netron可视化查看	缺少某些层或连接
输入输出形状	打印ONNX模型的输入输出信息	形状不匹配或动态维度问题
数值精度	对比TensorFlow和ONNX Runtime的推理结果	精度差异超过1e-5
推理速度	用ONNX Runtime跑100次取平均	比原模型慢很多

我曾经有一次，模型导出后精度差了0.5%，查了半天发现是某个算子的计算精度设置不对。所以这个验证环节，真的不能省。

4.7 小结

从TensorFlow导出ONNX模型，说白了就是一场「翻译」工作。你既要懂TensorFlow的「语法」，又要懂ONNX的「语法」。踩坑是难免的，但只要你掌握了方法，大部分问题都能解决。

我个人建议，刚开始接触时，先从简单的模型入手，比如一个单输入单输出的分类模型。等流程跑通了，再逐步增加复杂度。别一上来就搞那种几百层的Transformer模型，那纯粹是给自己找不痛快。

好，这一章就到这里。下一章咱们聊聊PyTorch模型的导出，那又是另一番风景了。