Callbacks in neural networks
By: Risto Hinno
Tags: Callback, Fastai, Dynamic Programming, Pytorch, Neural Networks

Intro

While training machine learning model you want to play with it. Use different parameters, architectures, tricks and everything else that might improve performance. How to implement those changes? You have to change source code. But that will make codebase bigger, more difficult to maintain, more error-prone. Is there a better way? Yes! We could use callbacks.

Wikipedia defines callback as a “… any executable code that is passed as an argument to other code”. To my understanding, the main idea of callbacks is to inject code into existing program (functions, classes). Let’s use a simple example to see what callbacks are.

Full code examples for this post are available here.

Training loop without callbacks

We want to train a neural network where we fit batches of data. Full example code is available here. Training loop consists of our example main fit function, which uses:

dataloaders to generate data batches

model architecture which holds trainable parameters

loss function to calculate loss

optimizer to calculate derivatives of loss and do the backward step to update model weights

To do all the work in each batch we use one_batch function which:

does the forward step and gets current model prediction for batch

calculates loss by comparing model output with real labels

does the backward step by adjusting model parameter values

So far this is a simple classical model training example. In each step in fit and one_batch function we might want to use different adjustments to have better overview of the training, increase speed of training etc. Simplistic approach would be to update existing functions by adding new lines of code. But problem is that we don’t know which of the changes are really improving our model. If we hardcode them into our model we might have a bunch of redundant code not to mention the errors that might creep in. Callbacks give us a great way to inject new code into training loop without touching source code of the training loop.

Overview of callbacks

In our case callbacks are classes which have methods which are called in right step of the fit function. There is a base CallBack class which defines methods for individual callback (callbacks from 1 to n in the following graph). There is a CallBackHandler class which combines individual callbacks into one object which is given as an argument for Learner class which combines data, optimizer etc. to simplify training.

Usage of callbacks

After everything needed for training is packed into Learner class we could feed it into our fit function. Let’s have an example.

Training loop with callbacks

With callbacks we could adjust our training loop (full example is here and here):

What we have added some mysterious calls to some learner.cb (learner callback) functions. Learner itself is just a class which holds loss function, optimizer, dataloaders and callbacks. I have made similar adjustments to one_batch function and introduced all_batch function to simplify fit function:

Adding calls to learner.cb functions give us now freedom to define custom functions which could be called during training time without changing training loop source code. Each step in training loop we call callbacks which contain custom functionality we want to activate.

Callback

Callback is a class which defines functions which are called different stages of the training loop. Base class for callback is following (source):

This class itself doesn’t do much, if we inserted this code into training loop, it wouldn’t change anything. Reason for this class is to use it as a template for specific callbacks that modify training loop. For example code for callback counting and displaying number of batches during training is following:

Callback uses inheritance by using base CallBack class and only re-defining functions that are necessary for specific callback. BatchCounter defines begin_epoch to initialize batch counter and after_step to increment batch count and display number of batches processed. We have our first callback, but usually we (may) want to insert multiple callbacks, how to do that? We could use CallbackHandler .

CallbackHandler

CallbackHandler is another class whose main purpose is to collect all our callbacks and call them in right training phase. Source code for CallbackHandler could be found here. Beginning of the CallbackHandler is following:

During initialization all callbacks are stored. In different training phases every callback appropriate function is called (remember every callback inherits from base CallBack class, so if we haven’t defined specific functions in our callback they are inherited from base CallBack with default behaviour doing basically nothing). CallbackHandler checks if we should stop training or skip some step. This a reason why CallBack should return True or False .

Also, we refer to learner so that each individual callback has access to learner and its internal states (see method set_learn ).

Usage

Now we just need to insert callbacks to fit function (full example is here):

We just added list of callback(s) to CallbackHandler and gave it as a argument to learner. If we want to add more callbacks then we could just define our custom callback class and give it as an argument to CallbackHandler . We don’t have to change fit function anymore.

Summary

I hope this article demonstrated what callbacks are and how to use them in (neural network) training. They require some changes in current fit function but offer a lot of flexibility later during training. We could inject a lot of functionality into our model without touching source code.

References