📚 课程目录第4章 SavedModel格式详解

4、SavedModel格式详解：理解SavedModel的目录结构（assets、variables、saved_model.pb），以及如何保存和加载SavedModel

在TFLite Micro的转换工具链中，SavedModel是TensorFlow 2.x推荐的标准模型序列化格式。它不仅仅是一个单一的文件，而是一个自包含的目录结构，旨在将模型的图定义、权重参数以及辅助资源完整地打包在一起。对于嵌入式开发者而言，理解SavedModel的内部构造是确保模型能够被正确加载、转换并最终部署到微控制器上的关键前提。本节将深入剖析SavedModel的三大核心组件，并演示其保存与加载的标准流程。

上图展示了SavedModel目录结构与TFLite Micro转换链路的关系。在转换过程中，saved_model.pb提供了计算图的拓扑结构，variables/目录提供了具体的数值参数，而assets/目录则补充了图执行时可能需要的辅助文件。三者缺一不可，共同构成了一个完整、可移植的模型包。

4.1 核心原理：SavedModel的三层架构

SavedModel的设计哲学是“自包含”与“版本无关”。它通过将模型的不同关注点分离到不同的子目录中，实现了模型在训练、评估、部署等不同生命周期阶段的平滑过渡。其核心架构可以拆解为以下三个层次：

图定义层（saved_model.pb）：这是一个序列化的GraphDef协议缓冲区文件。它包含了模型的完整计算图结构，即从输入张量到输出张量的所有操作（Op）及其连接关系。此外，它还存储了模型的签名（SignatureDef），签名定义了模型的输入/输出张量名称、数据类型和形状，这是后续转换器定位输入输出节点的关键依据。
状态管理层（variables/）：该目录存储了模型的所有可训练变量（如卷积核权重、偏置项）以及非可训练变量（如BatchNorm的移动均值/方差）。在TensorFlow 2.x中，变量被拆分为variables.index（包含变量名到分片位置的映射）和variables.data-*（实际存储的张量值，可能被分片为多个文件）。这种设计支持了分布式训练和超大模型的分片存储。
资源辅助层（assets/）：此目录用于存放计算图执行时可能需要的外部文件。例如，NLP模型中的词表文件、初始化查找表、或自定义Op所需的字典数据。这些文件不会被序列化到GraphDef中，而是以原始文件形式附带在模型包内。

关键理解：SavedModel与传统的H5格式（单文件）不同，它不是一个压缩包，而是一个明文目录。这意味着你可以直接使用文件系统工具（如ls、du）查看模型的大小和组成，也便于在部署流水线中单独替换或更新权重文件，而无需重新导出整个模型。

4.2 操作流程：保存与加载SavedModel的标准方法

在TensorFlow 2.x中，保存和加载SavedModel的API非常简洁。以下流程适用于从Keras模型或自定义tf.Module导出。

4.2.1 保存SavedModel

对于继承自tf.keras.Model的模型，推荐使用tf.saved_model.save()方法。该方法不仅保存了图结构，还保存了模型的调用约定（即__call__方法对应的签名）。

import tensorflow as tf

# 假设 model 是一个已经训练好的 Keras 模型
model = tf.keras.Sequential([
    tf.keras.layers.Input(shape=(32,)),
    tf.keras.layers.Dense(64, activation='relu'),
    tf.keras.layers.Dense(10, activation='softmax')
])

# 保存为 SavedModel 格式
tf.saved_model.save(model, "path/to/saved_model_dir")

# 保存后目录结构如下：
# path/to/saved_model_dir/
# ├── assets/
# ├── saved_model.pb
# └── variables/
#     ├── variables.data-00000-of-00001
#     └── variables.index

如果需要保存推理专用的签名（例如，指定输入特征名为"input_tensor"），可以在保存时通过signatures参数指定：

# 使用 @tf.function 定义推理函数
@tf.function(input_signature=[tf.TensorSpec(shape=[None, 32], dtype=tf.float32, name="input_tensor")])
def serving_fn(x):
    return {"output": model(x)}

tf.saved_model.save(model, "path/to/saved_model_dir", signatures={"serving_default": serving_fn})

4.2.2 加载SavedModel

加载SavedModel使用tf.saved_model.load()方法。加载后返回一个SavedModel对象，可以通过其属性访问恢复的模型。

# 加载模型
loaded_model = tf.saved_model.load("path/to/saved_model_dir")

# 查看模型签名
infer = loaded_model.signatures["serving_default"]

# 执行推理
import numpy as np
sample_input = np.random.rand(1, 32).astype(np.float32)
output = infer(tf.constant(sample_input))
print(output["output"].numpy())

经验提示：在加载后，务必检查loaded_model.signatures字典中的键。默认情况下，Keras模型会有一个名为serving_default的签名。如果签名缺失或名称不符，TFLite转换器将无法自动定位输入输出张量，导致转换失败。

4.3 案例拆解：从Keras到SavedModel的实战转换

为了更直观地理解SavedModel在TFLite Micro工具链中的作用，我们以一个简单的语音唤醒词识别模型（DNN）为例，演示完整的保存、加载与验证流程。

背景与约束：目标设备是一块具有256KB Flash和64KB RAM的MCU。模型输入为40维MFCC特征，输出为3个类别（"silence"、"unknown"、"yes"）。由于资源受限，最终需要转换为TFLite Micro格式，因此中间产物必须是结构清晰的SavedModel。

初始现象：在直接使用model.save("model.h5")保存H5格式后，尝试用tf.lite.TFLiteConverter.from_keras_model()转换时，虽然成功，但无法在TFLite Micro中指定自定义的输入输出名称，且无法附加词表等资源文件。

分析判断：H5格式仅包含模型权重和结构，不包含服务签名和资源文件。对于嵌入式部署，我们需要精确控制输入张量的名称（例如input_mfcc）和输出张量的名称（例如output_scores），以便在C++代码中绑定缓冲区。因此，必须改用SavedModel格式。

实际操作步骤：

定义带签名的导出函数：使用@tf.function包装模型调用，并指定输入张量名为input_mfcc。
保存为SavedModel：调用tf.saved_model.save()，并传入signatures参数。
验证加载：使用tf.saved_model.load()加载，并打印签名信息，确认输入输出名称符合预期。
转换TFLite：使用tf.lite.TFLiteConverter.from_saved_model()进行转换，并设置输入输出张量名称。

# 步骤1：定义签名
@tf.function(input_signature=[tf.TensorSpec(shape=[1, 40], dtype=tf.float32, name="input_mfcc")])
def model_serve(x):
    return {"output_scores": model(x)}

# 步骤2：保存
tf.saved_model.save(model, "kws_saved_model", signatures={"serving_default": model_serve})

# 步骤3：验证
loaded = tf.saved_model.load("kws_saved_model")
print(loaded.signatures["serving_default"].structured_input_signature)
# 输出: ({}, {'x': TensorSpec(shape=(1, 40), dtype=tf.float32, name='input_mfcc')})

# 步骤4：转换
converter = tf.lite.TFLiteConverter.from_saved_model("kws_saved_model")
converter.input_arrays = ["input_mfcc"]
converter.output_arrays = ["output_scores"]
tflite_model = converter.convert()

结果对比：转换后的TFLite模型在MCU上成功运行，输入输出缓冲区与C++代码完美对齐。而使用H5格式时，输入张量名默认为serving_default_x:0，导致需要额外的映射代码，增加了出错风险。

风险与代价：SavedModel目录比单个H5文件占用更多存储空间（因为包含多个文件），但在模型转换阶段，该目录通常位于开发机而非MCU上，因此不构成部署负担。主要代价是需要额外编写签名定义代码。

常见误区：不要试图手动修改saved_model.pb文件来改变模型结构。该文件是二进制协议缓冲区，任何手动编辑都会导致文件损坏。如需修改模型结构，应回到Keras代码层面修改后重新导出。

4.4 常见误区与总结复盘

在工程实践中，开发者经常在SavedModel的使用上遇到以下问题，值得特别关注：

误区/问题	现象	根因分析	解决方案
忽略签名定义	转换TFLite时提示找不到输入/输出	未使用`signatures`参数，导致默认签名不明确	使用`@tf.function`定义明确的输入输出名称
混淆H5与SavedModel	使用`from_keras_model`转换时无法附加资源文件	H5格式不支持`assets/`目录	改用`from_saved_model`转换器
加载后无法调用	`loaded_model()`报错	直接调用对象而非通过签名调用	使用`loaded_model.signatures['serving_default'](x)`
变量分片过多	转换速度慢，且`variables.data-*`文件数量多	训练时使用了分布式策略导致变量分片	在保存前使用`tf.saved_model.save`的`variables_to_save`参数合并分片