TensorFlow 篇 | TensorFlow 2.x 基于 Keras 的模型保存及重建
「导语」模型训练完成后一项十分重要的步骤是对模型信息进行持久化保存,以用于后续的再训练以及线上 Serving。保存后的模型文件除了个人使用外,还可以将其分享到 TensorFlow Hub ,从而让他人可以很方便地在该预训练模型的基础上进行再次开发与训练。
Keras 模型概览
一个常见的 Keras
模型通常由以下几个部分组成:
模型的结构或配置:它指定了模型中包含的层以及层与层之间的连接方式。
一系列权重矩阵的值:用于记录模型的状态。
优化器:用于优化损失函数,可以在模型编译 (
compile
) 时指定。
一系列损失以及指标:用于调控训练过程,既包括在编译时指定的损失和指标,也包括通过
add_loss()
以及add_metric()
方法添加的损失和指标。
使用 Keras API
可以将上述模型组成部分全部保存到磁盘,或者保存其中的部分。 Keras API
提供了以下三种选择:
保存模型的全部内容:通常以
TensorFlow SavedModel
格式或者Keras H5
格式进行保存,这也是最为常用的模型保存方式。
仅保存模型的结构:通常以
json
文件的形式进行保存。
仅保存模型的权重矩阵的值:通常以
numpy
数组的形式进行保存,一般会在模型再训练或迁移学习时使用这种保存方式。
模型构建与训练
在进行模型保存之前,我们需要先构建模型并进行训练,使得模型损失和指标达到特定的状态。考虑下面这个简单的模型,它是使用 Functional API
实现的一个 3
层的深度神经网络。
在使用随机数据对该模型进行一轮的迭代训练后,模型权重矩阵的值和优化器的状态都发生了变化,此时我们可以对该模型进行保存操作了,当然,你也可以训练之前进行模型保存,不过那样做并没有什么实际意义。
保存整个模型
Keras
可以将模型的全部信息保存到一个文件目录下,并且可以在之后使用该文件目录重建模型,即使没有模型的源码也可以完成重建的过程。
保存后的模型文件中,包含了以下信息:
模型的结构。
模型的权重矩阵,这是模型在训练过程中学到的信息。
模型的编译 (
compile
) 信息。
模型的优化器状态信息,该状态信息可以使你能够在上次训练终止的位置继续进行训练。
在 Keras
中,可以使用 model.save()
方法或者 tf.keras.models.save_model()
方法对模型进行保存,使用 tf.keras.models.load_model()
方法来对保存的模型文件进行加载并重建模型。
模型文件保存的格式可以有 2
种,一种为 TensorFlow SavedModel
格式,另一种是 Keras H5
格式,官方推荐使用 SavedModel
格式进行模型保存,它也是 model.save()
方法默认使用的保存格式。可以通过设置 save()
方法中的参数 format='h5'
或者指定 save()
方法中的文件名后缀为 .h5
或 .keras
来切换为使用 H5
格式进行模型保存。
SavedModel 格式保存
使用 SavedModel
格式进行模型保存及重建的代码如下所示:
执行上述代码会生成一个名为 path_to_saved_model
的目录,模型的全部信息会保存在该目录下,目录结构如下所示:
其中 assets
是一个可选的目录,用于存放模型运行所需的辅助文件,比如字典文件等。 variables
目录下存放的是模型权重的检查点文件,模型的所有权重信息均保存在该目录下。 saved_model.pb
文件中包含了模型的结构以及训练的配置信息如优化器,损失以及指标等信息。
H5 格式保存
Keras
也支持使用 H5
格式进行模型保存及重建,保存后的文件中包含有模型的结构,权重矩阵的值以及编译信息等,相比于 SavedModel
格式, H5
格式是一个更为轻量的替代方案。
使用 H5
格式进行模型保存及重建的代码如下所示:
但使用 H5
格式有它本身的局限性,相比于 SavedModel
, H5
格式的模型文件中不能保存以下两种信息:
外部添加的损失以及指标信息:通过
model.add_loss()
方法以及model.add_metric()
方法添加的损失和指标信息不会保存到H5
文件中。如果你想加载模型后继续在训练中使用它们,需要重新通过上述方法进行添加。需要注意的是在Keras
层内通过self.add_loss()
或self.add_metric()
添加的损失和指标是会被H5
文件自动保存的。
自定义对象的计算图信息:比如自定义的层 (
layers
) 不会被保存到H5
文件中。此时如果使用H5
文件直接进行模型加载,会报ValueError: Unknown layer
错误。
因此在保存整个模型时,尤其是对于自定义 (Subclassed
) 的模型,我们最好使用 SavedModel
格式来进行模型保存以及重建。
另外从代码里可以看到,由于 model.save()
方法保存了模型的全部信息,所以在重建模型后,可以在原有模型的基础上继续进行训练,而无需进行其它额外的设置。
仅保存模型结构
有时我们可能只对模型的结构感兴趣,而不想保存模型的权重值以及优化器的状态等信息。在这种情况下,我们可以借助于模型的配置方法
(config
) 来对模型的结构进行保存和重建。
Sequential/Functional
对于 Sequential
模型和 Functional API
模型,因为它们大多是由预定义的层 (layers
) 组成的,所以它们的配置信息都是以结构化的形式存在的,可以很方便地执行模型结构的保存操作。我们可以使用模型的 get_config()
方法来获取模型的配置,然后通过 from_config(config)
方法来重建模型。
对于最开始介绍的
Functional API
的模型而言,其结构保存与重建的代码如下所示:
其中 get_config()
方法返回的结果是一个 python
字典,它包含了该模型结构的全部信息,因此可以很方便地进行模型重建。 python
字典的内容如下所示:
对于
Sequential
模型而言,模型结构保存与重建的示例代码如下所示:
为了方便模型结构的持久化存储,可以直接使用模型的 to_json()
方法将模型的信息序列化为 json
格式然后保存到本地,重建模型时可以使用 from_json()
方法解析该 json
文件内容以完成重建工作。示例代码如下所示:
需要注意的是因为只保存了模型的结构,所以在模型重建后,新模型并不包含之前模型的编译信息和状态信息,因此需要重新对新模型进行编译以进行后续的训练操作。
Subclassed 模型和层
Subclassed
模型的结构信息是在 __init__
和 call
方法中定义的,它们会被视为 python
字节码而无法将其序列化为 json
格式。截至本文完成前,还不存在一种方法可以直接
对自定义模型(继承自 keras.Model
)的结构信息
进行保存和重建操作。
而对于 Subclassed
层(继承自 keras.Layer
)来说是可以进行结构保存和重建的,我们需要重写其 get_config()
方法和 from_config()
(可选的)方法来达成这一目标。
其中 get_config()
方法需要返回一个可被 json
序列化的字典,以便适配 Keras
的结构以及模型保存的接口。 from_config(config)
类方法需要根据 config
参数返回一个新的层对象。
定义好层的 config
方法之后,我们可以使用 serialize
和 deserialize
方法来序列化(保存)和反序列化(重建)层的结构信息,注意,在重建层时需要以 custom_objects
方式提供自定义层的信息。示例代码如下所示:
其中 serialize
方法返回的结果如下所示:
对于使用自定义层和自定义函数构造的 Functional API
模型,既可以使用上述的 serialize/deserialize
方法来保存和重建模型结构,也可以通过 get_config/keras.Model.from_config
方法来完成该保存和重建操作,不过需要借助于 keras.utils.custom_object_scope
方法来指明自定义的对象。示例代码如下所示:
另外对于上述的 Functional API
模型,还可以直接在内存中进行模型的拷贝,这与获取配置然后再通过配置进行模型重建的流程基本一致。示例代码如下所示:
仅保存模型权重
除了可以选择只保存模型的结构信息外,你也可以只保存模型的权重信息。它在以下场景中会比较有用:
你只需要使用模型进行预测:在这种情况下,你不需要继续训练现有模型,此时就没有必要保存模型的编译信息以及优化器的状态信息。
你正在进行迁移学习:在这种情况下,你会使用现有模型的状态(权重矩阵)来训练一个新的模型,因此也不需要现有模型的编译信息。
在 Keras
中,可以通过模型的 get_weights()
和 set_weights()
方法来获取和设置权重矩阵的值。代码如下所示:
模型的 get_weights()
方法返回的是一个 numpy array
组成的 list
,里面记录了模型的所有权重矩阵的信息,后面可以将保存的权重信息作为 set_weights(weights)
方法的参数来设置新模型的状态,也可以将这项操作理解为模型间的权重迁移。
除了可以在同类模型间进行权重迁移外,具有兼容体系结构的模型之间同样也可以进行权重迁移操作,比如我们可以先获取前面 Functional
模型的权重矩阵信息,然后将该权重矩阵的值通过 set_weights()
方法迁移到相应的 Subclassed
模型中。不过需要注意的是,两个模型中的权重矩阵的顺序,数量以及大小要保持一致。示例代码如下所示:
另外,对于无状态的层,由于其不包含权重矩阵,所以它不会改变模型中权重矩阵的顺序以及数量,因此即使不同的模型间存在额外的无状态层(如 Dropout
层),它们也可以拥有兼容的模型结构,同样可以使用上述方式进行权重矩阵的迁移操作。
如果要将权重矩阵信息持久化到本地,可以使用 save_weights(fpath)
方法进行本地保存,后面可以使用 load_weights(fpath)
方法从本地文件中加载权重信息来恢复模型的状态。示例代码如下:
注意,save_weights
方法既可以保存为 Keras HDF5
格式的文件,也可以保存为 TensorFlow Checkpoint
格式的文件。与 model.save
方法类似,它也包含一个参数 save_format
,同样有两种取值。默认为 tf
表示使用 Checkpoint
格式保存, h5
表示使用 HDF5
格式保存。如果不明确指定,它会根据文件的后缀名去推测,以 .h5
或 .hdf5
结尾的文件名会以 HDF5
格式进行保存。
可以结合使用 get_config()/from_config()
和 get_weights()/set_weights()
来重建模型并恢复原有的状态。但是与使用 model.save()
不同,这种方式因为不能保存模型的训练配置以及优化器的状态,所以在继续训练时,需要重新调用模型的 compile
方法来指定训练过程所需的一些配置。代码如下所示:
TensorFlow Hub
TensorFlow Hub
是一个用于可重用机器学习的模型存储库和函数库。 TensorFlow Hub
中的模型可以在不同的学习任务中被重用 (即迁移学习
),它能让我们在训练时仅使用很小的数据集就能取得不错的效果,同时还能提升模型的泛化能力以及加快模型的学习速度。
当我们在实现一个较为常见的模型时,可以优先考虑使用 TensorFlow Hub
中已有的预训练模型,并在该模型基础上继续建模与训练,以减少程序的代码量和模型的训练时间。
TensorFlow 2.x
推荐使用 SavedModel
格式的文件在 TensorFlow Hub
上分享预训练的模型。 tensorflow_hub
函数库提供了一个类 hub.KerasLayer
,它可以从远程或本地的 SavedModel
文件中读取模型信息并以 Keras Layer
的方式重建模型,重建后的模型中包含有预训练的权重矩阵信息。
使用 hub.KerasLayer
加载模型并预测的代码如下所示:
在预测时,要注意保存的模型文件中权重矩阵的数据类型与输入数据的类型相匹配。在 TensorFlow 2.x
中,模型文件的数据类型由 tf.keras.backend.floatx()
获得,默认为 float32
。可以通过 tf.keras.backend.set_floatx('float64')
对模型的数据类型进行修改或者将输入数据进行类型转化从而避免因类型不匹配而出错。
使用该方法加载模型与使用 keras.models.load_model
加载模型相比有几点不同:
其一,它本质上是一个
Keras Layer
,因此不具有模型的一些方法,比如summary
方法。
其二,其内部结构已经被隐藏了,因此不能查看其具体的组成结构 (
layers
) 。
其三,在训练时,需要将其作为模型的一层,并创建一个新的模型才能继续训练。
在 Keras
代码中使用 hub.KerasLayer
进行训练的代码如下:
从代码中可以看到,我们可以像使用 Keras Layer
一样使用 hub.KerasLayer
。默认情况下, hub.KerasLayer
中的权重矩阵的值是不可训练的,即在新的训练过程中保持不变,如果需要对原来的模型做微调,可以将 trainable
参数设置为 True
。
另外,如果需要将微调后的模型重新导出,则可以使用上面介绍的各种方法如 model.save()
或 tf.keras.models.save_model()
等按需导出。
如果你想将自己的预训练模型分享到 TensorFlow Hub
,则首先需要将模型导出为 SavedModel
格式的文件,然后再推送到线上进行模型分享。在导出模型文件时需要注意,如果是要分享整个模型,在调用模型的 save
方法时需要将 include_optimizer
参数设置为 False
,也就是说分享的模型中不能包含优化器的状态信息;如果只是要分享模型的某一部分,则可以在构建模型前,将待分享的部分抽离出来,并在训练完成后保存即可。示例代码如下所示:
参考资料
版权声明: 本文为 InfoQ 作者【Alex】的原创文章。
原文链接:【http://xie.infoq.cn/article/25976d50894fcf9dee6364e16】。文章转载请联系作者。
评论