Learning Rate Schedule in Practice: an example with Keras and TensorFlow 2.0 | by B. Chen | Towards Data Science

One of the painful things about training a neural network is the sheer number of hyperparameters we have to deal with. For example
Learning rate
Momentum or the hyperparameters for Adam optimization algorithm
Number of layers
Number of hidden units
Mini-batch size
Activation function
etc
Among them, the most important parameter is the learning rate. If your learning rate is set to low, training will progress very slowly as you are making very tiny updates to the weights in your network. However, if your learning rate is set too high, it can cause undesirable divergent behavior in your loss function.
When training a neural network, it is often useful to reduce the learning rate as the training progresses. This can be done by using learning rate schedules or adaptive learning rate. In this article, we will focus on adding and customizing learning rate schedule in our machine learning model and look at examples of how we do them in practice with Keras and TensorFlow 2.0
Learning Rate Schedules seek to adjust the learning rate during training by reducing the learning rate according to a pre-defined schedule. The popular learning rate schedules include
Constant learning rateTime-based decayStep decayExponential decay
Constant learning rate
Time-based decay
Step decay
Exponential decay
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.
Please check out my Github repo for source code.
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'),    ])    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()
The constant learning rate is the default schedule in all Keras Optimizers. For example, in the SGD optimizer, the learning rate defaults to 0.01.
To use a custom learning rate, simply instantiate an SGD optimizer and pass the argument learning_rate=0.01 .
sgd = tf.keras.optimizers.SGD(learning_rate=0.01)model.compile(    optimizer=sgd,     loss='sparse_categorical_crossentropy',     metrics=['accuracy'])
And to fit the model to training data:
history_constant = model.fit(    X_train,     y_train,     epochs=100,     validation_split=0.2,    batch_size=64)
Let’s plot the model accuracy and this can serve as a baseline for us to experiment with other learning rate schedules.
Time-based decay is one of the most popular learning rate schedules. Formally, the time-based decay is defined as:
learning_rate = lr * 1 / (1 + decay * epoch)
where lr is the previous learning rate, decay is a hyperparameter and epoch is the iteration number. When the decay is zero, this has no effect on changing the learning rate. When the decay is specified, it will decrease the learning rate from the previous epoch by the given fixed amount. The value of decay is normally implemented as
decay = initial_learning_rate / num_of_epoches 
In Keras, one way to implement the time-based decay is by defining a time-based decay function lr_time_based_decay() and pass it to LearningRateScheduler callback.
initial_learning_rate = 0.01epochs = 100decay = initial_learning_rate / epochsdef lr_time_based_decay(epoch, lr):    return lr * 1 / (1 + decay * epoch)history_time_based_decay = model.fit(    X_train,     y_train,     epochs=100,     validation_split=0.2,    batch_size=64,    callbacks=[LearningRateScheduler(lr_time_based_decay, verbose=1)],)
And below are the plots of accuracy and learning rate.
Another popular learning rate schedule is to systematically drop the learning rate at specific times during training. Formally, it is defined as:
learning_rate = initial_lr * drop_rate^floor(epoch / epochs_drop)
Where initial_lr is the initial learning rate such as 0.01, the drop_rate is the amount that the learning rate is modified each time if it is changed, epoch is the current epoch number, and epochs_drop is how often to change the learning rate such as 10 epochs. Similarly, we can implement this by defining a step decay function lr_step_decay() and pass it to LearningRateScheduler callback.
initial_learning_rate = 0.01def lr_step_decay(epoch, lr):    drop_rate = 0.5    epochs_drop = 10.0    return initial_learning_rate * math.pow(drop_rate, math.floor(epoch/epochs_drop))# Fit the model to the training datahistory_step_decay = model.fit(    X_train,     y_train,     epochs=100,     validation_split=0.2,    batch_size=64,    callbacks=[LearningRateScheduler(lr_step_decay, verbose=1)],)
And below are the plots of the accuracy and learning rate.
Another popular learning rate schedule is to drop the learning rate at an exponential rate. Formally, it is defined as:
learning_rate = initial_lr * e^(−k * epoch)
Where initial_lr is the initial learning rate such as 0.01, k is a hyperparameter, and epoch is the current epoch number. Similarly, we can implement this by defining an exponential decay function lr_exp_decay() and pass it to LearningRateScheduler callback.
initial_learning_rate = 0.01def lr_exp_decay(epoch, lr):    k = 0.1    return initial_learning_rate * math.exp(-k*epoch)# Fit the model to the training datahistory_exp_decay = model.fit(    X_train,     y_train,     epochs=100,     validation_split=0.2,    batch_size=64,    callbacks=[LearningRateScheduler(lr_exp_decay, verbose=1)],)
Finally, let us compare the model accuracy using different learning rate schedules.
Looks like Constant and Time-based learning rates have better performance than Step decay and Exponential decay for this particular tutorial. Bear in mind that, this tutorial only uses the first 10,000 images with some arbitrary value forinitial_learning_rate=0.01, validation_split=0.2 and batch_size=64 .
In a real-world application, they are a lot more to consider for tunning the learning rate. Please check out the paper “Practical Recommendations for Gradient-based Training of Deep Architectures” for some best practices.
Thanks for reading. This article has covered the most popular Learning Rate Schedules. Next time, we will take a look at the Adaptive Learning Rate.
Please check out the notebook on my Github for the source code.
Stay tuned if you are interested in the practical aspect of machine learning.
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