Keras 使用自定义文件名在训练回调中保存 protobuffer 和权重

Question

我正在尝试为我的模型拟合过程编写一个回调，如果在一个 Epoch 结束后，模型得到改进，我会将权重和保存为原型缓冲区。 最好是./tmp/weights.hdf5 和./tmp/model.pb 之类的东西。我为此使用了两个回调（下面的 mnist 示例，我使用的是 TF 2.6.0）：

import tensorflow as tf
import numpy as np

mnist = tf.keras.datasets.mnist
(x_train, y_train),(x_test, y_test) = mnist.load_data()
x_train, x_test = x_train / 255.0, x_test / 255.0

def create_model():
  return tf.keras.models.Sequential([
    tf.keras.layers.Flatten(input_shape=(28, 28)),
    tf.keras.layers.Dense(512, activation='relu'),
    tf.keras.layers.Dropout(0.2),
    tf.keras.layers.Dense(10, activation='softmax')
  ])

model = create_model()
model.compile(optimizer='adam',
              loss='sparse_categorical_crossentropy',
              metrics=['accuracy'],
              steps_per_execution=10)

callback_weights = tf.keras.callbacks.ModelCheckpoint(
    filepath="./tmp/weights.hdf5", save_weights_only=True, save_best_only=True)

callback_model = tf.keras.callbacks.ModelCheckpoint(
    filepath="./tmp/", save_weights_only=False, save_best_only=True)

model.fit(x=x_train,
          y=y_train,
          epochs=2,
          validation_data=(x_test, y_test),
          callbacks=[callback_weights, callback_model])

而ls tmp 给出：

assets keras_metadata.pb saved_model.pb variables weights.hdf5

我的问题：

• 我必须为此使用两个回调还是可以在一个回调中完成？
• 如何控制 protobuffer 文件的名称？在 save_weights_only=True 有效时指定文件名，但当我使用 save_weights_only=False 执行此操作时，它会从 filepath 参数创建一个目录。

Answer 1

要回答您上述具体问题，

(1)。我们不必为此使用两个回调（保存模型或权重）。它可以通过一个回调来完成。 (2)。当我们设置save_weights_only=False 时，意味着程序将保存整个模型及其当前状态（或权重）。为了更清楚的了解，请看下文

# it'll save only weight 
callback_weights = tf.keras.callbacks.ModelCheckpoint(
    filepath="weights.h5", 
    save_weights_only=True,  
    save_best_only=True)

# it'll save model config + weight = entire trained model 
callback_weights_model = tf.keras.callbacks.ModelCheckpoint(
    filepath="model.h5", 
    save_weights_only=False, # entire model (config + weight)
    save_best_only=True)

# it'll also save model config + weight = entire trained model 
callback_model = tf.keras.callbacks.ModelCheckpoint(
    filepath="./tmp/", 
    save_weights_only=False, # entire model (config + weight)
    save_best_only=True)

在这里，callback_weights_model 和 callback_model 将使用不同的输出格式完成相同的工作。但callback_weights 只会保存当前状态或训练过的权重文件。所以，如果我们在训练时使用它们如下：

model.fit(...
          callbacks=[callback_weights, 
                     callback_model, 
                     callback_weights_model])

那么我们将有以下文件。

tmp/ [asset, variable, .pb]
model.h5
weight.h5

我们来看看

loaded_model = tf.keras.models.load_model('./tmp/')
assert np.allclose(model.predict(x_test), loaded_model.predict(x_test))

loaded_model = tf.keras.models.load_model('/content/model.h5')
assert np.allclose(model.predict(x_test), loaded_model.predict(x_test))

#  ValueError: No model config found in the file at /content/weights.h5.
# loaded_model = tf.keras.models.load_model('/content/weights.h5') 

loaded_model = create_model()
loaded_model.load_weights('/content/weights.h5')
assert np.allclose(model.predict(x_test), loaded_model.predict(x_test))

所以，如果我们需要保存整个模型，我们可以选择上面的 callback_weights_model 或 callback_model 回调。否则，如果我们只需要保存权重文件，我们可以使用callback_weights。