Tensorflow Callback：如何将最佳模型保存在内存而不是磁盘上

Question

我使用 Tensorflow 使用以下函数进行回归

import tensorflow as tf

def ff(*args, **kwargs):
    model = tf.keras.models.Sequential()
    model.add(tf.keras.Input(shape=[inp_train.shape[-1],]))
    for i in range(n_layer):
        model.add(tf.keras.layers.Dense(n_unit, activation=act))
    model.add(tf.keras.layers.Dense(out_train.shape[1]))
    model.compile(optimizer=opt, loss='mae')
    early_stop  = tf.keras.callbacks.EarlyStopping(monitor='val_loss', patience=100)
    check_point = tf.keras.callbacks.ModelCheckpoint('best_model.h5', monitor='val_loss', save_best_only=True)
    model.fit(inp_train, out_train, epochs=n_epoch, batch_size=s_batch, validation_data=(inp_val, out_val), callbacks=[early_stop, check_point], verbose=0)
    best_model = tf.keras.models.load_model('best_model.h5')
    return model, best_mode

如您所见，我通过check_point 回调保存了最佳模型，并在以后使用它进行预测。问题是这样我必须先将最好的模型保存在磁盘上，然后再从磁盘加载。如果我想并行运行几次，因为每次运行都会创建一个同名文件，所以它不起作用。

那么，如何在变量中分配最佳模型而无需将其保存在磁盘上？

Answer 1

注意：我修复了一个错误并且未经测试

我必须为自己做这件事，我想我会分享：

回调：

class SaveBestModel(tf.keras.callbacks.Callback):
    def __init__(self, save_best_metric='val_loss', this_max=False):
        self.save_best_metric = save_best_metric
        self.max = this_max
        if this_max:
            self.best = float('-inf')
        else:
            self.best = float('inf')

    def on_epoch_end(self, epoch, logs=None):
        metric_value = logs[self.save_best_metric]
        if self.max:
            if metric_value > self.best:
                self.best = metric_value
                self.best_weights = self.model.get_weights()

        else:
            if metric_value < self.best:
                self.best = metric_value
                self.best_weights= self.model.get_weights()

用法：

save_best_model = SaveBestModel()
model.fit(data, callbacks=[save_best_model]
#set best weigts
model.set_weights(save_best_model.best_weights)

Answer 2

这里是创建回调并在回调时将模型保存到外部list 的基本示例。它必须是一个列表（或允许使用方法进行修改的类型）。在回调类__init___ 方法中，基类tf.keras.callbacks.Callback 使用附加参数列表进行了扩展。这个例子表明它有效。当在 training_end 上调用回调时，它会将当前模型附加到列表中。

import tensorflow as tf
from tensorflow.python.keras.models import Model

# define a custom callback
class MyCustomCallback(tf.keras.callbacks.Callback):

  def __init__(self, external_list):
      self.list_obj = external_list

  def on_train_end(self, logs=None):
      self.list_obj.append(self.model)

# test the idea works
model_save_list = []
my_callback = MyCustomCallback(model_save_list)

model1 = Model()
my_callback.set_model(model1)
my_callback.on_train_end()

print(model_save_list)

运行这个，你会看到内部模型被添加到你的列表对象中：

[<tensorflow.python.keras.engine.training.Model object at 0x10d230b50>]

通过将新的回调添加到回调中来修改您的训练，如下所示：

model.fit(inp_train, out_train, epochs=n_epoch, batch_size=s_batch, validation_data=(inp_val, out_val), callbacks=[early_stop, my_callback], verbose=0)