I want to save my neural network such that I can make a video on how the neural network trains (plotting results throughout the training process). Since the most difference in results occur at the start, I would like to save the model more often at the start of training. For example, I might want to save my model after the following epochs:

1, 2, 4, 8, 16, 32, 64, 128, ...

With tf.keras.callbacks.ModelCheckpoint I can only find methods for saving at a regular interval, for example after every epoch. However my computer does not have enough space to save the network that often. If I would take a longer interval, then I would lose all the large improvements the model makes at the start of training. Therefore I would like a different schedule (for example the one mentioned above). Preferably I would just use a list of the desired batches (or epochs) after which I want to save the network.

Is this possible with Tensorflow?

CodePudding user response：

You need to extend the Callback class. This gives you access to the on_epoch_end() function which is what you want. You can use my code below as a reference

Here's a link to the documentation for more info: Writing your own callbacks

class CustomCallback(tf.keras.callbacks.Callback):
    def __init__(self, save_epoch_list):
        super(CustomCallback, self).__init__()
        self.save_epoch_list = save_epoch_list 


    def on_epoch_end(self, epoch, logs=None):
        if epoch in self.epoch_save_list:
            # insert save model code here
            pass

epoch_save_list = [2, 4, 8, 16, 32]

model.fit(
    x_train,
    y_train,
    batch_size=64,
    epochs=30,
    callbacks=[CustomCallback(epoch_save_list)],
)