Improved ModelCheckpoint Callback

Remove previously saved models with early training stop feature

We all have come to appreciate the flexibility afforded to us by keras’ various in-built callbacks like TensorBoard, LearningRateScheduler etc. during model training. However, it is not uncommon either for them to occasionally leave us wanting. In one such instance, I looked for functionality in ModelCheckpoint to delete the saved model files.

The motivation being that I was working on a shared machine and running multiple training experiments, each running 100s of training epochs. It did not take long for me to end up filling GBs and GBs of storage space. So, to avoid this, I thought it will be nice to have this added feature in this callback but with a quick internet search I found that users have already been asking for it.

Simultaneously, I felt that there is no point in having EarlyStopping separate from ModelCheckpoint, since they are both executing similar checks. Hence, I share my version of ModelCheckpoint with the following added functionalities:

  • Combine the functionality of EarlyStopping within the same callback
  • Option to remove the saved model files from previous epochs for the same training session

Below is the code. Please feel free to share and experiment.

The above code block is bare minimum. One can add further features like

  • how many model files to save (eg. last three best models)
  • save models with multiple different criteria (multiple performance metrics)

Budding AI Researcher

Get the Medium app

A button that says 'Download on the App Store', and if clicked it will lead you to the iOS App store
A button that says 'Get it on, Google Play', and if clicked it will lead you to the Google Play store