4、模型导出进阶：使用TensorFlow/Keras导出ONNX模型（tf2onnx），处理动态轴与批处理

好，咱们接着聊模型导出。前面几章我们主要讲了 PyTorch 的导出，但实际工作中，TensorFlow 和 Keras 的用户群体也非常庞大。我自己的团队里，就有好几个同事是 TF 的忠实用户。说实话，TF 的导出流程比 PyTorch 要稍微绕一点，但掌握了套路，其实也挺顺手的。

今天我们就来啃这块硬骨头——用 tf2onnx 把 TensorFlow/Keras 模型转成 ONNX，并且重点解决动态轴和批处理的问题。嗯，这里要注意，动态轴处理不好，模型部署到生产环境时，遇到不同大小的输入就会直接崩掉。

4.1 为什么需要 tf2onnx？

TensorFlow 官方其实提供了 TF-TRT、TFLite 等部署方案，但 ONNX 的优势在于它的「中间人」角色。你想想看，一个模型如果转成 ONNX，就能在 PyTorch、TensorRT、OpenVINO 甚至 Apple CoreML 之间自由切换。说白了，ONNX 就是深度学习模型的通用语言。

我个人习惯用 tf2onnx 这个工具，因为它对 Keras 和 TF2 的支持比较完善。我记得最早接触它的时候，还是 TF1.x 时代，那时候踩了不少坑。现在 TF2 出来之后，整个流程简化了很多。

4.2 环境准备与安装

在开始之前，先把环境搭好。我建议你创建一个干净的虚拟环境，避免依赖冲突。

# 创建虚拟环境
conda create -n tf2onnx_env python=3.8
conda activate tf2onnx_env

# 安装 TensorFlow 和 tf2onnx
pip install tensorflow==2.10.0
pip install tf2onnx==1.14.0
pip install onnxruntime  # 用于验证导出结果

⚠️ 版本兼容性提醒

我曾经因为 TensorFlow 版本和 tf2onnx 版本不匹配，折腾了一整天。建议使用 TensorFlow 2.8~2.12 之间的版本，配合 tf2onnx 1.13~1.15。太新的 TF 版本（比如 2.15+）可能会有一些算子兼容性问题。

4.3 基础导出流程：从 Keras 模型到 ONNX

我们先从一个最简单的例子开始。假设你训练了一个 Keras 的 Sequential 模型，用来做 MNIST 手写数字识别。

import tensorflow as tf
import tf2onnx
import onnx

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

# 2. 编译并训练（这里省略训练过程，假设模型已经训练好）
# model.compile(...)
# model.fit(...)

# 3. 导出为 ONNX
spec = (tf.TensorSpec((None, 28, 28), tf.float32, name="input"),)
output_path = "./mnist_model.onnx"

model_proto, _ = tf2onnx.convert.from_keras(
    model, 
    input_signature=spec, 
    output_path=output_path
)

print(f"模型已导出到: {output_path}")

这里有几个关键点：

  • input_signature：这是 tf2onnx 的「眼睛」，它告诉转换器输入的形状和数据类型。注意我用了 None 作为 batch size，这就是动态轴的起点。
  • name="input"：给输入节点起个名字，方便后续推理时引用。

4.4 处理动态轴：让模型支持可变 batch size

为什么动态轴这么重要？我举个例子：你在训练时用的 batch size 是 32，但部署到生产环境后，客户端可能一次只发来 1 张图片，或者并发请求时一次发来 64 张。如果模型固定了 batch size，那就只能重新训练或者做各种 hack。

在 tf2onnx 中，处理动态轴其实很简单——你只需要在 input_signature 中把对应维度设为 None 即可。但这里有个坑：不是所有算子都天然支持动态轴。

💡 核心原则

在定义 input_signature 时，把 batch size 维度设为 None，ONNX 模型就会自动生成动态轴。但要注意，模型内部的某些操作（比如 Reshape、Flatten）可能会破坏动态性。

我们来看一个更复杂的例子——处理多输入、多输出的模型，并且所有输入都支持动态 batch。

import tensorflow as tf
import tf2onnx

# 构建一个多输入模型
input_a = tf.keras.Input(shape=(224, 224, 3), name="image")
input_b = tf.keras.Input(shape=(10,), name="metadata")

x = tf.keras.layers.Conv2D(32, 3, activation='relu')(input_a)
x = tf.keras.layers.GlobalAveragePooling2D()(x)
x = tf.keras.layers.Concatenate()([x, input_b])
output = tf.keras.layers.Dense(1, activation='sigmoid', name="score")(x)

model = tf.keras.Model(inputs=[input_a, input_b], outputs=output)

# 定义动态输入签名
spec = [
    tf.TensorSpec((None, 224, 224, 3), tf.float32, name="image"),
    tf.TensorSpec((None, 10), tf.float32, name="metadata")
]

# 导出
model_proto, _ = tf2onnx.convert.from_keras(
    model,
    input_signature=spec,
    output_path="./multi_input_model.onnx"
)

# 验证动态轴
onnx_model = onnx.load("./multi_input_model.onnx")
for inp in onnx_model.graph.input:
    print(f"输入: {inp.name}, 形状: {[dim.dim_value if dim.dim_value else 'dynamic' for dim in inp.type.tensor_type.shape.dim]}")

运行这段代码，你会看到输出中 batch size 维度显示为 dynamic。嗯，这就是我们想要的效果。

4.5 处理批处理：固定 batch vs 动态 batch 的取舍

有些场景下，你可能希望固定 batch size。比如你的推理服务是批量处理的，一次固定处理 32 张图。固定 batch 的好处是：推理性能更稳定，内存分配更可控。

我个人建议：

  • 在线推理（实时响应）：使用动态 batch，因为请求量不可预测。
  • 离线批处理：使用固定 batch，可以做一些算子融合优化。

固定 batch 的导出方式很简单：

# 固定 batch size 为 32
spec = (tf.TensorSpec((32, 28, 28), tf.float32, name="input"),)
model_proto, _ = tf2onnx.convert.from_keras(model, input_signature=spec, output_path="./fixed_batch.onnx")

📌 小技巧

如果你不确定该用动态还是固定，我建议先导出动态版本。因为 ONNX Runtime 支持动态 batch，而且性能损失很小。我曾经在项目中做过对比，动态 batch 比固定 batch 慢不到 5%，但灵活性高了很多。

4.6 常见问题与避坑指南

这部分是我最想分享的。我在实际项目中遇到过不少坑，这里挑几个典型的说说。

4.6.1 问题一：Keras 的 Input Layer 名称冲突

如果你在 Keras 中定义了多个 Input 层，但没有显式指定 name，tf2onnx 可能会报错。我曾经遇到一个模型，两个输入都叫 "input_1"，结果导出时直接崩溃。

解决方案：给每个 Input 层显式命名，并且在 input_signature 中使用相同的名字。

4.6.2 问题二：自定义层不支持导出

如果你在 Keras 中写了自定义层（继承自 tf.keras.layers.Layer），tf2onnx 可能无法自动转换。这时候需要手动注册算子映射。

我记得有一次，团队里有人写了一个自定义的 Attention 层，导出时各种报错。最后我们不得不把自定义层改写成标准 Keras 层的组合，才顺利导出。

4.6.3 问题三：动态轴导致 Reshape 出错

当 batch size 是动态时，某些 Reshape 操作会出问题。比如你写了个 Reshape((-1, 64))，ONNX 可能无法正确推断形状。

解决方案：使用 tf.reshape 时，尽量用 tf.shape 动态获取维度信息，而不是硬编码。

# 错误的做法
x = tf.reshape(x, (-1, 64))

# 正确的做法
batch_size = tf.shape(x)[0]
x = tf.reshape(x, (batch_size, -1, 64))

4.7 验证导出结果

导出完成后，一定要验证。我习惯用 ONNX Runtime 跑一遍推理，对比原始 Keras 模型的输出。

import onnxruntime as ort
import numpy as np

# 加载 ONNX 模型
ort_session = ort.InferenceSession("./mnist_model.onnx")

# 构造输入数据（batch size = 1）
input_data = np.random.randn(1, 28, 28).astype(np.float32)

# 运行推理
onnx_output = ort_session.run(None, {"input": input_data})

# 对比 Keras 输出
keras_output = model.predict(input_data)

# 检查误差
print(f"最大误差: {np.max(np.abs(onnx_output[0] - keras_output))}")

如果误差在 1e-5 以内，基本就稳了。如果误差很大，可能是某些算子转换出了问题，需要检查 tf2onnx 的日志。

4.8 总结与建议

好了，这一章的内容就到这里。我们来快速回顾一下：

  • 使用 tf2onnx.convert.from_keras 导出 Keras 模型
  • 通过 input_signature 中的 None 实现动态 batch
  • 固定 batch 和动态 batch 各有适用场景，按需选择
  • 注意自定义层、Reshape 操作等常见坑点

下一章我们会深入 ONNX 的算子层面，看看如何手动修改 ONNX 图结构。如果你在导出过程中遇到任何问题，欢迎在课程群里交流。我个人觉得，模型导出这个环节，多踩几次坑就熟练了——嗯，毕竟我也是这么过来的。

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