4、模型导出进阶：TensorFlow模型导出为ONNX（tf2onnx工具）

好，咱们继续往下走。

前面几章我们把PyTorch模型转ONNX聊得比较透了。这一章，咱们来聊聊TensorFlow。说实话，我在工业界看到的场景里，TensorFlow模型转ONNX的需求一点都不比PyTorch少。尤其是那些老项目，或者一些用TF Serving部署的系统，现在都想往ONNX Runtime上迁移。

嗯，今天的主角是tf2onnx这个工具。它是个开源项目，专门干这件事的。

4.1 为什么需要tf2onnx？

你可能会问：TensorFlow自己不是有tf.saved_model吗？为什么还要转ONNX？

说白了，就是为了跨平台和加速。ONNX Runtime在CPU和GPU上的优化做得相当不错，而且它支持更多的硬件后端。我有个项目，原来用TF Serving跑BERT模型，延迟一直下不来。后来转成ONNX，用ONNX Runtime跑，延迟直接降了30%。

另外，ONNX的生态更开放。你转成ONNX之后，可以无缝切换到其他推理引擎，比如TensorRT、OpenVINO等等。这个灵活性，是TensorFlow原生格式给不了的。

4.2 环境准备

开始之前，先把环境搭好。我个人习惯用pip安装，干净利落。

pip install tensorflow
pip install onnx
pip install tf2onnx

注意一下版本兼容性。我建议TensorFlow用2.x版本，tf2onnx用1.15以上。我曾经遇到过因为版本不匹配，导出时各种报错，折腾了大半天。后来发现就是tf2onnx版本太旧了。

小提示：如果你用的是TensorFlow 1.x，建议先升级到2.x。实在不行，tf2onnx也支持1.x，但功能会少一些。

4.3 基本导出流程

咱们先走一个最简单的例子。假设你有一个训练好的Keras模型，想转成ONNX。

import tensorflow as tf
import tf2onnx
import onnx

# 1. 加载或创建你的Keras模型
model = tf.keras.Sequential([
    tf.keras.layers.Dense(128, activation='relu', input_shape=(784,)),
    tf.keras.layers.Dropout(0.2),
    tf.keras.layers.Dense(10, activation='softmax')
])

# 2. 指定输入输出的名字和形状
spec = (tf.TensorSpec((None, 784), tf.float32, name="input"),)

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

# 4. 验证一下
onnx_model = onnx.load(output_path)
onnx.checker.check_model(onnx_model)
print("模型导出成功！")

你看，核心就三步：定义输入签名、调用from_keras、验证。这里有个关键点——input_signature。它告诉tf2onnx你的输入长什么样。我刚开始用的时候，经常忘记指定name参数，结果导出的ONNX里输入节点名字是默认的，后面推理时还得去查，很麻烦。

重点：一定要给输入输出节点起个有意义的名字。比如"input_ids"、"logits"之类的。这样后面做推理时，一眼就能看懂。

4.4 处理SavedModel格式

很多时候，你拿到的模型是SavedModel格式的。这也很常见，尤其是从TF Serving上下来的模型。

import tf2onnx

# 假设你的SavedModel在 ./saved_model 目录下
model_path = "./saved_model"
output_path = "./model.onnx"

# 转换
model_proto, _ = tf2onnx.convert.from_saved_model(model_path, 
                                                   output_path=output_path)

这里有个坑。SavedModel里可能包含多个签名（signature），默认情况下tf2onnx会使用默认的serving_default。如果你的模型有自定义签名，需要手动指定：

model_proto, _ = tf2onnx.convert.from_saved_model(
    model_path,
    output_path=output_path,
    signature_def="my_custom_signature"  # 指定签名
)

我记得有一次，同事导出的ONNX模型推理结果全是错的。查了半天，发现他导出的签名不对，模型输出的是中间层的特征，而不是最终的预测结果。嗯，这个细节很容易被忽略。

4.5 处理动态输入形状

TensorFlow模型经常需要处理动态batch size。比如你的模型输入是(None, 224, 224, 3)，这个None表示batch size可变。

tf2onnx默认会保留动态轴。但有时候，ONNX Runtime对动态形状的支持不是特别好，尤其是某些算子。我建议你尽量固定batch size，除非你确实需要动态。

# 固定batch size为1
spec = (tf.TensorSpec((1, 224, 224, 3), tf.float32, name="input"),)

# 或者保留动态
spec = (tf.TensorSpec((None, 224, 224, 3), tf.float32, name="input"),)
注意：如果你的模型里有Reshape、Gather等操作，动态形状可能会导致导出失败。这时候需要检查一下模型结构，或者用--inputs-as-nchw等参数调整。

4.6 常见问题与避坑指南

做这个导出，我踩过的坑真不少。挑几个典型的说说。

4.6.1 算子不支持

tf2onnx虽然支持大部分TensorFlow算子，但总有一些冷门算子没覆盖到。比如tf.raw_ops里的某些操作。

遇到这种情况，我的做法是：

  • 先查tf2onnx的官方文档，看看这个算子是否在支持列表里
  • 如果不在，考虑用tf.compat.v1里的替代算子
  • 实在不行，就自定义算子，或者修改模型结构

4.6.2 精度问题

我曾经遇到过一个模型，转成ONNX后精度下降了0.5%。排查了很久，发现是某个tf.nn.softmax在转换时被优化成了不同的实现。

解决方案：

  • 导出时加上--opset 13或更高版本，新版本的opset对精度控制更好
  • 用onnxruntime和tensorflow分别跑一遍，逐层对比输出

4.6.3 控制流问题

如果你的模型里有tf.while_loop或tf.cond，导出时可能会报错。tf2onnx对控制流的支持有限。

我的建议是：尽量把控制流逻辑移到模型外部。比如在预处理阶段做判断，而不是让模型内部去处理。

4.7 验证对齐

导出完成后，一定要做精度验证。这是我最强调的一点。

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

# 生成测试数据
test_input = np.random.randn(1, 784).astype(np.float32)

# TensorFlow推理
tf_output = model.predict(test_input)

# ONNX Runtime推理
ort_session = ort.InferenceSession("./model.onnx")
ort_input = {ort_session.get_inputs()[0].name: test_input}
ort_output = ort_session.run(None, ort_input)

# 对比
diff = np.abs(tf_output - ort_output[0])
print(f"最大误差: {diff.max()}")
print(f"平均误差: {diff.mean()}")

一般来说，误差在1e-5以内就算对齐了。如果误差很大，那就得逐层排查了。

小技巧：可以用onnxruntime的InferenceSession的get_outputs()方法，拿到所有中间层的输出，和TensorFlow的中间层对比。这样能快速定位问题出在哪一层。

4.8 总结

好了，这一章的内容就这些。总结一下：

  • tf2onnx是TensorFlow转ONNX的主流工具，支持Keras和SavedModel
  • 输入签名要明确，名字要有意义
  • 动态形状要谨慎处理，尽量固定batch size
  • 导出后一定要做精度验证，逐层对比

下一章，咱们聊聊PyTorch模型导出时的一些高级技巧。到时候见。

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