Writing your own callbacks

Authors: Rick Chao, Francois Chollet Date created: 2019/03/20 Last modified: 2023/06/25 Description: Complete guide to writing new Keras callbacks.
View in Colab • GitHub source
A callback is a powerful tool to customize the behavior of a Keras model during training, evaluation, or inference. Examples include keras.callbacks.TensorBoard to visualize training progress and results with TensorBoard, or keras.callbacks.ModelCheckpoint to periodically save your model during training.
In this guide, you will learn what a Keras callback is, what it can do, and how you can build your own. We provide a few demos of simple callback applications to get you started.
All callbacks subclass the keras.callbacks.Callback class, and override a set of methods called at various stages of training, testing, and predicting. Callbacks are useful to get a view on internal states and statistics of the model during training.
You can pass a list of callbacks (as the keyword argument callbacks) to the following model methods:
Called at the beginning of fit/evaluate/predict.
Called at the end of fit/evaluate/predict.
Called right before processing a batch during training/testing/predicting.
Called at the end of training/testing/predicting a batch. Within this method, logs is a dict containing the metrics results.
Called at the beginning of an epoch during training.
Called at the end of an epoch during training.
Let's take a look at a concrete example. To get started, let's import tensorflow and define a simple Sequential Keras model:
Then, load the MNIST data for training and testing from Keras datasets API:
Now, define a simple custom callback that logs:
The logs dict contains the loss value, and all the metrics at the end of a batch or epoch. Example includes the loss and mean absolute error.
In addition to receiving log information when one of their methods is called, callbacks have access to the model associated with the current round of training/evaluation/inference: self.model.
Here are a few of the things you can do with self.model in a callback:
Let's see this in action in a couple of examples.
This first example shows the creation of a Callback that stops training when the minimum of loss has been reached, by setting the attribute self.model.stop_training (boolean). Optionally, you can provide an argument patience to specify how many epochs we should wait before stopping after having reached a local minimum.
keras.callbacks.EarlyStopping provides a more complete and general implementation.
In this example, we show how a custom Callback can be used to dynamically change the learning rate of the optimizer during the course of training.
See callbacks.LearningRateScheduler for a more general implementations.
Be sure to check out the existing Keras callbacks by reading the API docs. Applications include logging to CSV, saving the model, visualizing metrics in TensorBoard, and a lot more!