Keras Callbacks and How to Save Your Model from Overtraining | by Josh Johnson | Towards Data Science

In this article, you will learn how to use the ModelCheckpoint callback in Keras to save the best version of your model during training.
I love building predictive deep learning models. I love watching the training outputs, seeing the loss fall and watching for the diverging losses between training and validation sets that indicate overfitting. But sometimes a model finds a great solution...and keeps training to a solution that only works for the training set. Now, if I’m there, staring like it’s a fish tank, I can interrupt the training before too much damage is done. But, who wants to sit there and stare at models training all day? I mean, I do, but my time is better spent doing other work during long training sessions.
Then I learned about Keras callbacks, and specifically ModelCheckpoint!
Top athletes always ride the edge of over-training. The right about of training will optimize an athlete’s performance, but training too much leads to injury and reduced performance. Deep-learning models are similar. The right amount of training makes a strong model, but too much and performance can drop off on new data.
During training deep learning models seek to minimize their loss, to be more accurate according to a given loss function. However, they judge that accuracy on the set of data they are training on. Imagine a school-kid who takes a practice test home and memorizes every problem on it and every answer. If they haven’t found the underlying patterns, the strategies that the test is meant to assess, they will be at a loss because none of the questions on the practice test are on the exam!
Deep learning models can do that same thing. If they train too much on a dataset, they can learn that dataset specifically, rather than picking up the underlying functions connecting features and labels. The model can overfit the data.
There are many strategies to combat overfitting in predictive modeling, and those should be deployed as necessary, but overfitting from overtraining is a unique hazard.
Luckily, Keras has some help for us in the form of callbacks! Before I knew about callbacks I thought I had to guess the right number of training epochs, or use trial and error to tune them. Luckily, we don’t! We can use the Keras callback keras.callbacks.ModelCheckpoint() to save the model at its best performing epoch.
A callback is a function that can run after each epoch. It takes as arguments the epoch number and any metrics you have your model keeping track of. They can be used to do such useful things as scheduling reductions in the learning rate (I love a well-tuned decaying learning rate, don’t you?), early stopping of training, or saving the model between epochs. You can even code your own callback functions for your own special purposes.
If you want to know more about all of the available callbacks, check out the Keras documentation here.
One of my favorite callbacks is ModelCheckpoint. I like to multi-task and often open a few Google Colab windows, set some deep models to train in them, and then work on other tasks on my local machine (I don’t have the GPU hardware to do efficient Tensorflow model training like Google does). I can set my models to train for a little longer than I think they might need and come back later to load them from the point before they began overfitting and losing accuracy on the validation set.
The ModelCheckpoint callback can be loaded from keras.callbacks
from keras.callbacks import ModelCheckpoint
We initialize the class object with the filepath to which to save, the conditions under which we want it saved, and how transparent the process should be. For example, let’s say we want only the very best version of the model and we define ‘best’ as the one with the lowest validation loss. You can choose as a trigger either your loss function or any of the metrics you passed to the metrics argument when you compiled the model.
We would initialize our callback like so:
checkpoint = ModelCheckpoint(filepath=filepath,                              monitor=’val_loss’,                             verbose=1,                              save_best_only=True,                             mode=’min’)
One cool trick for saving multiple times is to append the epoch number and/or current metric you are monitoring to the filepath, since the epoch will be passed to the ModelCheckpoint object.
filepath = 'my_best_model.epoch{epoch:02d}-loss{val_loss:.2f}.hdf5'
The above will save your model to disk as something like:
my_best_model.epoch_01-loss_13.31.hdf5
Keras models take a list of callbacks as an argument in the .fit() call. The argument expects a list, even if you are passing only one callback. However, you can deck your model out with all kinds of fancy callbacks. My favorite combo is ModelCheckpoint and ReduceLROnPlateau. The second one monitors a metric and reduces the learning rate by a given amount when that metric plateaus.
We pass our list of callbacks to the .fit() method like so:
callbacks = [checkpoint, <other_callbacks_you_want>]model.fit(X_train,          y_train,          epochs=epochs          callbacks=callbacks)
A note about saving models: models saved in .hdf5 format are great because the whole model is one place and can be loaded somewhere else, such as in deployment. However the files can get large, and saving your model at every epoch can get storage intensive fast. One option available in the ModelCheckpoint callback constructor is save_weights_only=True. This will save space, but will not save the entire model architecture. In order to recover it, you would rebuild the model and then assign the saved weights, rather than just loading it all in one step.
Another quirk I’ve found is that not every layer likes to be saved in .hdf5 format. In a previous blog post I showed you how to add a TextVectorization layer to a NLP model to do preprocessing as part of the model itself. However, I found that that I couldn’t save my model in .hdf5 format if I used that layer. However, if I leave off the .hdf5 extension, then keras saves the model as a file directory of assets, and this works for the TextVectorization layer.
After fitting, we can reload our model for evaluation at its best performing epoch with:
model = keras.models.load_model(filepath)
Let’s see it all in action!
Data provided by Carnegie Mellon University via the Statlib library. Embedded Gist by author
I hope the ModelCheckpoint and the other callbacks that Keras provides and the ones you devise yourself will help you make the best predictive models!
In this article, you learned about callbacks for deep models using the Keras api for Tensorflow. Specifically, you learned how to use the ModelCheckpoint callback to save the best version of your model before it over-trains and a few ways to customize the callback.
As always, use your data science powers for good, and have lots of fun! Please leave any questions, comments, or kind and polite feedback in the comments.