开发者指南 / 自定义保存和序列化

自定义保存和序列化

作者:Neel Kovelamudi
创建日期 2023/03/15
最后修改日期 2023/03/15
描述:关于自定义层和模型保存的更高级指南。

在 Colab 中查看 GitHub 源代码


简介

本指南涵盖了可以在 Keras 保存中自定义的高级方法。对于大多数用户而言,主要序列化、保存和导出指南中概述的方法就足够了。

API

我们将介绍以下 API

  • save_assets()load_assets()
  • save_own_variables()load_own_variables()
  • get_build_config()build_from_config()
  • get_compile_config()compile_from_config()

在恢复模型时,这些方法按以下顺序执行

  • build_from_config()
  • compile_from_config()
  • load_own_variables()
  • load_assets()

设置

import os
import numpy as np
import keras

状态保存自定义

这些方法决定了在调用 model.save() 时如何保存模型层的状态。您可以覆盖它们以完全控制状态保存过程。

save_own_variables()load_own_variables()

当分别调用 model.save()keras.models.load_model() 时,这些方法保存和加载层的 state 变量。默认情况下,保存和加载的 state 变量是层的权重(可训练和不可训练)。以下是 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()

这些方法使用的存储是一个字典,可以用层变量填充。让我们看一个自定义它的示例。

示例

@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(),
)
 1/1 ━━━━━━━━━━━━━━━━━━━━ 0s 101ms/step - loss: 0.2908

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."
)

buildcompile 保存自定义

get_build_config()build_from_config()

这些方法协同工作以保存层的构建状态并在加载时恢复它们。

默认情况下,这仅包括一个包含层输入形状的 build 配置字典,但覆盖这些方法可以用于包含其他变量和查找表,这些变量和查找表对于恢复已构建的模型很有用。

示例

@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)

get_compile_config()compile_from_config()

这些方法协同工作以保存模型编译时使用信息(优化器、损失函数等),并使用此信息恢复和重新编译模型。

覆盖这些方法对于使用自定义优化器、自定义损失函数等编译恢复的模型很有用,因为在 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)
 1/1 ━━━━━━━━━━━━━━━━━━━━ 0s 79ms/step - accuracy: 0.0000e+00 - loss: 0.0627 - mean_metric_wrapper: 0.2500

结论

使用本教程中学习的方法可以满足各种用例,允许保存和加载具有奇特资产和状态元素的复杂模型。总结如下

  • save_own_variablesload_own_variables 决定如何保存和加载状态。
  • save_assetsload_assets 可以用来存储和加载模型所需的任何其他信息。
  • get_build_configbuild_from_config 保存和恢复模型的构建状态。
  • get_compile_configcompile_from_config 保存和恢复模型的编译状态。