作者： Neel Kovelamudi
 在 TensorFlow.org 上檢視
|
 在 Google Colab 中執行
|
 在 GitHub 上檢視原始碼
|
 在 keras.io 上檢視
|
簡介
本指南涵蓋 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 tensorflow as tf
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 LayerWithCustomVariables(keras.layers.Dense):
def __init__(self, units, **kwargs):
super().__init__(units, **kwargs)
self.stored_variables = tf.Variable(
np.random.random((10,)), name="special_arr", dtype=tf.float32
)
def save_own_variables(self, store):
super().save_own_variables(store)
# Stores the value of the `tf.Variable` upon saving
store["variables"] = self.stored_variables.numpy()
def load_own_variables(self, store):
# Assigns the value of the `tf.Variable` upon loading
self.stored_variables.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):
return super().call(inputs) * self.stored_variables
model = keras.Sequential([LayerWithCustomVariables(1)])
ref_input = np.random.random((8, 10))
ref_output = np.random.random((8,))
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].stored_variables.numpy(),
restored_model.layers[0].stored_variables.numpy(),
)
1/1 [==============================] - 4s 4s/step - loss: 0.2723
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."
)
build 和 compile 儲存自訂
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 tf.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 customization in overriding `build()` adds 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.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 = super().get_build_config() # only gives `input_shape`
build_config.update(
{"layer_init": self.layer_init} # 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 = tf.math.squared_difference(y_pred, y_true)
loss = loss / 10.0
loss = tf.reduce_sum(loss, axis=1)
return loss
@keras.saving.register_keras_serializable(package="my_custom_package")
def mean_pred(y_true, y_pred):
return tf.reduce_mean(y_pred)
@keras.saving.register_keras_serializable(package="my_custom_package")
class ModelWithCustomCompile(keras.Model):
def __init__(self):
super().__init__()
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 [==============================] - 1s 651ms/step - loss: 0.0616 - accuracy: 0.0000e+00 - mean_pred: 0.2500 WARNING:absl:Skipping variable loading for optimizer 'SGD', because it has 1 variables whereas the saved optimizer has 5 variables.
結論
使用本教學課程中學到的方法，可以實現各種使用案例，允許儲存和載入具有異國情調的資產和狀態元素的複雜模型。總而言之
save_own_variables和load_own_variables決定如何儲存和載入您的狀態。save_assets和load_assets可以新增來儲存和載入模型需要的任何額外資訊。get_build_config和build_from_config儲存和還原模型的建構狀態。get_compile_config和compile_from_config儲存和還原模型的編譯狀態。