Keras深度学习框架第十讲：自定义保存和序列化

1、绪论

自定义保存和序列化

保持和导入模型的处置，虽然在Keras中提供了默认的模型保存和序列化方法，但在某些情况下，脚本可能需要自定义这些过程以满足特定的需求。

自定义保存：除了使用model.save()将模型保存为HDF5或SavedModel格式外，还可以自定义保存过程，比如将模型结构、权重和元数据保存到自定义的二进制格式或数据库中。这可以通过实现自定义的保存函数来完成，该函数将负责将模型的各个部分转换为所需的格式，并写入到指定的位置。

自定义序列化：对于更复杂的模型结构或自定义层，默认的序列化方法可能无法满足需求。在这种情况下，可以通过实现自定义的序列化方法来处理这些复杂的对象。这通常涉及到在类的定义中添加额外的序列化逻辑，以确保对象的所有必要信息都能被正确地捕获并序列化。

自定义保存和序列化提供了更大的灵活性和控制权，使你能够根据自己的需求来定制模型的保存和加载过程。

2、自定义保持和序列化操作

2.1 基础准备

本文讨论了在Keras中可以自定义的高级保存方法。对于大多数用户来说，本文描述的序列化、保存和导出方法是足够的。然而，当需要更高级的控制或自定义模型保存和加载过程时，可以使用以下介绍的方法。

相关的API介绍

在本文的示例中，我们将使用到一下的API接口函数：

• save_assets() 和 load_assets()：

这两个函数允许你在保存和加载模型时处理额外的资源或文件，比如词汇表文件、预处理脚本等。你可以在自定义的模型类中重写这两个方法，以在保存时保存这些资源，并在加载时重新加载它们。

• save_own_variables() 和 load_own_variables()：

默认情况下，Keras会自动保存和加载模型的权重。然而，如果你需要保存或加载模型的其他变量（非权重），你可以重写这两个方法。这在你需要保存一些在训练过程中计算的特定变量或状态时特别有用。

• get_build_config() 和 build_from_config()：

这些方法允许你自定义模型的构建过程。get_build_config()方法返回一个包含构建模型所需信息的字典，而build_from_config()方法则使用此字典来重新构建模型。你可以在你的自定义模型类中重写这些方法，以添加或修改构建过程中的任何步骤。

• get_compile_config() 和 compile_from_config()：

类似地，这些方法允许你自定义模型的编译过程。get_compile_config()方法返回一个包含编译模型所需信息的字典，而compile_from_config()方法则使用此字典来重新编译模型。你可以在你的自定义模型类中重写这些方法，以添加或修改编译过程中的任何步骤。

恢复模型时的执行顺序

当你从配置中恢复一个模型时，上述方法将按照以下顺序执行：

• build_from_config()：首先，根据提供的配置信息重新构建模型的结构。
• compile_from_config()：然后，使用提供的编译配置信息重新编译模型（设置优化器、损失函数等）。
• load_own_variables()：如果重写了此方法，则在此步骤中加载模型的其他变量（非权重）。
• 加载权重：最后，加载模型的权重（这是由Keras自动处理的，不需要你重写任何方法）。

通过自定义这些方法，就可以对Keras模型的保存和加载过程进行更高级的控制和定制。

设置

import os
import numpy as np
import keras

2.2 自定义状态保存

自定义状态保存方法决定了在调用 model.save() 时如何保存你模型层的状态。用户可以重写这些方法以完全控制状态保存过程。

2.2.1 save_own_variables() 和 load_own_variables()

这两个方法分别在调用 model.save() 和 keras.models.load_model() 时保存和加载层的状态变量。默认情况下，保存和加载的状态变量是层的权重（包括可训练权重和不可训练权重）。下面是 save_own_variables() 的默认实现：

def save_own_variables(self, store):
    all_vars = self._trainable_weights + self._non_trainable_weights
    for i, v in enumerate(all_vars):
        store[f"{i}"] = v.numpy()

当使用 save_own_variables() 和 load_own_variables() 方法时，它们使用的是一个字典来存储层的变量。下面是一个如何自定义这些方法的示例，以便在保存和加载时包含额外的层变量。

@keras.utils.register_keras_serializable(package="my_custom_package")
class LayerWithCustomVariable(keras.layers.Dense):
    def __init__(self, units, **kwargs):
        super().__init__(units, **kwargs)
        self.my_variable = keras.Variable(
            np.random.random((units,)), name="my_variable", dtype="float32"
        )

    def save_own_variables(self, store):
        super().save_own_variables(store)
        # Stores the value of the variable upon saving
        store["variables"] = self.my_variable.numpy()

    def load_own_variables(self, store):
        # Assigns the value of the variable upon loading
        self.my_variable.assign(store["variables"])
        # Load the remaining weights
        for i, v in enumerate(self.weights):
            v.assign(store[f"{i}"])
        # Note: You must specify how all variables (including layer weights)
        # are loaded in `load_own_variables.`

    def call(self, inputs):
        dense_out = super().call(inputs)
        return dense_out + self.my_variable


model = keras.Sequential([LayerWithCustomVariable(1)])

ref_input = np.random.random((8, 10))
ref_output = np.random.random((8, 10))
model.compile(optimizer="adam", loss="mean_squared_error")
model.fit(ref_input, ref_output)

model.save("custom_vars_model.keras")
restored_model = keras.models.load_model("custom_vars_model.keras")

np.testing.assert_allclose(
    model.layers[0].my_variable.numpy(),
    restored_model.layers[0].my_variable.numpy(),
)

2.2.2 save_assets()和load_assets()

以上方法可以添加到模型类定义中，以便在保存和加载时存储和加载模型所需的任何额外信息。

例如，在自然语言处理（NLP）领域的层，如 TextVectorization 层和 IndexLookup 层，在保存时可能需要将其相关的词汇表（或查找表）存储在文本文件中。

以下我们通过一个简单的 assets.txt 文件来看看这个工作流程的基础。

@keras.saving.register_keras_serializable(package="my_custom_package")
class LayerWithCustomAssets(keras.layers.Dense):
    def __init__(self, vocab=None, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.vocab = vocab

    def save_assets(self, inner_path):
        # Writes the vocab (sentence) to text file at save time.
        with open(os.path.join(inner_path, "vocabulary.txt"), "w") as f:
            f.write(self.vocab)

    def load_assets(self, inner_path):
        # Reads the vocab (sentence) from text file at load time.
        with open(os.path.join(inner_path, "vocabulary.txt"), "r") as f:
            text = f.read()
        self.vocab = text.replace("<unk>", "little")


model = keras.Sequential(
    [LayerWithCustomAssets(vocab="Mary had a <unk> lamb.", units=5)]
)

x = np.random.random((10, 10))
y = model(x)

model.save("custom_assets_model.keras")
restored_model = keras.models.load_model("custom_assets_model.keras")

np.testing.assert_string_equal(
    restored_model.layers[0].vocab, "Mary had a little lamb."
)

2.3 构建和编译保存的自定义模型

2.3.1 get_build_config()和build_from_config()

在 Keras 中，get_build_config() 和 build_from_config() 方法通常用于自定义层（Custom Layers），以便在保存和加载模型时能够保存和恢复层的构建状态。这两个方法配合工作，使得自定义层能够保存其特有的状态，并在加载时恢复这些状态。

get_build_config()

这个方法应当返回一个字典，包含了足够的信息来重建该层的状态。默认情况下，它应该包含层的输入形状（input_shape），但是在编写脚本时可以覆盖这个方法以包含更多信息，比如自定义层特有的变量（Variables）或查找表（Lookup Tables）。

build_from_config(config)

这个方法接受一个配置字典（通常是由 get_build_config() 返回的），并使用这个字典来重建层的状态。在脚本中，对自定义保持和加载应该覆盖这个方法，以便它可以根据配置字典中的信息来重新创建自定义层所需的所有变量和查找表。

@keras.saving.register_keras_serializable(package="my_custom_package")
class LayerWithCustomBuild(keras.layers.Layer):
    def __init__(self, units=32, **kwargs):
        super().__init__(**kwargs)
        self.units = units

    def call(self, inputs):
        return keras.ops.matmul(inputs, self.w) + self.b

    def get_config(self):
        return dict(units=self.units, **super().get_config())

    def build(self, input_shape, layer_init):
        # Note the overriding of `build()` to add an extra argument.
        # Therefore, we will need to manually call build with `layer_init` argument
        # before the first execution of `call()`.
        super().build(input_shape)
        self._input_shape = input_shape
        self.w = self.add_weight(
            shape=(input_shape[-1], self.units),
            initializer=layer_init,
            trainable=True,
        )
        self.b = self.add_weight(
            shape=(self.units,),
            initializer=layer_init,
            trainable=True,
        )
        self.layer_init = layer_init

    def get_build_config(self):
        build_config = {
            "layer_init": self.layer_init,
            "input_shape": self._input_shape,
        }  # Stores our initializer for `build()`
        return build_config

    def build_from_config(self, config):
        # Calls `build()` with the parameters at loading time
        self.build(config["input_shape"], config["layer_init"])


custom_layer = LayerWithCustomBuild(units=16)
custom_layer.build(input_shape=(8,), layer_init="random_normal")

model = keras.Sequential(
    [
        custom_layer,
        keras.layers.Dense(1, activation="sigmoid"),
    ]
)

x = np.random.random((16, 8))
y = model(x)

model.save("custom_build_model.keras")
restored_model = keras.models.load_model("custom_build_model.keras")

np.testing.assert_equal(restored_model.layers[0].layer_init, "random_normal")
np.testing.assert_equal(restored_model.built, True)

2.3.2 get_compile_config() 和compile_from_config()

在 Keras 中，get_compile_config() 和 compile_from_config() 这样的方法并不是 Keras API 官方提供的方法。然而，从概念上讲，如果你想要实现类似的功能来保存和恢复模型编译时的配置（如优化器、损失函数等），你可以自定义这些方法或者通过其他方式来实现。

这两个方法一起工作，以保存模型编译时所使用的信息（如优化器、损失函数等），并使用这些信息来恢复和重新编译模型。

重写这两个方法对于使用自定义优化器、自定义损失函数等来编译恢复后的模型是有用的，因为这些需要在调用 compile_from_config() 中的 model.compile 之前进行反序列化。

@keras.saving.register_keras_serializable(package="my_custom_package")
def small_square_sum_loss(y_true, y_pred):
    loss = keras.ops.square(y_pred - y_true)
    loss = loss / 10.0
    loss = keras.ops.sum(loss, axis=1)
    return loss


@keras.saving.register_keras_serializable(package="my_custom_package")
def mean_pred(y_true, y_pred):
    return keras.ops.mean(y_pred)


@keras.saving.register_keras_serializable(package="my_custom_package")
class ModelWithCustomCompile(keras.Model):
    def __init__(self, **kwargs):
        super().__init__(**kwargs)
        self.dense1 = keras.layers.Dense(8, activation="relu")
        self.dense2 = keras.layers.Dense(4, activation="softmax")

    def call(self, inputs):
        x = self.dense1(inputs)
        return self.dense2(x)

    def compile(self, optimizer, loss_fn, metrics):
        super().compile(optimizer=optimizer, loss=loss_fn, metrics=metrics)
        self.model_optimizer = optimizer
        self.loss_fn = loss_fn
        self.loss_metrics = metrics

    def get_compile_config(self):
        # These parameters will be serialized at saving time.
        return {
            "model_optimizer": self.model_optimizer,
            "loss_fn": self.loss_fn,
            "metric": self.loss_metrics,
        }

    def compile_from_config(self, config):
        # Deserializes the compile parameters (important, since many are custom)
        optimizer = keras.utils.deserialize_keras_object(config["model_optimizer"])
        loss_fn = keras.utils.deserialize_keras_object(config["loss_fn"])
        metrics = keras.utils.deserialize_keras_object(config["metric"])

        # Calls compile with the deserialized parameters
        self.compile(optimizer=optimizer, loss_fn=loss_fn, metrics=metrics)


model = ModelWithCustomCompile()
model.compile(
    optimizer="SGD", loss_fn=small_square_sum_loss, metrics=["accuracy", mean_pred]
)

x = np.random.random((4, 8))
y = np.random.random((4,))

model.fit(x, y)

model.save("custom_compile_model.keras")
restored_model = keras.models.load_model("custom_compile_model.keras")

np.testing.assert_equal(model.model_optimizer, restored_model.model_optimizer)
np.testing.assert_equal(model.loss_fn, restored_model.loss_fn)
np.testing.assert_equal(model.loss_metrics, restored_model.loss_metrics)

3、总结

通过本文的学习和讨论，通过前文列出的方法可以实现各种用例，允许保存和加载具有特殊资产和状态元素的复杂模型。总结如下：

• save_own_variables 和 load_own_variables 方法决定了如何保存和加载模型的自定义状态。
• save_assets 和 load_assets 方法可以添加，用于存储和加载模型所需的任何额外信息。
• get_build_config 和 build_from_config 方法用于保存和恢复模型的构建状态。这通常涉及模型的结构和参数设置，但不包括训练配置（如优化器和损失函数）。
• get_compile_config 和 compile_from_config 方法用于保存和恢复模型的编译状态。这包括模型的训练配置，如优化器、损失函数和评估指标。

通过实现这些方法，编程人员可以对模型的保存和加载过程进行更细粒度的控制。例如，编程者可能想要保存一些在模型训练过程中计算得到的特定统计数据，或者可能想要使用自定义的序列化/反序列化逻辑来处理模型的某些部分。这些方法为这些高级用例提供了灵活性。

请注意，虽然 get_compile_config 和 compile_from_config 方法在标准的 Keras API 中并不直接存在，但它们代表了一种概念，程序员可以在自定义模型类中实现类似的功能。同样地，save_own_variables、load_own_variables、save_assets 和 load_assets 也可能需要根据具体需求进行自定义实现。