Member-only story

A Keyframe-style Learning Rate Scheduler for PyTorch

When it comes to defining learning rate schedules in PyTorch, you have plenty of options. 15 different scheduler classes, to be exact. This offers good flexibility, but it sometimes means that making a small change to your learning rate schedule requires switching to a different scheduler or combining multiple schedulers.

David Gilbertson
Better Programming
10 min readMay 28, 2023

--

Photo by John Moeses Bauan on Unsplash

It struck me one day that defining how a learning rate should change over the course of a training run is similar to defining an animation, with a combination of:

  • Key frames: at timestep x the value should be y
  • Transitions: how to interpolate values between timesteps

So, I’ve created an LR scheduler that uses this keyframe concept to define a schedule as a sequence of frames and transitions.

I created this for personal use and have been using it for a little while. I find it quite useful and wonder if others might too, hence this article.

I won’t explain the code in this post, and instead just focus on what it does and how to use it. Here’s a gist with the code if you’d like to try it out. If you…

--

--

Responses (1)