汤不热吧SavedModel 格式详解：为什么它是 TensorFlow 生产环境下模型持久化的唯一真神
在 TensorFlow 生态系统中，模型持久化有两种常见方式：Keras H5 格式（.h5）和 SavedModel 格式。虽然 H5 格式简单易用，但它本质上只保存了模型的权重和网络结构配置。而在生产环境中，SavedModel 格式则是公认的黄金标准。为什么呢？
SavedModel 不仅仅保存了权重，它还保存了完整的 TensorFlow 运行图（Computation Graph）、所有变量（Variables）、资产文件（Assets）以及最重要的：具体的签名函数（Signatures）。这意味着 SavedModel 是一个完全自包含的、独立于原始训练代码的部署包，可以直接被 TensorFlow Serving、TensorFlow Lite Converter 或其他推理系统无缝加载和使用。
SavedModel 的三大核心优势
- 自包含的执行图 (Graph Def)： SavedModel 保存了模型的整个计算图作为 tf.function 的具体化（Concrete Functions），脱离了 Python 环境也能执行，这对 C++ 或 Go 等生产环境部署至关重要。
- 兼容性与标准化： 它是 TensorFlow Serving 和 TFLite/TFJS 等所有下游工具链的官方输入格式。
- 签名支持 (Signatures)： 允许定义明确的输入和输出接口，确保推理时接口稳定。
实践操作：创建、保存与加载 SavedModel
我们将演示如何创建一个简单的 Keras 模型，并将其保存为 SavedModel，然后使用 TensorFlow 的低级别 API 加载并进行推理。
步骤一：创建并训练模型
我们创建一个简单的线性回归模型，用于演示。
import tensorflow as tf
import numpy as np
import os
# 确保导出目录存在且为空
EXPORT_PATH = './my_production_model/1'
if tf.io.gfile.exists(EXPORT_PATH):
tf.io.gfile.rmtree(EXPORT_PATH)
# 1. 创建 Keras 模型
model = tf.keras.Sequential([
tf.keras.layers.Dense(64, activation='relu', input_shape=(5,)),
tf.keras.layers.Dense(1)
])
model.compile(optimizer='adam', loss='mse')
# 2. 模拟训练数据
X_train = np.random.rand(100, 5).astype(np.float32)
Y_train = X_train.sum(axis=1) * 3 + 10
print("开始训练模型...")
model.fit(X_train, Y_train, epochs=5, verbose=0)
# 验证模型在 SavedModel 导出前的预测能力
input_data = np.array([[1.0, 2.0, 3.0, 4.0, 5.0]], dtype=np.float32)
original_prediction = model.predict(input_data)
print(f"原始模型预测结果: {original_prediction[0][0]:.4f}")
步骤二：导出为 SavedModel 格式
使用 tf.saved_model.save 或 Keras 的 save 方法（推荐 Keras 的方法，因为它会自动生成 serving_default 签名）。
# 3. 导出模型到 SavedModel 格式
# 注意：导出的路径通常包含版本号，便于部署管理
tf.saved_model.save(model, EXPORT_PATH)
print(f"\n模型已成功导出到: {EXPORT_PATH}")
# 检查导出的文件结构：
# my_production_model/1/
# ├── assets/ (辅助文件，如词汇表)
# ├── variables/ (所有权重和优化器状态)
# └── saved_model.pb (核心文件：计算图和签名定义)
步骤三：脱离 Keras 环境加载并推理
现在，我们不再依赖原始的 model Python 对象或 Keras 定义，仅使用 tf.saved_model.load 加载模型，这模拟了生产环境（如 TF Serving）的加载过程。
# 4. 使用低级 API 加载 SavedModel
# 注意：我们不需要知道原始模型是用 Keras 还是 tf.Module 构建的
loaded_model = tf.saved_model.load(EXPORT_PATH)
# 5. 调用 'serving_default' 签名进行推理
# SavedModel 暴露的接口是 Concrete Function
infer_function = loaded_model.signatures["serving_default"]
# 将 NumPy 数据转换为 Tensor
tensor_input = tf.constant(input_data)
# 执行推理
# SavedModel 返回一个字典，包含输出张量
output_dict = infer_function(tensor_input)
# 获取输出结果（通常键为 'output_0' 或 'output'）
key = list(output_dict.keys())[0]
loaded_prediction = output_dict[key].numpy()[0][0]
print(f"\n加载的模型 (通过签名调用) 预测结果: {loaded_prediction:.4f}")
# 验证结果一致性
assert abs(original_prediction[0][0] - loaded_prediction) < 1e-4
print("SavedModel 验证通过：预测结果与原始模型一致。")
总结
SavedModel 格式通过将模型的权重、计算图以及清晰的输入/输出签名封装在一个自包含的目录中，解决了生产环境中部署依赖性、兼容性和效率的问题。对于任何需要投入生产或进行推理加速（如转 TFLite、ONNX 转换）的 TensorFlow 模型，SavedModel 都是唯一推荐和标准的持久化格式。