在 TensorFlow 2 中使用 TF Hub 中的 SavedModel

建议使用 TensorFlow 2 的 SavedModel 格式在 TensorFlow Hub 上共享预训练模型和模型部分。该格式取代了旧的 TF1 Hub 格式并提供了一组新的 API。

本页将介绍如何使用低级 hub.load() API 及其 hub.KerasLayer 封装容器在 TensorFlow 2 程序中重用 TF2 SavedModel。（通常，hub.KerasLayer 将与其他 tf.keras.layers 组合以构建 Keras 模型或 TF2 Estimator 的 model_fn。）这些 API 还可以在有限条件下加载 TF1 Hub 格式的旧模型，请参阅兼容性指南。

TensorFlow 1 的用户可以更新到 TF 1.15，并使用相同的 API。更早版本的 TF1 则无法运行。

使用 TF Hub 的 SavedModel

在 Keras 中使用 SavedModel

Keras 是 TensorFlow 的高级 API，用于通过构成 Keras 层对象来构建深度学习模型。tensorflow_hub 库提供了使用 SavedModel 的网址（或文件系统路径）初始化的 hub.KerasLayer 类，然后提供了 SavedModel 中的计算，包括其预训练权重。

以下示例使用了预训练的文本嵌入向量：

import tensorflow as tf
import tensorflow_hub as hub

hub_url = "https://tfhub.dev/google/nnlm-en-dim128/2"
embed = hub.KerasLayer(hub_url)
embeddings = embed(["A long sentence.", "single-word", "http://example.com"])
print(embeddings.shape, embeddings.dtype)

由此，可以使用通常的 Keras 方式构建文本分类器：

model = tf.keras.Sequential([
    embed,
    tf.keras.layers.Dense(16, activation="relu"),
    tf.keras.layers.Dense(1, activation="sigmoid"),
])

文本分类 Colab 完整地展示了如何训练和评估此类分类器。

hub.KerasLayer 中的模型权重默认设置为不可训练。请参阅下文中有关微调的部分来了解如何更改该设置。按照 Keras 惯例，会在同一层对象的所有应用之间共享权重。

在 Estimator 中使用 SavedModel

使用 TensorFlow 的 Estimator API 进行分布式训练的用户可以通过在 hub.KerasLayer 及其他 tf.keras.layers 方面编写 model_fn 来使用 TF Hub 中的 SavedModel。

幕后工作：SavedModel 下载和缓存

使用 TensorFlow Hub（或实现其托管协议的其他 HTTPS 服务器）中的 SavedModel 时，如果尚不存在，需要将模型下载并解压到本地文件系统中。可以设置环境变量 TFHUB_CACHE_DIR 以重写用于缓存下载并解压的 SavedModel 的默认临时位置。有关详情，请参阅缓存。

在低级 TensorFlow 中使用 SavedModel

hub.load(handle) 函数可以下载并解压 SavedModel（除非 handle 已经是文件系统路径），然后返回使用 TensorFlow 的内置函数 tf.saved_model.load() 加载该模型的结果。因此，hub.load() 可以处理任何有效的 SavedModel（与其 TF1 的前身 hub.Module 不同）。

高级主题：完成加载后对 SavedModel 的期望

根据 SavedModel 的内容，obj = hub.load(...) 的结果可通过多种方式调用（TensorFlow 的 SavedModel 指南中提供了更多详细信息）：

SavedModel 的应用签名（如有）表示为具体函数的字典，可以像 tensors_out = obj.signatures["serving_default"](**tensors_in) 一样调用，张量的字典由相应的输入和输出名称键控，并受制于签名的形状和数据类型约束。
已保存对象（如有）的 @tf.function 装饰方法恢复为 tf.function 对象，这些对象可以通过在保存之前已跟踪 tf.function 的张量和非张量参数的所有组合进行调用。特别要指出的是，如果存在带有适当跟踪的 obj.__call__ 方法，则 obj 本身可以像 Python 函数一样调用。output_tensor = obj(input_tensor, training=False) 便是一个简单的示例。

这对 SavedModel 可以实现的接口提供了极大的自由度。obj 的可重用 SavedModel 接口建立了惯例，使客户端代码（包括 hub.KerasLayer 之类的适配器）知道如何使用 SavedModel。

某些 SavedModel 可能不遵循该惯例，仅提供服务签名，特别是不打算在较大的模型中重用的整个模型。

SavedModel 中的可训练变量重新加载为可训练变量，tf.GradientTape 将在默认情况下对其进行监视。请参阅下文中有关微调的部分来了解一些注意事项，初学者应避免使用微调。如果您想使用微调，也可以查看 obj.trainable_variables 是否建议仅重新训练最初可训练变量的子集。

为 TF Hub 创建 SavedModel

概述

SavedModel 是 TensorFlow 用于已训练模型或模型部分的标准序列化格式。它存储了模型的训练权重以及用于执行其计算的确切 TensorFlow 运算。SavedModel 可独立于创建它的代码单独使用。特别值得一提的是，它可以在不同的高级建模 API（例如 Keras）之间重用，因为 TensorFlow 运算是它们的通用基本语言。

从 Keras 保存

自 TensorFlow 2 起，tf.keras.Model.save() 和 tf.keras.models.save_model() 默认为 SavedModel 格式（非 HDF5）。生成的 SavedModel 可与 hub.load()、hub.KerasLayer 和其他待提供的高级 API 的类似适配器组合使用。

要共享完整的 Keras 模型，只需在保存它时设置 include_optimizer=False。

要共享一部分 Keras 模型，请将该部分制作成模型，然后保存。您可以从头编写模型代码…

piece_to_share = tf.keras.Model(...)
full_model = tf.keras.Sequential([piece_to_share, ...])
full_model.fit(...)
piece_to_share.save(...)

…也可以在事后提取出要共享的部分（如果它与完整模型的层次一致）：

full_model = tf.keras.Model(...)
sharing_input = full_model.get_layer(...).get_output_at(0)
sharing_output = full_model.get_layer(...).get_output_at(0)
piece_to_share = tf.keras.Model(sharing_input, sharing_output)
piece_to_share.save(..., include_optimizer=False)

GitHub 上的 TensorFlow 模型对 BERT 使用前一种方法（请参见 nlp/tools/export_tfhub_lib.py，注意用于导出的 core_model 和用于恢复检查点的 pretrainer 之间的区别），对 ResNet 使用后一种方法（请参见 legacy/image_classification/tfhub_export.py）。

从低级 TensorFlow 保存

此操作需要充分熟悉 TensorFlow 的 SavedModel 指南。

如果您希望不只是提供应用签名，则应实现可重用 SavedModel 接口。从概念上讲，如下所示：

class MyMulModel(tf.train.Checkpoint):
  def __init__(self, v_init):
    super().__init__()
    self.v = tf.Variable(v_init)
    self.variables = [self.v]
    self.trainable_variables = [self.v]
    self.regularization_losses = [
        tf.function(input_signature=[])(lambda: 0.001 * self.v**2),
    ]

  @tf.function(input_signature=[tf.TensorSpec(shape=None, dtype=tf.float32)])
  def __call__(self, inputs):
    return tf.multiply(inputs, self.v)

tf.saved_model.save(MyMulModel(2.0), "/tmp/my_mul")

layer = hub.KerasLayer("/tmp/my_mul")
print(layer([10., 20.]))  # [20., 40.]
layer.trainable = True
print(layer.trainable_weights)  # [2.]
print(layer.losses)  # 0.004

微调

对导入的 SavedModel 的已训练变量及其周围模型的变量共同执行训练的过程被称为微调 SavedModel。微调可以实现更好的训练质量，但通常要求也会因此提高（可能需要更多时间，更加依赖优化器及其超参数，提高了过拟合的风险，并且需要扩充数据集，尤其是对于 CNN）。我们建议 SavedModel 使用者仅在建立良好的训练制度后，并且仅在 SavedModel 发布者建议的情况下实施微调。

微调会更改已训练的“连续”模型参数。它不会改变硬编码转换，例如词例化文本输入以及将词例映射到其嵌入向量矩阵中的对应条目。

针对 SavedModel 使用者

创建 hub.KerasLayer，如

layer = hub.KerasLayer(..., trainable=True)

将对由层加载的 SavedModel 启用微调。这会将 SavedModel 中声明的可训练权重和权重正则化器添加到 Keras 模型中，并在训练模式下运行 SavedModel 的计算（随机失活等）。

图像分类 Colab 中提供了包含可选微调的端到端训练示例。

重新导出微调结果

高级用户可能希望将微调的结果重新保存为 SavedModel 以便使用，取代最初加载的模型。可使用以下代码完成此操作：

loaded_obj = hub.load("https://tfhub.dev/...")
hub_layer = hub.KerasLayer(loaded_obj, trainable=True, ...)

model = keras.Sequential([..., hub_layer, ...])
model.compile(...)
model.fit(...)

export_module_dir = os.path.join(os.getcwd(), "finetuned_model_export")
tf.saved_model.save(loaded_obj, export_module_dir)

针对 SavedModel 创建者

创建要在 TensorFlow Hub 上共享的 SavedModel 时，请提前考虑使用者是否以及应如何微调模型，并在文档中提供指导。

从 Keras 模型保存 SavedModel 时，应使微调的所有机制全部生效（保存权重正则化损失、声明可训练变量、在 training=True 和 training=False 情况下跟踪 __call__ 等）。

选择能够适应梯度流（例如，输出 logits 而非 softmax 概率或 top-k 预测）的模型接口。

如果模型使用随机失活、批次归一化或涉及超参数的类似训练技术，请将超参数设置为对多种预期目标问题和批次大小均有意义的值。（截至撰写本文时，从 Keras 保存 SavedModel 尚不便于使用者进行调整。）

各个层上的权重正则化器（及其正则化强度系数）将得到保存，但优化器内部的权重正则化（例如 tf.keras.optimizers.Ftrl.l1_regularization_strength=...)）将丢失。请向您的 SavedModel 的使用者提出相应建议。