4. TensorFlow 模型导出 ONNX：tf2onnx 工具实战

TensorFlow 转 ONNX，说白了就是给模型换个「通用格式」。我刚开始接触时也觉得挺玄乎，后来发现核心工具就一个——tf2onnx。今天咱们就把 Keras 模型和 SavedModel 的转换流程彻底捋清楚。

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

4.2 Keras 模型转换：最常用的路径

Keras 模型转换有两种主流方式：命令行和 Python API。我推荐后者，调试起来更灵活。

4.2.1 方式一：Python API 转换

import tf2onnx
import tensorflow as tf

# 构建或加载 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')
])

# 指定输入签名（关键！）
spec = (tf.TensorSpec((None, 784), tf.float32, name="input"),)

# 转换
output_path = "./model.onnx"
model_proto, _ = tf2onnx.convert.from_keras(
    model, 
    input_signature=spec, 
    opset=15,
    output_path=output_path
)
print(f"ONNX 模型已保存至: {output_path}")

这里有个细节：input_signature 必须显式指定。为什么？因为 Keras 模型在转换时需要知道输入张量的形状和数据类型。我曾经漏掉这一步，结果导出的 ONNX 模型输入形状全是 unknown，推理时直接报错。

4.2.2 方式二：命令行转换

适合快速验证：

python -m tf2onnx.convert \
    --saved-model ./saved_model \
    --output ./model.onnx \
    --opset 15 \
    --inputs input:0 \
    --outputs output:0

4.3 SavedModel 转换：生产环境首选

SavedModel 是 TF 官方推荐的部署格式。我参与的一个工业项目，模型就是用 SavedModel 格式分发的。转换时要注意 SignatureDef 的匹配。

import tf2onnx
import tensorflow as tf

# 加载 SavedModel
model = tf.saved_model.load("./saved_model")

# 获取 serving signature
signature = model.signatures["serving_default"]

# 转换
spec = (tf.TensorSpec((None, 224, 224, 3), tf.float32, name="input"),)
output_path = "./saved_model.onnx"

model_proto, _ = tf2onnx.convert.from_function(
    signature,
    input_signature=spec,
    output_path=output_path
)

嗯，这里要注意：from_function 接收的是 ConcreteFunction，不是原始的 tf.function。我刚开始搞混过，转换出来的模型推理结果全是 NaN。

4.4 转换流程全景图

下面这张图总结了 TF→ONNX 的完整转换链路：

4.5 常见转换问题与避坑指南

转换过程中，我遇到过不少「坑」。下面这张表总结了最常见的三类问题：

问题类型	典型表现	解决方案
算子不支持	转换时报错：Unsupported Op: ...	升级 opset 版本；或使用 custom_ops 注册自定义算子
输入形状不匹配	ONNX 模型输入形状为 unknown	显式指定 input_signature，包含 batch 维度
动态轴丢失	推理时 batch size 固定为 1	转换时设置 dynamic_axes 参数

4.6 精度验证：转换后的第一道防线

模型转换完，第一件事就是验证精度。我通常用这个脚本做快速比对：

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

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

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

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

# 计算误差
diff = np.abs(tf_output - ort_output).max()
print(f"最大误差: {diff:.6f}")

if diff < 1e-4:
    print("✅ 精度验证通过")
else:
    print("⚠️ 精度偏差较大，需要排查")

4.7 动态轴配置：让模型更灵活

很多场景下，我们需要模型支持动态 batch size 或动态输入尺寸。tf2onnx 通过 dynamic_axes 参数实现：

import tf2onnx
import tensorflow as tf

model = tf.keras.Sequential([
    tf.keras.layers.Conv2D(32, 3, input_shape=(None, None, 3)),
    tf.keras.layers.GlobalAvgPool2D(),
    tf.keras.layers.Dense(10)
])

spec = (tf.TensorSpec((None, None, None, 3), tf.float32, name="input"),)

# 配置动态轴
model_proto, _ = tf2onnx.convert.from_keras(
    model,
    input_signature=spec,
    output_path="./dynamic_model.onnx",
    custom_ops=None,
    extra_args={
        "dynamic_axes": {
            "input": {0: "batch_size", 1: "height", 2: "width"},
            "output": {0: "batch_size"}
        }
    }
)

配置动态轴后，ONNX 模型在推理时可以接受不同尺寸的输入。不过要注意：动态轴会增加推理框架的调度开销。如果生产环境输入尺寸固定，建议用静态形状以获得最佳性能。

4.8 转换后的模型检查

转换完成后，我习惯用 ONNX 自带的工具做一次完整性检查：

import onnx

# 加载模型
model = onnx.load("./model.onnx")

# 检查模型结构
onnx.checker.check_model(model)
print("✅ 模型结构检查通过")

# 打印输入输出信息
print("输入:", model.graph.input)
print("输出:", model.graph.output)
print("算子数:", len(model.graph.node))

这一步能发现很多隐藏问题，比如输入输出名称不匹配、形状信息丢失等。我建议把它加到 CI 流程里，每次模型更新后自动跑一遍。