Path: blob/master/site/zh-cn/guide/keras/save_and_serialize.ipynb
25118 views
Copyright 2020 The TensorFlow Authors.
保存和加载 Keras 模型
简介
Keras 模型由多个组件组成:
架构或配置,指定模型包含的层及其连接方式。
优化器(通过编译模型来定义)。
优化器(通过编译模型来定义)。
一组损失和指标(通过编译模型或调用
add_loss()
或add_metric()
定义)。
您可以通过 Keras API 将这些片段一次性保存到磁盘,或仅选择性地保存其中一些片段:
将所有内容以 TensorFlow SavedModel 格式(或较早的 Keras H5 格式)保存到单个存档。这是标准做法。
仅保存架构/配置,通常保存为 JSON 文件。
仅保存权重值。通常在训练模型时使用。
我们来看看每个选项。什么时候使用哪个选项?它们是如何工作的?
如何保存和加载模型
如果您只有 10 秒钟来阅读本指南,则您只需了解以下内容。
保存 Keras 模型
重新加载模型:
现在,我们来查看详细信息。
安装
全模型保存和加载
您可以将整个模型保存到单个工件中。它将包括:
模型的架构/配置
模型的权重值(在训练过程中学习)
模型的编译信息(如果调用了
compile()
)优化器及其状态(如果有,这使您可以从中断的地方重新启动训练)
API
model.save()
或tf.keras.models.save_model()
tf.keras.models.load_model()
您可以使用两种格式将整个模型保存到磁盘:TensorFlow SavedModel 格式和较早的 Keras H5 格式。推荐使用 SavedModel 格式。它是使用 model.save()
时的默认格式。
您可以通过以下方式切换到 H5 格式:
将
save_format='h5'
传递给save()
。将以
.h5
或.keras
结尾的文件名传递给save()
。
SavedModel 格式
SavedModel 是更全面的保存格式,它可以保存模型架构、权重和调用函数的跟踪 Tensorflow 子计算图。这使 Keras 能够恢复内置层和自定义对象。
示例:
SavedModel 包含的内容
调用 model.save('my_model')
会创建一个名为 my_model
的文件夹,其包含以下内容:
以下示例演示了在没有重写配置方法的情况下,从 SavedModel 格式加载自定义层所发生的情况。
有关 SavedModel 格式的详细信息,请参阅 SavedModel 指南(磁盘上的 SavedModel 格式)。
SavedModel 处理自定义对象的方式
保存模型和模型的层时,SavedModel 格式会存储类名称、调用函数、损失和权重(如果已实现,还包括配置)。调用函数会定义模型/层的计算图。
在没有模型/层配置的情况下,调用函数用于创建一个与原始模型一样存在的模型,可以训练、评估该模型以及将其用于推断。
不过,在编写自定义模型或层类时,定义 get_config
和 from_config
方法始终是一个好习惯。这使您可以在需要时轻松更新计算。如需了解详情,请参阅有关自定义对象的部分。
示例:
第一个加载的模型是使用配置和 CustomModel
类加载的。第二个模型是通过动态创建类似于原始模型的模型类来加载的。
序贯模型示例:
TensoFlow 2.4 中的新功能参数 save_traces
已添加到 model.save
,它允许您切换 SavedModel 函数跟踪。保存函数以允许 Keras 在没有原始类定义的情况下重新加载自定义对象,因此当 save_traces=False
时,所有自定义对象必须已定义 get_config
/ from_config
方法。加载时,必须将自定义对象传递给 custom_objects
参数。save_traces=False
会减少 SavedModel 使用的磁盘空间并节省时间。
如上例所示,加载器动态地创建了一个与原始模型行为类似的新模型。
Keras 还支持保存单个 HDF5 文件,其中包含模型的架构、权重值和 compile()
信息。它是 SavedModel 的轻量化替代选择。
示例:
函数式模型示例:
与 SavedModel 格式相比,H5 文件不包括以下两方面内容:
通过
model.add_loss()
和model.add_metric()
添加的外部损失和指标不会被保存(这与 SavedModel 不同)。如果您的模型有此类损失和指标且您想要恢复训练,则您需要在加载模型后自行重新添加这些损失。请注意,这不适用于通过self.add_loss()
和self.add_metric()
在层内创建的损失/指标。只要该层被加载,这些损失和指标就会被保留,因为它们是该层call
方法的一部分。已保存的文件中不包含自定义对象(如自定义层)的计算图。在加载时,Keras 需要访问这些对象的 Python 类/函数以重建模型。请参阅自定义对象。
保存架构
模型的配置(或架构)指定模型包含的层,以及这些层的连接方式*。如果您有模型的配置,则可以使用权重的新初始化状态创建模型,而无需编译信息。
*请注意,这仅适用于使用函数式或序列式 API 定义的模型,不适用于子类化模型。
序贯模型或函数式 API 模型的配置
这些类型的模型是显式的层计算图:它们的配置始终以结构化形式提供。
API
get_config()
和from_config()
tf.keras.models.model_to_json()
和tf.keras.models.model_from_json()
get_config()
和 from_config()
调用 config = model.get_config()
将返回一个包含模型配置的 Python 字典。然后可以通过 Sequential.from_config(config)
(针对 Sequential
模型)或 Model.from_config(config)
(针对函数式 API 模型)重建同一模型。
相同的工作流也适用于任何可序列化的层。
层示例:
示例:
函数式模型示例:
to_json()
和 tf.keras.models.model_from_json()
这与 get_config
/ from_config
类似,不同之处在于它会将模型转换成 JSON 字符串,之后该字符串可以在没有原始模型类的情况下进行加载。它还特定于模型,不适用于层。
示例:
仅加载 TensorFlow 计算图
模型和层
子类化模型和层的架构在 __init__
和 call
方法中进行定义。它们被视为 Python 字节码,无法将其序列化为与 JSON 兼容的配置。您可以尝试将字节码序列化(例如通过 pickle
),但这样做极不安全,因为模型将无法在其他系统上进行加载。
为了保存/加载带有自定义层的模型或子类化模型,您应该重写 get_config
和 from_config
(可选)方法。此外,您还应该注册自定义对象,以便 Keras 能够感知它。
自定义函数
自定义函数(如激活损失或初始化)不需要 get_config
方法。只需将函数名称注册为自定义对象,就足以进行加载。
仅加载 TensorFlow 计算图
您可以加载由 Keras 生成的 TensorFlow 计算图。要进行此类加载,您无需提供任何 custom_objects
。您可以执行以下代码进行加载:
请注意,此方式有几个缺点:
tf.saved_model.load
返回的对象不是 Keras 模型,因此不太容易使用。例如,您将无法访问.predict()
或.fit()
。tf.saved_model.load
返回的对象不是 Keras 模型,因此不太容易使用。例如,您将无法访问.predict()
或.fit()
。
虽然不鼓励使用此方式,但当您遇到棘手问题(例如,您丢失了自定义对象的代码,或在使用 tf.keras.models.load_model()
加载模型时遇到问题)时,它还是能够提供帮助。
有关详细信息,请参阅 tf.saved_model.load
相关页面。
定义配置方法
规范:
get_config
应该返回一个 JSON 可序列化字典,以便兼容 Keras 节省架构和模型的 API。from_config(config)
(classmethod
) 应返回从配置创建的新层或模型对象。默认实现返回cls(**config)
。
示例:
注册自定义对象
Keras 会记录哪个类生成了配置。在上面的示例中,tf.keras.layers.serialize
会生成自定义层的序列化形式:
Keras 会维护一份所有内置层、模型、优化器和指标类的主列表,用于查找正确的类以调用 from_config
。如果找不到该类,则会引发错误 (Value Error: Unknown layer
)。可以通过几种方式将自定义类注册到此列表中:
在加载函数中设置
custom_objects
参数。(请参阅上文”定义配置方法“部分中的示例)tf.keras.utils.custom_object_scope
或者tf.keras.utils.CustomObjectScope
tf.keras.utils.register_keras_serializable
示例:
内存中模型克隆
您还可以通过 tf.keras.models.clone_model()
在内存中克隆模型。这相当于获取模型的配置,然后通过配置重建模型(因此它不会保留编译信息或层的权重值)。
示例:
您只需使用模型进行推断:在这种情况下,您无需重新开始训练,因此不需要编译信息或优化器状态。
在内存中将权重从一层转移到另一层
您只需使用模型进行推断:在这种情况下,您无需重新开始训练,因此不需要编译信息或优化器状态。
您正在进行迁移学习:在这种情况下,您需要重用先验模型的状态来训练新模型,因此不需要先验模型的编译信息。
用于内存中权重迁移的 API
您可以使用 get_weights
和 set_weights
在不同对象之间复制权重:
tf.keras.layers.Layer.get_weights()
:返回 Numpy 数组列表。tf.keras.layers.Layer.set_weights()
:将模型权重设置为weights
参数中的值。
示例如下。
通常建议使用相同的 API 来构建模型。如果您在序贯模型和函数式模型之间,或在函数式模型和子类化模型等之间进行切换,请始终重新构建预训练模型并将预训练权重加载到该模型。
在内存中将权重从一个模型转移到另一个具有兼容架构的模型
无状态层的情况
无状态层不会改变权重的顺序或数量,因此即便存在额外的/缺失的无状态层,模型也可以具有兼容架构。
用于将权重保存到磁盘并将其加载回来的 API
可以用以下格式调用 model.save_weights
,将权重保存到磁盘:
TensorFlow 检查点
HDF5
model.save_weights
的默认格式是 TensorFlow 检查点。可以通过以下两种方式指定保存格式:
save_format
参数:将值设置为save_format="tf"
或save_format="h5"
。path
参数:如果路径以.h5
或.hdf5
结束,则使用 HDF5 格式。除非设置了save_format
,否则对于其他后缀,将使用 TensorFlow 检查点格式。
您还可以选择将权重作为内存中 Numpy 数组取回。每个 API 都有自己的优缺点,详情如下。
TF 检查点格式
示例:
格式详细信息
TensorFlow 检查点格式使用对象特性名称来保存和恢复权重。以 tf.keras.layers.Dense
层为例。该层包含两个权重:dense.kernel
和 dense.bias
。将层保存为 tf
格式后,生成的检查点会包含 "kernel"
和 "bias"
键及其对应的权重值。有关详情,请参阅 TF 检查点指南中的“加载机制”。
请注意,特性/计算图边缘根据父对象中使用的名称而非变量的名称进行命名。请考虑下面示例中的 CustomLayer
。变量 CustomLayer.var
是将 "var"
而非 "var_a"
作为键的一部分来保存的。
迁移学习示例
本质上,只要两个模型具有相同的架构,它们就可以共享同一个检查点。
示例:
通常建议使用相同的 API 来构建模型。如果您在序贯模型和函数式模型之间切换,或在函数式模型和子类化模型等之间切换,请始终重新构建预训练模型并将预训练权重加载到该模型。
下一个问题是,如果模型架构截然不同,如何保存权重并将其加载到不同模型?解决方案是使用 tf.train.Checkpoint
来保存和恢复确切的层/变量。
示例:
HDF5 格式
HDF5 格式包含按层名称分组的权重。权重是通过将可训练权重列表与不可训练权重列表连接起来进行排序的列表(与 layer.weights
相同)。因此,如果模型的层和可训练状态与保存在检查点中的相同,则可以使用 HDF5 检查点。
示例:
请注意,当模型包含嵌套层时,更改 layer.trainable
可能导致 layer.weights
的顺序不同。
迁移学习示例
从 HDF5 加载预训练权重时,建议将权重加载到设置了检查点的原始模型中,然后将所需的权重/层提取到新模型中。
示例: