逑识

吾生也有涯，而知也无涯，以无涯奉有涯，其易欤？

0%

TensorFlow 2.x 基于 Keras 的模型保存及重建

发表于 分类于
本文字数： 14k 阅读时长 ≈ 12 分钟

「导语」模型训练完成后一项十分重要的步骤是对模型信息进行持久化保存，以用于后续的再训练以及线上 Serving。保存后的模型文件除了个人使用外，还可以将其分享到 TensorFlow Hub ，从而让他人可以很方便地在该预训练模型的基础上进行再次开发与训练。

Keras 模型概览

一个常见的 Keras 模型通常由以下几个部分组成：

  1. 模型的结构或配置：它指定了模型中包含的层以及层与层之间的连接方式。
  2. 一系列权重矩阵的值：用于记录模型的状态。
  3. 优化器：用于优化损失函数，可以在模型编译 (compile) 时指定。
  4. 一系列损失以及指标：用于调控训练过程，既包括在编译时指定的损失和指标，也包括通过 add_loss() 以及 add_metric() 方法添加的损失和指标。

使用 Keras API 可以将上述模型组成部分全部保存到磁盘，或者保存其中的部分。 Keras API 提供了以下三种选择：

  1. 保存模型的全部内容：通常以 TensorFlow SavedModel 格式或者 Keras H5 格式进行保存，这也是最为常用的模型保存方式。
  2. 仅保存模型的结构：通常以 json 文件的形式进行保存。
  3. 仅保存模型的权重矩阵的值：通常以 numpy 数组的形式进行保存，一般会在模型再训练或迁移学习时使用这种保存方式。

模型构建与训练

在进行模型保存之前，我们需要先构建模型并进行训练，使得模型损失和指标达到特定的状态。考虑下面这个简单的模型，它是使用 Functional API 实现的一个 3 层的深度神经网络。

from tensorflow import keras
from tensorflow.keras import layers
import numpy as np

inputs = keras.Input(shape=(784, ), name='digits')
x = layers.Dense(64, activation='relu', name='dense_1')(inputs)
x = layers.Dense(64, activation='relu', name='dense_2')(x)
outputs = layers.Dense(10, name='predictions')(x)

model = keras.Model(inputs=inputs, outputs=outputs, name='3_layer_mlp')
model.summary()

x_train, y_train = (
    np.random.random((60000, 784)),
    np.random.randint(10, size=(60000, 1)),
)
x_test, y_test = (
    np.random.random((10000, 784)),
    np.random.randint(10, size=(10000, 1)),
)

model.compile(
    loss=keras.losses.SparseCategoricalCrossentropy(from_logits=True),
    optimizer=keras.optimizers.RMSprop())
history = model.fit(x_train, y_train, batch_size=64, epochs=1)

# Save predictions for future checks
predictions = model.predict(x_test)

在使用随机数据对该模型进行一轮的迭代训练后，模型权重矩阵的值和优化器的状态都发生了变化，此时我们可以对该模型进行保存操作了，当然，你也可以训练之前进行模型保存，不过那样做并没有什么实际意义。

保存整个模型

Keras 可以将模型的全部信息保存到一个文件目录下，并且可以在之后使用该文件目录重建模型，即使没有模型的源码也可以完成重建的过程。

保存后的模型文件中，包含了以下信息：

  1. 模型的结构。
  2. 模型的权重矩阵，这是模型在训练过程中学到的信息。
  3. 模型的编译 (compile) 信息。
  4. 模型的优化器状态信息，该状态信息可以使你能够在上次训练终止的位置继续进行训练。

在 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 格式进行模型保存及重建的代码如下所示：

# Export the model to a SavedModel
model.save('path_to_saved_model')

# Recreate the exact same model
new_model = keras.models.load_model('path_to_saved_model')

# Check that the state is preserved
new_predictions = new_model.predict(x_test)
np.testing.assert_allclose(predictions, new_predictions, rtol=1e-6, atol=1e-6)

# Note that the optimizer state is preserved as well:
# you can resume training where you left off.
new_model.fit(x_train, y_train, batch_size=64, epochs=1)

执行上述代码会生成一个名为 path_to_saved_model 的目录，模型的全部信息会保存在该目录下，目录结构如下所示：

.
├── assets
├── saved_model.pb
└── variables
    ├── variables.data-00000-of-00001
    └── variables.index

其中 assets 是一个可选的目录，用于存放模型运行所需的辅助文件，比如字典文件等。 variables 目录下存放的是模型权重的检查点文件，模型的所有权重信息均保存在该目录下。 saved_model.pb 文件中包含了模型的结构以及训练的配置信息如优化器，损失以及指标等信息。

H5 格式保存

Keras 也支持使用 H5 格式进行模型保存及重建，保存后的文件中包含有模型的结构，权重矩阵的值以及编译信息等，相比于 SavedModel 格式， H5 格式是一个更为轻量的替代方案。

使用 H5 格式进行模型保存及重建的代码如下所示：

# Save the model
model.save('path_to_my_model.h5')

# Recreate the exact same model purely from the file
new_model = keras.models.load_model('path_to_my_model.h5')

# Check that the state is preserved
new_predictions = new_model.predict(x_test)
np.testing.assert_allclose(predictions, new_predictions, rtol=1e-6, atol=1e-6)

# Note that the optimizer state is preserved as well:
# you can resume training where you left off.
new_model.fit(x_train, y_train, batch_size=64, epochs=1)

但使用 H5 格式有它本身的局限性，相比于 SavedModel ， H5 格式的模型文件中不能保存以下两种信息：

  1. 外部添加的损失以及指标信息：通过 model.add_loss() 方法以及 model.add_metric() 方法添加的损失和指标信息不会保存到 H5 文件中。如果你想加载模型后继续在训练中使用它们，需要重新通过上述方法进行添加。需要注意的是在 Keras 层内通过 self.add_loss() 或 self.add_metric() 添加的损失和指标是会被 H5 文件自动保存的。
  2. 自定义对象的计算图信息：比如自定义的层 (layers) 不会被保存到 H5 文件中。此时如果使用 H5 文件直接进行模型加载，会报 ValueError: Unknown layer 错误。

因此在保存整个模型时，尤其是对于自定义 (Subclassed) 的模型，我们最好使用 SavedModel 格式来进行模型保存以及重建。

另外从代码里可以看到，由于 model.save() 方法保存了模型的全部信息，所以在重建模型后，可以在原有模型的基础上继续进行训练，而无需进行其它额外的设置。

仅保存模型结构

有时我们可能只对模型的结构感兴趣，而不想保存模型的权重值以及优化器的状态等信息。在这种情况下，我们可以借助于模型的配置方法 (config) 来对模型的结构进行保存和重建。

Sequential/Functional

对于 Sequential 模型和 Functional API 模型，因为它们大多是由预定义的层 (layers) 组成的，所以它们的配置信息都是以结构化的形式存在的，可以很方便地执行模型结构的保存操作。我们可以使用模型的 get_config() 方法来获取模型的配置，然后通过 from_config(config) 方法来重建模型。

  1. 对于最开始介绍的 Functional API 的模型而言，其结构保存与重建的代码如下所示：

    config = model.get_config()
    new_model = keras.Model.from_config(config)

    # Note that the model state is not preserved! We only saved the architecture.
    new_predictions = new_model.predict(x_test)
    assert abs(np.sum(predictions - new_predictions)) > 0.

    其中 get_config() 方法返回的结果是一个 python 字典，它包含了该模型结构的全部信息，因此可以很方便地进行模型重建。 python 字典的内容如下所示：

    {
        'name':
        '3_layer_mlp',
        'layers': [{
            'class_name': 'InputLayer',
            'config': {
                'batch_input_shape': (None, 784),
                'dtype': 'float32',
                'sparse': False,
                'ragged': False,
                'name': 'digits'
            },
            'name': 'digits',
            'inbound_nodes': []
        }, {
            'class_name': 'Dense',
            'config': {
                'name': 'dense_1',
                'trainable': True,
                'dtype': 'float32',
                'units': 64,
                'activation': 'relu',
                'use_bias': True,
                'kernel_initializer': {
                    'class_name': 'GlorotUniform',
                    'config': {
                        'seed': None
                    }
                },
                'bias_initializer': {
                    'class_name': 'Zeros',
                    'config': {}
                },
                'kernel_regularizer': None,
                'bias_regularizer': None,
                'activity_regularizer': None,
                'kernel_constraint': None,
                'bias_constraint': None
            },
            'name': 'dense_1',
            'inbound_nodes': [[['digits', 0, 0, {}]]]
        }, {
            'class_name': 'Dense',
            'config': {
                'name': 'dense_2',
                'trainable': True,
                'dtype': 'float32',
                'units': 64,
                'activation': 'relu',
                'use_bias': True,
                'kernel_initializer': {
                    'class_name': 'GlorotUniform',
                    'config': {
                        'seed': None
                    }
                },
                'bias_initializer': {
                    'class_name': 'Zeros',
                    'config': {}
                },
                'kernel_regularizer': None,
                'bias_regularizer': None,
                'activity_regularizer': None,
                'kernel_constraint': None,
                'bias_constraint': None
            },
            'name': 'dense_2',
            'inbound_nodes': [[['dense_1', 0, 0, {}]]]
        }, {
            'class_name': 'Dense',
            'config': {
                'name': 'predictions',
                'trainable': True,
                'dtype': 'float32',
                'units': 10,
                'activation': 'linear',
                'use_bias': True,
                'kernel_initializer': {
                    'class_name': 'GlorotUniform',
                    'config': {
                        'seed': None
                    }
                },
                'bias_initializer': {
                    'class_name': 'Zeros',
                    'config': {}
                },
                'kernel_regularizer': None,
                'bias_regularizer': None,
                'activity_regularizer': None,
                'kernel_constraint': None,
                'bias_constraint': None
            },
            'name': 'predictions',
            'inbound_nodes': [[['dense_2', 0, 0, {}]]]
        }],
        'input_layers': [['digits', 0, 0]],
        'output_layers': [['predictions', 0, 0]]
    }

  2. 对于 Sequential 模型而言，模型结构保存与重建的示例代码如下所示：

    from tensorflow import keras

    model = keras.Sequential([keras.Input((32, )), keras.layers.Dense(1)])
    config = model.get_config()
    new_model = keras.Sequential.from_config(config)
    new_model.summary()

为了方便模型结构的持久化存储，可以直接使用模型的 to_json() 方法将模型的信息序列化为 json 格式然后保存到本地，重建模型时可以使用 from_json() 方法解析该 json 文件内容以完成重建工作。示例代码如下所示：

json_config = model.to_json()
new_model = keras.models.model_from_json(json_config)

需要注意的是因为只保存了模型的结构，所以在模型重建后，新模型并不包含之前模型的编译信息和状态信息，因此需要重新对新模型进行编译以进行后续的训练操作。

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 方式提供自定义层的信息。示例代码如下所示：

from tensorflow import keras

class ThreeLayerMLP(keras.layers.Layer):
    def __init__(self, hidden_units):
        super().__init__()
        self.hidden_units = hidden_units
        self.dense_layers = [keras.layers.Dense(u) for u in hidden_units]

    def call(self, inputs):
        x = inputs
        for layer in self.dense_layers:
            x = layer(x)
        return x

    def get_config(self):
        return {"hidden_units": self.hidden_units}

layer = ThreeLayerMLP([64, 64, 10])
serialized_layer = keras.layers.serialize(layer)
print(serialized_layer)
new_layer = keras.layers.deserialize(
    serialized_layer,
    custom_objects={"ThreeLayerMLP": ThreeLayerMLP},
)

其中 serialize 方法返回的结果如下所示：

{'class_name': 'ThreeLayerMLP', 'config': {'hidden_units': ListWrapper([64, 64, 10])}}

对于使用自定义层和自定义函数构造的 Functional API 模型，既可以使用上述的 serialize/deserialize 方法来保存和重建模型结构，也可以通过 get_config/keras.Model.from_config 方法来完成该保存和重建操作，不过需要借助于 keras.utils.custom_object_scope 方法来指明自定义的对象。示例代码如下所示：

def custom_activation(x):
    return tf.nn.tanh(x)**2

inputs = keras.Input((64, ))
x = layer(inputs)
outputs = keras.layers.Activation(custom_activation)(x)
model = keras.Model(inputs, outputs)

config = model.get_config()
print(config)
custom_objects = {
    "ThreeLayerMLP": ThreeLayerMLP,
    "custom_activation": custom_activation
}
with keras.utils.custom_object_scope(custom_objects):
    new_model = keras.Model.from_config(config)
    new_model.summary()

另外对于上述的 Functional API 模型，还可以直接在内存中进行模型的拷贝，这与获取配置然后再通过配置进行模型重建的流程基本一致。示例代码如下所示：

with keras.utils.custom_object_scope(custom_objects):
    new_model = keras.models.clone_model(model)
    new_model.summary()

仅保存模型权重

除了可以选择只保存模型的结构信息外，你也可以只保存模型的权重信息。它在以下场景中会比较有用：

  1. 你只需要使用模型进行预测：在这种情况下，你不需要继续训练现有模型，此时就没有必要保存模型的编译信息以及优化器的状态信息。
  2. 你正在进行迁移学习：在这种情况下，你会使用现有模型的状态（权重矩阵）来训练一个新的模型，因此也不需要现有模型的编译信息。

在 Keras 中，可以通过模型的 get_weights() 和 set_weights() 方法来获取和设置权重矩阵的值。代码如下所示：

weights = model.get_weights()  # Retrieves the state of the model.
model.set_weights(weights)  # Sets the state of the model.

模型的 get_weights() 方法返回的是一个 numpy array 组成的 list ，里面记录了模型的所有权重矩阵的信息，后面可以将保存的权重信息作为 set_weights(weights) 方法的参数来设置新模型的状态，也可以将这项操作理解为模型间的权重迁移。

除了可以在同类模型间进行权重迁移外，具有兼容体系结构的模型之间同样也可以进行权重迁移操作，比如我们可以先获取前面 Functional 模型的权重矩阵信息，然后将该权重矩阵的值通过 set_weights() 方法迁移到相应的 Subclassed 模型中。不过需要注意的是，两个模型中的权重矩阵的顺序，数量以及大小要保持一致。示例代码如下所示：

from tensorflow import keras
from tensorflow.keras import layers

class ThreeLayerMLP(keras.Model):
    def __init__(self, hidden_units):
        super().__init__()
        self.hidden_units = hidden_units
        self.dense_layers = [layers.Dense(u) for u in hidden_units]

    def call(self, inputs):
        x = inputs
        for layer in self.dense_layers:
            x = layer(x)
        return x

    def get_config(self):
        return {"hidden_units": self.hidden_units}

subclassed_model = ThreeLayerMLP([64, 64, 10])
subclassed_model(tf.ones((1, 784)))
subclassed_model.set_weights(functional_model.get_weights())

assert len(functional_model.weights) == len(subclassed_model.weights)
for a, b in zip(functional_model.weights, subclassed_model.weights):
  np.testing.assert_allclose(a.numpy(), b.numpy())

另外，对于无状态的层，由于其不包含权重矩阵，所以它不会改变模型中权重矩阵的顺序以及数量，因此即使不同的模型间存在额外的无状态层（如 Dropout 层），它们也可以拥有兼容的模型结构，同样可以使用上述方式进行权重矩阵的迁移操作。

如果要将权重矩阵信息持久化到本地，可以使用 save_weights(fpath) 方法进行本地保存，后面可以使用 load_weights(fpath) 方法从本地文件中加载权重信息来恢复模型的状态。示例代码如下：

# Save weights to disk
model.save_weights('path_to_my_weights.h5')
new_model = ThreeLayerMLP([64, 64, 10])
new_model(tf.ones((1, 784)))
new_model.load_weights('path_to_my_weights.h5')

# Check that the state is preserved
new_predictions = new_model.predict(x_test)
np.testing.assert_allclose(predictions, new_predictions, rtol=1e-6, atol=1e-6)

# Note that the optimizer was not preserved.
new_model.compile(
    loss=keras.losses.SparseCategoricalCrossentropy(from_logits=True),
    optimizer=keras.optimizers.RMSprop())
new_model.fit(x_train, y_train, batch_size=64, epochs=1)

注意，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 方法来指定训练过程所需的一些配置。代码如下所示：

config = model.get_config()
weights = model.get_weights()

new_model = keras.Model.from_config(config)
new_model.set_weights(weights)

# Check that the state is preserved
new_predictions = new_model.predict(x_test)
np.testing.assert_allclose(predictions, new_predictions, rtol=1e-6, atol=1e-6)

# Note that the optimizer was not preserved,
# so the model should be compiled anew before training
# (and the optimizer will start from a blank state).
new_model.compile(
    loss=keras.losses.SparseCategoricalCrossentropy(from_logits=True),
    optimizer=keras.optimizers.RMSprop())
new_model.fit(x_train, y_train, batch_size=64, epochs=1)

TensorFlow Hub

TensorFlow Hub 是一个用于可重用机器学习的模型存储库和函数库。 TensorFlow Hub 中的模型可以在不同的学习任务中被重用 (即迁移学习)，它能让我们在训练时仅使用很小的数据集就能取得不错的效果，同时还能提升模型的泛化能力以及加快模型的学习速度。

当我们在实现一个较为常见的模型时，可以优先考虑使用 TensorFlow Hub 中已有的预训练模型，并在该模型基础上继续建模与训练，以减少程序的代码量和模型的训练时间。

TensorFlow 2.x 推荐使用 SavedModel 格式的文件在 TensorFlow Hub 上分享预训练的模型。 tensorflow_hub 函数库提供了一个类 hub.KerasLayer ，它可以从远程或本地的 SavedModel 文件中读取模型信息并以 Keras Layer 的方式重建模型，重建后的模型中包含有预训练的权重矩阵信息。

使用 hub.KerasLayer 加载模型并预测的代码如下所示：

import tensorflow_hub as hub

new_layer = hub.KerasLayer(
    'path_to_my_model',
    trainable=True,
    input_shape=[784],
    dtype=tf.float32,
)

# Check that the state is preserved
new_predictions = new_layer(x_test.astype('float32'))
np.testing.assert_allclose(predictions, new_predictions, rtol=1e-6, atol=1e-6)

在预测时，要注意保存的模型文件中权重矩阵的数据类型与输入数据的类型相匹配。在 TensorFlow 2.x 中，模型文件的数据类型由 tf.keras.backend.floatx() 获得，默认为 float32。可以通过 tf.keras.backend.set_floatx('float64') 对模型的数据类型进行修改或者将输入数据进行类型转化从而避免因类型不匹配而出错。

使用该方法加载模型与使用 keras.models.load_model 加载模型相比有几点不同：

  1. 其一，它本质上是一个 Keras Layer，因此不具有模型的一些方法，比如 summary 方法。
  2. 其二，其内部结构已经被隐藏了，因此不能查看其具体的组成结构 (layers) 。
  3. 其三，在训练时，需要将其作为模型的一层，并创建一个新的模型才能继续训练。

在 Keras 代码中使用 hub.KerasLayer 进行训练的代码如下：

new_layer = hub.KerasLayer(
    'path_to_my_model',
    trainable=True,
    input_shape=[784],
    dtype=tf.float32,
)

new_model = tf.keras.Sequential([new_layer])
new_model.compile(
    loss=keras.losses.SparseCategoricalCrossentropy(from_logits=True),
    optimizer=keras.optimizers.RMSprop())
new_model.fit(x_train, y_train, batch_size=64, epochs=1)

从代码中可以看到，我们可以像使用 Keras Layer 一样使用 hub.KerasLayer 。默认情况下， hub.KerasLayer 中的权重矩阵的值是不可训练的，即在新的训练过程中保持不变，如果需要对原来的模型做微调，可以将 trainable 参数设置为 True 。

另外，如果需要将微调后的模型重新导出，则可以使用上面介绍的各种方法如 model.save() 或 tf.keras.models.save_model() 等按需导出。

如果你想将自己的预训练模型分享到 TensorFlow Hub ，则首先需要将模型导出为 SavedModel 格式的文件，然后再推送到线上进行模型分享。在导出模型文件时需要注意，如果是要分享整个模型，在调用模型的 save 方法时需要将 include_optimizer 参数设置为 False ，也就是说分享的模型中不能包含优化器的状态信息；如果只是要分享模型的某一部分，则可以在构建模型前，将待分享的部分抽离出来，并在训练完成后保存即可。示例代码如下所示：

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

参考资料

  1. Save and load Keras models
  2. SavedModels from TF Hub in TensorFlow 2
推荐阅读
  • TensorFlow 2.x 基于 Keras 的模型构建

    「导语」TensorFlow 2.0 版发布以来，Keras 已经被深度集成于 TensorFlow 框架中，Keras API 也成为了构建深度网络模型的第一选择。使用 Keras 进行模型开发与迭代是每一个数据开发人员都需要掌握的一项基本技能，让我们一起走进 Keras 的世界一探究竟。

欢迎关注我的其它发布渠道