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

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

下面的代码可以保存你的模型，然后再加载......

 import pickle
 filename = 'finalized_model.sav'
 pickle.dump(model, open(filename, 'wb'))    
 loaded_model = pickle.load(open(filename, 'rb'))

下面的部分是完整的代码...

import tensorflow as to
import pickle

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')
    filename = 'finalized_model.sav'
    pickle.dump(best_model, open(filename, 'wb'))

    loaded_model = pickle.load(open(filename, 'rb'))
    return loaded_model

Answer 2

下面是一个创建回调并在回调时将模型保存到外部的基本例子。list. 它必须是一个列表（或者是一个允许使用方法修改的类型）。基本的 tf.keras.callbacks.Callback 在回调类中，该类扩展了一个额外的参数，即 list。__init___ 方法。这个例子说明了它的工作原理。当回调被调用在 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)