TensorFlow 2 的 SavedModel 格式 是在 TensorFlow Hub 上共享预训练模型和模型片段的推荐方法。它取代了旧的 TF1 Hub 格式 并提供了一组新的 API。
本页介绍了如何使用低级 hub.load()
API 及其 hub.KerasLayer
包装器在 TensorFlow 2 程序中重用 TF2 SavedModels。(通常,hub.KerasLayer
与其他 tf.keras.layers
结合使用,以构建 Keras 模型或 TF2 Estimator 的 model_fn
。)这些 API 也可以在一定范围内加载 TF1 Hub 格式的旧模型,请参阅 兼容性指南。
TensorFlow 1 的用户可以更新到 TF 1.15,然后使用相同的 API。旧版本的 TF1 不起作用。
使用来自 TF Hub 的 SavedModels
在 Keras 中使用 SavedModel
Keras 是 TensorFlow 的高级 API,用于通过组合 Keras Layer 对象来构建深度学习模型。 tensorflow_hub
库提供了类 hub.KerasLayer
,它使用 SavedModel 的 URL(或文件系统路径)进行初始化,然后提供来自 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
。
幕后:SavedModel 下载和缓存
使用来自 TensorFlow Hub(或其他实现其 托管 协议的 HTTPS 服务器)的 SavedModel 会将其下载并解压缩到本地文件系统(如果不存在)。环境变量 TFHUB_CACHE_DIR
可以设置为覆盖缓存下载和解压缩的 SavedModels 的默认临时位置。有关详细信息,请参阅 缓存。
在低级 TensorFlow 中使用 SavedModel
模型句柄
SavedModels 可以从指定的 handle
加载,其中 handle
是文件系统路径、有效的 TFhub.dev 模型 URL(例如,“https://tfhub.dev/...”)。Kaggle 模型 URL 根据我们的条款和与模型资产相关的许可证镜像 TFhub.dev 句柄,例如,“https://www.kaggle.com/...”。来自 Kaggle 模型的句柄等效于其相应的 TFhub.dev 句柄。
函数 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)
一样调用,其中包含以各自输入和输出名称为键的张量字典,并受签名的形状和 dtype 约束。已保存对象的
@tf.function
装饰方法(如果有)将作为 tf.function 对象还原,这些对象可以通过所有 Tensor 和非 Tensor 参数组合调用,这些参数在保存之前已针对 tf.function 进行 跟踪。特别是,如果存在具有适当跟踪的obj.__call__
方法,则可以像 Python 函数一样调用obj
本身。一个简单的示例可能如下所示:output_tensor = obj(input_tensor, training=False)
。
这使得 SavedModels 可以实现的接口具有极大的自由度。用于 obj
的 可重用 SavedModels 接口 建立了一些约定,以便客户端代码(包括 hub.KerasLayer
等适配器)知道如何使用 SavedModel。
一些 SavedModels 可能不遵循该约定,尤其是那些不打算在更大的模型中重用,而只是提供服务签名的完整模型。
SavedModel 中的可训练变量将重新加载为可训练的,并且 tf.GradientTape
默认情况下会监视它们。有关一些注意事项,请参阅下面的微调部分,并考虑在开始时避免这样做。即使您想进行微调,您可能也希望查看 obj.trainable_variables
是否建议仅重新训练原始可训练变量的子集。
为 TF Hub 创建 SavedModels
概述
SavedModel 是 TensorFlow 用于训练模型或模型片段的标准序列化格式。它将模型的训练权重与执行其计算的精确 TensorFlow 操作一起存储。它可以独立于创建它的代码使用。特别是,它可以在不同的高级模型构建 API(如 Keras)之间重复使用,因为 TensorFlow 操作是它们的通用基本语言。
从 Keras 保存
从 TensorFlow 2 开始,tf.keras.Model.save()
和 tf.keras.models.save_model()
默认使用 SavedModel 格式(而不是 HDF5)。生成的 SavedModels 可以与 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)
TensorFlow 模型 在 GitHub 上使用前一种方法来处理 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 使用者
创建类似于
layer = hub.KerasLayer(..., trainable=True)
的 hub.KerasLayer
使得可以微调由该层加载的 SavedModel。它将 SavedModel 中声明的可训练权重和权重正则化器添加到 Keras 模型中,并在训练模式下运行 SavedModel 的计算(考虑 dropout 等)。
图像分类协作 包含一个具有可选微调的端到端示例。
重新导出微调结果
高级用户可能希望将微调结果保存回可以代替原始加载的 SavedModel 的 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 模型保存应该使微调的所有机制正常工作(保存权重正则化损失,声明可训练变量,跟踪 __call__
以便同时使用 training=True
和 training=False
等)。
选择一个与梯度流配合良好的模型接口,例如输出 logits 而不是 softmax 概率或 top-k 预测。
如果模型使用 dropout、批次归一化或类似的涉及超参数的训练技术,请将它们设置为对许多预期目标问题和批次大小都有意义的值。(截至撰写本文时,从 Keras 保存并不容易让使用者调整它们。)
单个层上的权重正则化器将被保存(及其正则化强度系数),但来自优化器内部的权重正则化(如 tf.keras.optimizers.Ftrl.l1_regularization_strength=...)
)将丢失。相应地建议您的 SavedModel 的使用者。