跳到主要内容

TensorFlow SavedModel 格式

TensorFlow SavedModel 是 TensorFlow 中用于保存和加载模型的标准化格式。它不仅是模型的序列化形式，还包含了模型的计算图、权重、优化器状态以及元数据。SavedModel 格式的设计目标是跨平台、跨语言兼容，并且支持多种部署场景。

什么是 SavedModel 格式？​

SavedModel 是 TensorFlow 2.x 中推荐的模型保存格式。它可以将整个模型（包括架构、权重和优化器状态）保存为一个目录，便于后续的加载、部署和优化。与早期的 checkpoint 或 HDF5 格式相比，SavedModel 提供了更强大的功能和更好的兼容性。

备注

SavedModel 格式不仅适用于 TensorFlow，还可以与其他框架（如 TensorFlow.js、TensorFlow Lite 和 TensorFlow Serving）无缝集成。

如何保存模型为 SavedModel 格式？​

在 TensorFlow 中，保存模型为 SavedModel 格式非常简单。以下是一个示例：

python
import tensorflow as tf

# 创建一个简单的模型
model = tf.keras.Sequential([
    tf.keras.layers.Dense(10, activation='relu', input_shape=(10,)),
    tf.keras.layers.Dense(1, activation='sigmoid')
])

# 编译模型
model.compile(optimizer='adam', loss='binary_crossentropy')

# 保存模型为 SavedModel 格式
model.save('my_model')

执行上述代码后，TensorFlow 会在当前目录下创建一个名为 my_model 的文件夹，其中包含以下内容：

my_model/
├── assets/
├── variables/
│   ├── variables.data-00000-of-00001
│   └── variables.index
└── saved_model.pb
  • saved_model.pb：包含模型的计算图和元数据。
  • variables/：包含模型的权重。
  • assets/：包含模型所需的额外文件（如词汇表）。

如何加载 SavedModel 格式的模型？​

加载 SavedModel 格式的模型同样非常简单。以下是一个示例：

python
import tensorflow as tf

# 加载 SavedModel 格式的模型
loaded_model = tf.keras.models.load_model('my_model')

# 使用加载的模型进行预测
input_data = tf.random.normal([1, 10])
predictions = loaded_model(input_data)
print(predictions)
提示

加载 SavedModel 格式的模型时，TensorFlow 会自动恢复模型的计算图和权重，因此无需重新定义模型架构。

SavedModel 格式的实际应用场景​

SavedModel 格式在以下场景中非常有用：

  1. 模型部署：SavedModel 格式可以直接用于 TensorFlow Serving，支持高性能的模型推理服务。
  2. 跨平台兼容：SavedModel 格式可以转换为 TensorFlow.js 或 TensorFlow Lite 格式，适用于 Web 和移动端部署。
  3. 模型共享：SavedModel 格式是标准化的，便于团队成员之间共享和协作。

以下是一个使用 TensorFlow Serving 部署 SavedModel 的示例：

bash
# 启动 TensorFlow Serving 服务
tensorflow_model_server --rest_api_port=8501 --model_name=my_model --model_base_path=/path/to/my_model

然后，您可以通过 REST API 调用模型进行推理：

bash
curl -d '{"instances": [[1, 2, 3, 4, 5, 6, 7, 8, 9, 10]]}' -X POST http://localhost:8501/v1/models/my_model:predict

总结​

TensorFlow SavedModel 格式是保存、加载和部署模型的强大工具。它提供了跨平台兼容性、高性能推理支持以及简化的模型共享方式。通过掌握 SavedModel 格式，您可以更高效地管理和部署 TensorFlow 模型。

警告

在保存模型时，请确保模型目录路径唯一，以避免覆盖已有的模型文件。

附加资源与练习​

  • 官方文档：TensorFlow SavedModel 指南
  • 练习：尝试保存一个自定义的 Keras 模型为 SavedModel 格式，并使用 TensorFlow Serving 部署它。
  • 进阶阅读：了解如何将 SavedModel 转换为 TensorFlow.js 或 TensorFlow Lite 格式。

通过实践这些内容，您将更深入地理解 TensorFlow SavedModel 格式的强大功能和应用场景。