保存和加载高级模型格式

准备将训练好的模型用于开发环境之外，是一项主要任务。无论是部署到强大服务器集群，还是资源有限的移动设备，都需要一种标准化、自包含的方式来打包您的模型。TensorFlow正是为此提供了SavedModel格式。它不仅仅是保存模型权重 (weight)；它包含了完整的TensorFlow程序，包括计算图、变量值以及任何所需的资产或元数据，使模型可移植且可提供服务。

SavedModel格式：一种通用软件包

可以把SavedModel看作您整个TensorFlow模型的序列化表示，已准备好部署。它是生产环境的推荐格式，并用作将模型转换为TensorFlow Lite或TensorFlow.js的中间格式。

当您以这种格式保存模型时，TensorFlow会创建一个包含以下内容的目录：

1. saved_model.pb 或 saved_model.pbtxt: 这个文件存储了由tf.functions定义的计算图结构以及像签名这样的元数据。它使用Protocol Buffer格式（二进制.pb或文本.pbtxt）。
2. variables/ 目录: 这个子目录包含了模型变量的训练好的值（权重 (weight)、偏差等），通常存储在检查点文件中。
3. assets/ 目录: 一个可选子目录，保存了图所需的外部文件，例如用于文本处理的词汇文件或类别映射。
4. assets.extra/ 目录: 一个可选子目录，库或用户可以在其中添加自己的资产，与模型位于同一位置但不由TensorFlow图直接读取。
5. fingerprint.pb: 一个包含标识SavedModel的唯一指纹的文件。

一个TensorFlow SavedModel的典型目录结构。

使用Keras保存模型

创建SavedModel最简单的方法是使用tf.keras.Model实例的save方法。Keras会处理标准模型的序列化细节。

import tensorflow as tf

# 假设 'model' 是一个已编译的 tf.keras.Model 实例
# model = tf.keras.Sequential([...])
# model.compile(...)
# model.fit(...)

# 以SavedModel格式保存模型
model.save("my_keras_model")

# 您也可以明确指定签名（强烈建议用于服务）
@tf.function(input_signature=[tf.TensorSpec(shape=[None, 784], dtype=tf.float32)])
def serving_default(input_tensor):
    return {'output': model(input_tensor)}

model.save("my_keras_model_with_signature", signatures={'serving_default': serving_default})

调用model.save("directory_path")会创建一个包含SavedModel的目录。默认情况下，Keras会尝试保存模型的前向传播（call方法）以及可能还有其训练/评估配置。然而，明确定义并保存签名对部署来说很重要，我们很快就会看到。

使用tf.saved_model.save保存

为了获得更细致的控制，特别是在处理自定义tf.Module对象或需要定义要暴露的特定函数时，您可以使用更底层的tf.saved_model.save API。这也是model.save在Keras模型内部所使用的。

import tensorflow as tf

# 使用自定义 tf.Module 的示例
class MyCustomModule(tf.Module):
    def __init__(self, name=None):
        super().__init__(name=name)
        self.my_variable = tf.Variable(5.0, name="my_var")
        self.untracked_list = [] # 不会被保存

    @tf.function(input_signature=[tf.TensorSpec(shape=[], dtype=tf.float32)])
    def __call__(self, x):
        # 此函数将被追踪并保存为签名
        print("Tracing MyCustomModule.__call__")
        return x * self.my_variable

    @tf.function(input_signature=[tf.TensorSpec(shape=[], dtype=tf.float32)])
    def add_amount(self, amount):
        # 另一个要暴露的函数
        print("Tracing MyCustomModule.add_amount")
        return self.my_variable + amount

module_to_save = MyCustomModule()

# 保存模块，将其方法暴露为签名
tf.saved_model.save(module_to_save, "my_module_savedmodel",
                    signatures= {
                        'serving_default': module_to_save.__call__,
                        'add': module_to_save.add_amount
                    })

print(f"SavedModel 创建于：my_module_savedmodel")
print(f"可用签名: {tf.saved_model.load('my_module_savedmodel').signatures.keys()}")

在这里，我们明确传递了tf.Module实例以及一个定义我们希望在保存模型中可用的签名的字典。只有那些是tf.Variable、tf.Module或可追踪的数据结构（例如包含可追踪对象的列表/字典）的属性才会被保存。Python 基本类型或不可追踪的集合（例如上面self.untracked_list）将丢失。

了解签名

签名是您的SavedModel用于推断的入口点。它们定义了模型中特定函数的预期输入张量（形状和数据类型）和对应的输出张量。当您使用TensorFlow Serving或其他平台部署模型时，这些签名定义了您可以调用的API端点。

• 它们为何重要？ 签名提供了一个与语言无关的契约，用于与模型交互。发送请求的客户端不需要了解内部Python代码；它只需要遵循签名定义的输入/输出规范。它们还允许保存过程创建针对特定输入类型的优化、具体的TensorFlow图，可能会提高性能。
• 它们是如何创建的？ 您可以通过使用@tf.function装饰Python方法并提供input_signature来定义签名。input_signature是tf.TensorSpec对象的列表或元组，指定每个输入参数 (parameter)的形状和数据类型。
• serving_default: 按照惯例，许多服务系统会查找名为serving_default的签名。为模型最常见的推断任务定义此签名是一个好的做法。

在没有明确签名的情况下保存Keras模型时，Keras通常会根据模型的call方法创建默认签名，但依赖这种隐式行为有时可能导致意想不到的结果，特别是对于复杂的输入处理。在调用model.save或使用tf.saved_model.save之前，使用@tf.function和input_signature明确定义签名是部署的推荐方法。

加载SavedModels

保存后，您可以使用tf.saved_model.load将SavedModel重新加载到TensorFlow程序中。此函数恢复了tf.Module或tf.keras.Model对象、它的变量、资产，并且非常重要地，恢复了作为签名保存的由tf.function装饰的方法。

import tensorflow as tf

# 加载之前保存的Keras模型
loaded_keras_model = tf.saved_model.load("my_keras_model_with_signature")

# 访问签名字典
print(f"可用签名: {list(loaded_keras_model.signatures.keys())}")

# 获取 'serving_default' 签名的特定函数对象
inference_func = loaded_keras_model.signatures['serving_default']

# 准备一些与 input_signature 匹配的虚拟输入数据
dummy_input = tf.random.uniform([2, 784], dtype=tf.float32) # 2个批次，784个特征

# 调用函数
output_dict = inference_func(dummy_input)
print(f"输出张量形状: {output_dict['output'].shape}")

# 加载之前保存的自定义模块
loaded_module = tf.saved_model.load("my_module_savedmodel")

# 访问其签名
add_func = loaded_module.signatures['add']
default_func = loaded_module.signatures['serving_default']

# 调用函数
result1 = default_func(tf.constant(10.0))
print(f"默认函数输出: {result1['output'].numpy()}") # 基于保存的变量的输出

result2 = add_func(amount=tf.constant(3.0))
print(f"添加函数输出: {result2['output'].numpy()}")

加载的对象并非完全是原始Python对象。它是一个专门的内部用户对象（_UserObject），保存了恢复的状态和函数。您主要通过其signatures属性与其交互，该属性是一个将签名键（如'serving_default'）映射到相应具体TensorFlow函数的字典。调用这些函数会执行恢复的计算图。请注意，原始@tf.function中的Python代码（如print语句）仅在函数首次调用或保存时的初始追踪阶段运行，而不是在调用已加载签名时。

检查SavedModels

部署之前，通常有益于检查SavedModel的内容，特别是其可用签名及其输入/输出规范。TensorFlow为此提供了saved_model_cli命令行工具。

# 显示SavedModel的所有信息
saved_model_cli show --dir my_keras_model_with_signature --all

# 示例输出片段：
# 标签集为 'serve' 的 MetaGraphDef 包含以下 SignatureDefs：
#
# signature_def['__saved_model_init_op']:
#   被跳过的 SessionInit 操作。
#
# signature_def['serving_default']:
#   给定的 SavedModel SignatureDef 包含以下输入：
#     inputs['input_tensor'] tensor_info:
#         dtype: DT_FLOAT
#         shape: (-1, 784)
#         name: serving_default_input_tensor:0
#   给定的 SavedModel SignatureDef 包含以下输出：
#     outputs['output'] tensor_info:
#         dtype: DT_FLOAT
#         shape: (-1, 10) # 假设是一个10分类输出
#         name: StatefulPartitionedCall:0
#   方法名称是：tensorflow/serving/predict

# 仅显示特定标签集（通常是用于 TF Serving 的 'serve'）的签名
saved_model_cli show --dir my_module_savedmodel --tag_set serve

# 显示特定签名细节
saved_model_cli show --dir my_module_savedmodel --tag_set serve --signature_def add

saved_model_cli show命令显示了结构、可用的标签集（图的组，通常只有serve用于推断），以及指定标签集中每个签名的详细输入/输出张量信息。这使您能够验证模型保存正确，并了解如何构建推断请求。

自定义组件的考量

如果您的模型使用了自定义层、激活函数 (activation function)、损失函数 (loss function)，或通过子类化Keras或TensorFlow基类定义的其他自定义Python对象，那么保存和加载需要特别注意。

• Keras模型： 当使用model.save时，Keras通常会处理通过tf.keras.utils.register_keras_serializable注册或在加载期间通过custom_objects参数 (parameter)传递（tf.keras.models.load_model(..., custom_objects=...)）的自定义对象。然而，SavedModel格式旨在实现可移植性，这意味着加载环境（如TF Serving）可能无法访问您的自定义Python代码。
• 最佳实践： 只要可能，请在您的@tf.functions或自定义层/模型方法中直接使用标准TensorFlow操作实现自定义逻辑。这确保逻辑被捕获在SavedModel内部的TensorFlow图本身中。如果您必须依赖外部Python代码，您将需要确保该代码在加载SavedModel的环境中可用并正确注册。对于复杂的依赖，容器化（例如，使用Docker）对部署来说变得不可或缺。

使用SavedModel格式正确保存模型，特别是带有明确定义的签名，是可靠部署的起点。它打包您的图、权重 (weight)和所需资产，创建一个自包含的工件，准备好由TensorFlow Serving提供服务，使用TensorFlow Lite转换为边缘设备使用，或在其他生产流程中使用。