Skip to content




class LinearWithWarmup(PolynomialDecay):
 | def __init__(
 |     self,
 |     optimizer: torch.optim.Optimizer,
 |     num_epochs: int,
 |     num_steps_per_epoch: int,
 |     warmup_steps: int = 100,
 |     last_epoch: int = -1
 | ) -> None

Implements a learning rate scheduler that increases the learning rate to lr during the first warmup_steps steps, and then decreases it to zero over the rest of the training steps.

In practice, this is a wrapper of PolynomialDecay with power=1 and end_learning_rate=0.


  • optimizer : torch.optim.Optimizer
    This argument does not get an entry in a configuration file for the object.
  • num_epochs : int
    The number of epochs in the experiment. this does NOT get an entry in the config.
  • num_steps_per_epoch : int
    The number of steps per epoch. this does NOT get an entry in the config.
  • warmup_steps : int
    The number of steps to linearly increase the learning rate.


Config for using the LinearWithWarmup Learning Rate Scheduler with warmup_steps set 100.

        "learning_rate_scheduler": {
            "type": "linear_with_warmup",
Note that you do NOT pass a optimizer, num_epochs, nor num_steps_per_epoch key to the Learning rate scheduler.