Callbacks in Tensor Flow

This blog is going to be about using callbacks in Tensor Flow to do various things. Callbacks are basically special functions that you can invoke before, during and after the training process. We use them to improve our model.

Here are a few of the examples so you have a better idea of what these are,

Early Stopping: This callback can be used to stop training once a specific metric is reached. For example, say you want to create a model that has a training accuracy of 99.9%, so instead of manually running this and checking after each epoch, you can set a total number of epochs to run to say 1000 and stop the training once you achieve your desired accuracy.

Best model: Imagine you are trying to tweak different hyperparameters and you realize the best model was with say 55th epoch and you are still training after 60th. Would it be better if you could just save the best model as you train and use that later?

Learning rate Scheduler: As the name suggests why not try different learning rates during training, lets start with a big learning rate and we keep decreasing it either after every epoch or based on the error/accuracy rate. You can also quickly run this to figure out the best model based on different learning rates and use that in a much bigger training step.

Tensor Board: This in itself would take a whole blog (or even more) but I just want to include it here so we know how to see this in action.

I know we have a lot of things to cover, so let’s jump in. We follow our standard template that we have used till now.

1.Dataset: Load the data set, do some feature engineering if needed.
2. Build Model: Build a Tensor Flow model with various layers.
3. Compile Model: Here we compile the model, select the loss & Optimizer functions.
4. Fit Model: Here we finally train the model using the training data and get some metrics.
5. Evaluate Model: We check our model performance on the validation data.

Dataset:
We will use Fashion Minst dataset we used before.

Fashion Minst: This is a dataset of 70,000 images. These are small grey scaled images with a standard size of 28x28 pixels. Here are a few examples.

First, let’s look at how to load data. This is going to use

#Imports
 import tensorflow as tf
#Load Dataset
 mnist = tf.keras.datasets.fashion_mnist
(training_images, training_labels), 
        (test_images, test_labels) = mnist.load_data()
#Reshape & Scale
 training_images=training_images.reshape(60000, 28, 28, 1)
 training_images=training_images / 255.0
 test_images = test_images.reshape(10000, 28, 28, 1)
 test_images=test_images/255.0

This function loads all the 70k images and splits it into two parts 60k for training and 10k for testing.

Build Model:

model= Sequential([
                   Conv2D(16,(3,3),input_shape(28,28,1),
                          activation='relu'),
                   MaxPool2D(2,2),
                   Flatten(),
                   Dense(512,activation='relu'),
                   Dense(10,activation='softmax')
                  ])

This model using Conv2D layer and MaxPooling2D layer. To get a little bit more into the details:

Compile Model:

After we build the model we need to compile it. Here we need to select the loss functions and the optimizers. As you can see from the below code snippet this is very easy in Tensor Flow.

model.compile(
      optimizer=tf.keras.optimizers.SGD(lr=0.01),
      loss='sparse_categorical_crossentropy',
      metrics=['accuracy'])

Fit Model:

Before we run our Model.Fit we need to code up our callbacks.

Early Stopping: Here we stop the training after we achieve an accuracy of 95%. This is done using the custom callback function we define below.

#Custom callback function used to stop training based on accuracy
class myCallback(tf.keras.callbacks.Callback):
  def on_epoch_end(self, epoch, logs={}):
    if(logs.get('accuracy')>0.93):
      print("\nReached 93% accuracy so cancelling training!")
      self.model.stop_training = True
EarlystoppingAccuracy = myCallback()

Save Best Model: This callback saves the best model based on the accuracy.

SaveBestModel = tf.keras.callbacks.ModelCheckpoint(
                   filepath='/content/sample_data',
                   save_weights_only=True,
                   monitor='accuracy',
                   mode='max',
                   save_best_only=True)

Learning Rate Scheduler: For this one we first define how we want the learning rate to change based on epochs. We assume a big learning rate at the beginning and then we decrease the learning rate.

# Function to define the learning rate scheduler
def scheduler(epoch, lr):
  if epoch < 10:
     return lr
  else:
     return lr * tf.math.exp(-0.1)
# Here is the final callback
LearningRateScheduler = tf.keras.callbacks.
                           LearningRateScheduler(scheduler)

Tensor Board: Here is the code to create the tensor board callback.

import os
import time
root_logdir=os.path.join(os.curdir,'my_logs')
def get_run_logdir():
   run_id=time.strftime('run_%Y_%m_%d-%H_%M_%S')
   return os.path.join(root_logdir,run_id)
run_logdir = get_run_logdir()
Tensorboard= tf.keras.callbacks.
                        TensorBoard(run_logdir,histogram_freq=15)

Finally we call the fit method on the model with all the callbacks.

model.fit(train_rescale,train_label,epochs=100,
          validation_data=(test_rescale,test_label),
          callbacks=[
                      EarlystoppingAccuracy,
                      SaveBestModel,
                      LearningRateScheduler,
                      Tensorboard])

Here is what we have as the end results.

Showing Training stop as desired accuracy level is reached

You can see the tensor board using the following code in colab.

%load_ext tensorboard
%tensorboard --logdir=./my_logs --port=6006

Good luck !!!

References:

Keras callbacks:
https://www.tensorflow.org/api_docs/python/tf/keras/callbacks

Coursera link:
https://www.coursera.org/specializations/tensorflow-in-practice

Hands-On Machine Learning with Scikit-Learn, Keras & TensorFlow by Aurelien Geron