4. TensorFlow模型转换为ONNX

任务目标

将训练好的TensorFlow SavedModel格式模型转换为ONNX格式，并验证转换后模型的推理输出与原模型一致。本案例以MNIST手写数字分类模型为例，演示完整的转换与验证闭环。

环境依赖

• tensorflow 2.x
• tf2onnx
• onnxruntime
• numpy

安装命令：pip install tensorflow tf2onnx onnxruntime numpy

实现步骤

首先训练一个简单的MNIST分类模型并保存为SavedModel格式，然后使用tf2onnx将其转换为ONNX，最后分别用TensorFlow和ONNX Runtime加载模型对同一输入进行推理，对比输出差异。

关键完整代码

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

# 1. 训练并保存一个简单的MNIST模型（SavedModel格式）
(x_train, y_train), (x_test, y_test) = tf.keras.datasets.mnist.load_data()
x_train = x_train.reshape(-1, 784).astype(np.float32) / 255.0
x_test = x_test.reshape(-1, 784).astype(np.float32) / 255.0

model = tf.keras.Sequential([
    tf.keras.layers.Dense(128, activation='relu', input_shape=(784,)),
    tf.keras.layers.Dense(10, activation='softmax')
])
model.compile(optimizer='adam', loss='sparse_categorical_crossentropy', metrics=['accuracy'])
model.fit(x_train, y_train, epochs=1, batch_size=128, verbose=0)

# 保存为SavedModel
model.save('mnist_savedmodel')

# 2. 使用tf2onnx转换为ONNX
spec = (tf.TensorSpec((None, 784), tf.float32, name="input"),)
onnx_model, _ = tf2onnx.convert.from_keras(model, input_signature=spec, opset=13)
with open("mnist.onnx", "wb") as f:
    f.write(onnx_model.SerializeToString())

# 3. 验证：对同一输入比较TensorFlow与ONNX Runtime的推理结果
sample = x_test[:5]  # 取5个测试样本

# TensorFlow推理
tf_output = model.predict(sample, verbose=0)

# ONNX Runtime推理
sess = ort.InferenceSession("mnist.onnx")
onnx_output = sess.run(None, {"input": sample})[0]

# 比较结果
max_diff = np.max(np.abs(tf_output - onnx_output))
print(f"最大输出差异: {max_diff:.6e}")

# 判断是否通过（差异小于1e-5视为一致）
if max_diff < 1e-5:
    print("验证通过：ONNX模型推理结果与TensorFlow一致")
else:
    print("验证失败：输出差异过大")

运行方式

将上述代码保存为convert_mnist.py，直接执行：

python convert_mnist.py

验证方法与判定标准

程序运行后，会输出TensorFlow与ONNX Runtime对同一批5个测试样本的最大输出差异值。若该值小于1e-5，则打印“验证通过”，说明转换成功且推理结果一致。同时，当前目录下会生成mnist_savedmodel/文件夹和mnist.onnx文件，前者为原始TensorFlow模型，后者为转换产物。

预期输出

最大输出差异: 0.000000e+00
验证通过：ONNX模型推理结果与TensorFlow一致

注：由于浮点运算顺序差异，实际输出差异可能在1e-7量级，仍远小于判定阈值1e-5。

常见问题

• 转换时报错“Unsupported Op”：尝试降低opset版本（如改为opset=11），或检查模型中是否有自定义层。
• 输入名称不匹配：ONNX Runtime的输入名称需与转换时指定的name="input"一致，可通过sess.get_inputs()[0].name查看。
• 动态维度问题：转换时使用tf.TensorSpec((None, 784), ...)保留了batch维度的动态性，ONNX Runtime推理时支持任意batch大小。