目录
导言
涵盖的API
Setup
状态保存自定义
构建和编译保存自定义
结论
政安晨的个人主页：政安晨
欢迎 👍点赞✍评论⭐收藏
收录专栏: TensorFlow与Keras机器学习实战
希望政安晨的博客能够对您有所裨益，如有不足之处，欢迎在评论区提出指正！
相较于上一篇文章而言，这是一篇更高级的图层和模型自定义保存指南。
导言
本文涵盖可在 Keras 保存中自定义的高级方法。对于大多数用户来说，上文的初级序列化、保存和导出文章中概述的方法已经足够。
涵盖的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()
Setup
import os
import numpy as np
import keras状态保存自定义
这些方法决定了调用 model.save() 时如何保存模型图层的状态。您可以重载这些方法来完全控制状态保存过程。
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()这些方法使用的存储空间是一个字典，可以填充层变量。
下面我们来看一个自定义的示例：
示例：
@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.2908save_assets() 和 load_assets()
这些方法可以添加到模型类定义中，以存储和加载模型所需的任何附加信息。
例如，文本矢量化层（TextVectorization layer）和索引查找层（IndexLookup layer）等 NLP 领域层可能需要在保存时将其相关词汇（或查找表）存储到文本文件中。
让我们用一个简单的 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."
)构建和编译保存自定义
get_build_config() 和 build_from_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)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_variables 和 load_own_variables 决定了保存和加载状态的方式。
save_assets 和 load_assets 可用于存储和加载模型所需的任何附加信息。
get_build_config 和 build_from_config 用于保存和恢复模型的构建状态。
get_compile_config 和 compile_from_config 保存和恢复模型的编译状态。
