How to Use a Learning Rate Scheduler in Keras

This article provides a short tutorial on how you can use Learning Rate Scheduler's in Keras with code and interactive visualizations, using Weights & Biases.

Saurav Maheshkar

Created on July 1|Last edited on June 12

Comment

In this article, we'll look at how you can use a Learning Rate Scheduler in Keras for writing efficient and stable training loops.
Unlike PyTorch, which provides various classes, TensorFlow provides an easy class into which you can plug in your own functions to use various kinds of Learning Rate Schedulers, which we can be easily added into the model.fit() call!
For a closer look at the Learning Rate Scheduler callback available in Tensorflow you can refer to the official documentation.
💡
Here's what we'll be covering: 
Table of ContentsCodeSummaryRecommended Reading
﻿
CodeIt's very easy to implement custom schedules in TensorFlow because instead of providing strict class definitions, Tensorflow allows you to feed in custom functions into a Callback which is then fed into the .fit() function.
from wandb.keras import WandbCallback
﻿
# Model Definition
model = ...
﻿
# Compile the model
model.compile(optimizer = <>, loss = <>, metrics = <>)
﻿
# Custom Scheduler Function
lr_start   = 1e-4
lr_max     = 0.000015 * replicas * batch_size
lr_min     = 1e-7
r_ramp_ep = 3
lr_sus_ep  = 0
lr_decay   = 0.7
   
def lrfn(epoch):
    if epoch < lr_ramp_ep:
        lr = (lr_max - lr_start) / lr_ramp_ep * epoch + lr_start
            
    elif epoch < lr_ramp_ep + lr_sus_ep:
        lr = lr_max
            
    else:
        lr = (lr_max - lr_min) * lr_decay**(epoch - lr_ramp_ep - lr_sus_ep) + lr_min
            
    return lr
﻿
# Using this Custom Function, create a Callback
lr_callback = tf.keras.callbacks.LearningRateScheduler(lrfn, verbose=True)
﻿
# Train the Model
model.fit(..., callbacks=[lr_callback, WandbCallback()])
That's all it takes. You just need to create a function that returns the updated learning rate after each epoch, and then create a LearningRateScheduler Callback Instance.
SummaryIn this article, you saw how you could use a Learning Rate Scheduler in Keras-based deep learning models and how using Weights & Biases to monitor your metrics can lead to valuable insights.
To see the full suite of W&B features, please check out this short 5 minutes guide. If you want more reports covering the math and "from-scratch" code implementations, let us know in the comments down below or on our forum ✨!﻿
Check out these other reports on Fully Connected covering other fundamental development topics like GPU Utilization and Saving Models.
Recommended Reading
Setting Up TensorFlow And PyTorch Using GPU On Docker
A short tutorial on setting up TensorFlow and PyTorch deep learning models on GPUs using Docker.
How to Compare Keras Optimizers in Tensorflow for Deep Learning
A short tutorial outlining how to compare Keras optimizers for your deep learning pipelines in Tensorflow, with a Colab to help you follow along.
Preventing The CUDA Out Of Memory Error In PyTorch
A short tutorial on how you can avoid the "RuntimeError: CUDA out of memory" error while using the PyTorch framework.
How To Calculate Number of Model Parameters for PyTorch and TensorFlow Models
This article provides a short tutorial on calculating the number of parameters for TensorFlow and PyTorch deep learning models, with examples for you to follow.
How to Initialize Weights in PyTorch
A short tutorial on how you can initialize weights in PyTorch with code and interactive visualizations.
How to Properly Use PyTorch's CosineAnnealingWarmRestarts Scheduler
This article provides a short tutorial on how to use the CosineAnnealingWarmRestarts Scheduler in PyTorch, along with code and interactive visualizations.
﻿
﻿

Add a comment

Tags: Articles, Domain Agnostic, Beginner, Tutorial, Keras

Iterate on AI agents and models faster. Try Weights & Biases today.