Building Custom Callbacks with Keras and TensorFlow 2 | by B. Chen | Towards Data Science

Callbacks are an important type of object in Keras and TensorFlow. They are designed to be able to monitor the model performance in metrics at certain points in the training run and perform some actions that might depend on those performances in metric values.
Keras has provided a number of built-in callbacks, for example, EarlyStopping, CSVLogger, ModelCheckpoint, LearningRateScheduler etc. Apart from these popular built-in callbacks, there is a base class called Callback which allows us to create our own callbacks and perform some custom actions. In this article, you will learn what is the Callback base class, what it can do, and how to build your own callbacks.
If you want to learn more about those built-in callbacks, please check out the following article:
towardsdatascience.com
Here is the outline of this article.
What is the Callback base class?
Creating a model for demonstration
1. Building a custom callback for training
2. Building a custom callback for testing
3. Building a custom callback for prediction
4. Custom callback applications
Please check out Notebook for the source code.
In TensorFlow, all callbacks are stored in the tensorflow.keras.callbacks module. Inside that module, there is a base class called Callback which all other callbacks inherit from. You can also subclass the Callback base class yourself to create your own callbacks.
The base class Callback is constructed with the following methods that will be called at the appropriate time.
Global methods
on_(train|test|predict)_begin(self, logs=None) : called at the beginning of fit(), evaluate(), and predict().
on_(train|test|predict)_end(self, logs=None) : called at the end of fit(), evaluate(), and predict().
Batch-level method (training, testing, and predicting)
on_(train|test|predict)_batch_begin(self, batch, logs=None): Called right before processing a batch during training/testing/predicting.
on_(train|test|predict)_batch_end(self, batch, logs=None): Called at the end of training/testing/predicting a batch. Within this method, logs is a dict containing the metrics results.
Epoch-level method (training only)
on_epoch_begin(self, epoch, logs=None): Called at the beginning of an epoch during training.
on_epoch_end(self, epoch, logs=None): Called at the end of an epoch during training.
For the demonstration purpose, we will build an image classifier to tackle Fashion MNIST, which is a dataset that has 70,000 grayscale images of 28-by-28 pixels with 10 classes.
Keras provides some utility functions to fetch and load common datasets, including Fashion MNIST. Let’s load Fashion MNIST
fashion_mnist = keras.datasets.fashion_mnist(X_train_full, y_train_full), (X_test, y_test) = fashion_mnist.load_data()
The dataset is already split into a training set and a test set. Here is the shape and data type of the training set:
>>> X_train_full.shape(60000, 28, 28)>>> X_train_full.dtypedtype('uint8')
We are going to train the neural network using Gradient Descent, we must scale the input feature down to the 0–1 range. And for faster training on a local machine, let’s just use the first 10,000 images.
X_train, y_train = X_train_full[:10000]/255.0, y_train_full[:10000]
Now let’s build the neural network! There are 3 ways to create a machine learning model with Keras and TensorFlow 2.0. Since we are building a simple fully connected neural network and for simplicity, let’s use the easiest way: Sequential Model with Sequential().
from tensorflow.keras.models import Sequentialfrom tensorflow.keras.layers import Dense, Flattendef create_model():     model = Sequential([        Flatten(input_shape=(28, 28)),        Dense(300, activation='relu'),        Dense(100, activation='relu'),        Dense(10, activation='softmax'),    ])    model.compile(        optimizer='sgd',         loss='sparse_categorical_crossentropy',         metrics=['accuracy']    )    return model
Our model has the following specifications:
The first layer (also known as the input layer) has the input_shape to set the input size (28, 28) which matches the training data. The input layer is a Flatten layer whose role is simply to convert each input image into a 1D array.
And then it is followed by 2 Dense layers, one with 300 units, and the other with 100 units. Both of them use the relu activation function.
The output Dense layer has 10 units and the softmax activation function.
model = create_model()model.summary()
Firstly, let’s import the base class Callback
from tensorflow.keras.callbacks import Callback
Then, we create a new class TrainningCallback() with the methods that will be called at the training time.
class TrainingCallback(Callback):    def on_train_begin(self, logs=None):        print("Starting training...")            def on_epoch_begin(self, epoch, logs=None):        print(f"Starting epoch {epoch}")    def on_train_batch_begin(self, batch, logs=None):        print(f"Training: Starting batch {batch}")            def on_train_batch_end(self, batch, logs=None):        print(f"Training: Finished batch {batch}, loss is {logs['loss']}")            def on_epoch_end(self, epoch, logs=None):        print(f"Finished epoch {epoch}, loss is {logs['loss']}, accuracy is {logs['accuracy']}")            def on_train_end(self, logs=None):        print("Finished training")
Next, to use it for training, we just need to pass TrainingCallback() to the callbacks argument in the model.fit() method.
history = model.fit(    X_train,     y_train,     epochs=2,     validation_split=0.20,     batch_size=4000,   // A large value for demo purpose    verbose=2,    callbacks=[TrainingCallback()])
We can see that TrainingCallback() get passed in a list to the callbacks argument. It is a list because in practice we might be passing a number of callbacks for performing different tasks, for example, debugging, learning rate schedules, or any other built-in callbacks.
By executing the statement, you should get an output like below:
Similarly, we can implement a custom callback TestingCallback() for testing.
class TestingCallback(Callback):        def on_test_begin(self, logs=None):        print("Starting testing ...")            def on_test_batch_begin(self, batch, logs=None):        print(f"Testing: Starting batch {batch}")        def on_test_batch_end(self, batch, logs=None):        print(f"Testing: Finished batch {batch}")            def on_test_end(self, logs=None):        print("Finished testing")
To use it for testing, pass it to the callbacks argument in the model.valutate() method.
model.evaluate(    X_test,     y_test,     verbose=False,     callbacks=[TestingCallback()],     batch_size=2000, // A large value for demo purpose)
By executing the statement, you should get an output like below:
Starting testing ...Testing: Starting batch 0Testing: Finished batch 0Testing: Starting batch 1Testing: Finished batch 1Testing: Starting batch 2Testing: Finished batch 2Testing: Starting batch 3Testing: Finished batch 3Testing: Starting batch 4Testing: Finished batch 4Finished testing[85.61210479736329, 0.8063]
Similarly, we can implement a custom callback PredictionCallback() for prediction.
class PredictionCallback(Callback):        def on_predict_begin(self, logs=None):        print("Starting prediction ...")        def on_predict_batch_begin(self, batch, logs=None):        print(f"Prediction: Starting batch {batch}")            def on_predict_batch_end(self, batch, logs=None):        print(f"Prediction: Finish batch {batch}")        def on_predict_end(self, logs=None):        print("Finished prediction")
To use it for prediction, we just need to pass it to the callbacks argument in the model.predict() method.
model.predict(    X_test,     verbose=False,     callbacks=[PredictionCallback()],     batch_size=2000,  // A large value for demo purpose)
By executing the statement, you should get an output like below:
Starting prediction ...Prediction: Starting batch 0Prediction: Finish batch 0Prediction: Starting batch 1Prediction: Finish batch 1Prediction: Starting batch 2Prediction: Finish batch 2Prediction: Starting batch 3Prediction: Finish batch 3Prediction: Starting batch 4Prediction: Finish batch 4Finished predictionarray([......])
A main application of callback is to perform some actions depend on performance metrics, for example:
Real-time plotting during training
Stop training when a metric has stopped improving
Save model at the end of every epoch
Adjust learning rate (or other hyperparameters) according to a defined schedule
etc
In this section, we are going to show you some examples of Keras custom callback applications.
Real-time plotting during training
Early stopping at minimum loss
Learning rate scheduling
This first example shows the creation of a Callback that shows a live, real-time update of loss as your training progresses.
import numpy as npimport tensorflow as tffig = plt.figure(figsize=(12,4))# Create plot inside the figureax = fig.add_subplot()ax.set_xlabel('Epoch #')ax.set_ylabel('loss')class TrainingPlot(tf.keras.callbacks.Callback):    def on_train_begin(self, logs={}):        # Initialize the lists for holding losses        self.losses = []        def on_epoch_end(self, epoch, logs={}):        # Append the losses to the lists        self.losses.append(logs['loss'])                # Plot        epochs = np.arange(0, len(self.losses))        ax.plot(epochs, self.losses, "b-")        fig.canvas.draw()
self.losses=[] is initialized in on_train_begin() for holding losses. The real-time plot is implemented at the end of each epochon_epoch_end() by calling ax.plot(epochs, self.losses, "b-") and fig.canvas.draw() .
To use it for training, we just need to pass it to the callbacks argument in the model.fit() method.
model = create_model()history = model.fit(    X_train,     y_train,     epochs=20,     validation_split=0.20,     batch_size=64,     verbose=2,    callbacks=[TrainingPlot()])
By executing the statement, you should get a live real-time output like below:
Note: this example is originally from Keras guide “Writing your own callbacks”, please check out the official documentation for details.
This 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.
To use it for training:
model = create_model()history = model.fit(    X_train,     y_train,     epochs=50,     validation_split=0.20,     batch_size=64,     verbose=2,    callbacks=[EarlyStoppingAtMinLoss()])
In this run, our training was stopped at epoch 28 as the minimum of loss has been reached
Note: this example is originally from Keras guide “Writing your own callbacks”, please check out the official documentation for details.
This example shows how a custom Callback can be used to dynamically change the learning rate of the optimizer during the course of training.
To use it
model = create_model()history = model.fit(    X_train,     y_train,     epochs=15,     validation_split=0.20,     batch_size=64,     verbose=2,    callbacks=[CustomLearningRateScheduler(lr_schedule)])
Keras provides a base class called Callback which allows us to subclass it and create our own callbacks. It is really useful for debugging and perform some actions depend on performance metrics.
I hope this article will help you to save time in creating your own custom callback and perform some custom actions. I recommend you to check out the documentation for the Callbacks API and to know about other things you can do.
Thanks for reading. Please check out the notebook for the source code and stay tuned if you are interested in the practical aspect of machine learning.
A practical introduction to Keras Callbacks in TensorFlow 2
Learning Rate Schedules in practice
The Google’s 7 steps of Machine Learning in practice: a TensorFlow example for structured data
3 ways to create a Machine Learning Model with Keras and TensorFlow 2.0
Batch normalization in practice: an example with Keras and TensorFlow 2.0
Early stopping in Practice: an example with Keras and TensorFlow
More can be found from my Github