How to set a learning rate scheduler in PyTorch

Learning rate schedules in PyTorch change an optimizer's parameter-group learning rates while training progresses. They are useful when early updates need a larger step size and later epochs should settle with smaller updates.

A scheduler wraps the optimizer, so the optimizer still applies gradients and owns the param_groups values. StepLR is a simple epoch-level scheduler that multiplies the current learning rate by gamma after a fixed number of scheduler steps.

Most PyTorch schedulers should step after optimizer.step(). The smoke run uses CPU tensors, a tiny linear model, and a five-epoch loop so the expected and observed learning-rate sequences can be compared before the pattern is copied into a full training loop.

1. Create a scheduler smoke script.

   scheduler_demo.py

   import torch
   from torch import nn
    
    
   torch.manual_seed(11)
    
   features = torch.tensor(
       [
           [0.0, 0.5],
           [1.0, 1.5],
           [2.0, 2.5],
           [3.0, 3.5],
       ],
       dtype=torch.float32,
   )
   weights = torch.tensor([[0.4], [-0.2]])
   targets = features @ weights + 0.1
    
   model = nn.Linear(2, 1)
   loss_fn = nn.MSELoss()
   optimizer = torch.optim.SGD(model.parameters(), lr=0.05)
   scheduler = torch.optim.lr_scheduler.StepLR(
       optimizer,
       step_size=2,
       gamma=0.5,
   )
    
   print(f"initial lr: {scheduler.get_last_lr()[0]:.4f}")
    
   observed_lrs = []
    
   for epoch in range(1, 6):
       optimizer.zero_grad(set_to_none=True)
       loss = loss_fn(model(features), targets)
       loss.backward()
       optimizer.step()
       scheduler.step()
       current_lr = scheduler.get_last_lr()[0]
       observed_lrs.append(round(current_lr, 4))
    
       print(
           f"epoch {epoch}: "
           f"loss={loss.item():.4f}, "
           f"lr={current_lr:.4f}"
       )
    
   expected_lrs = [0.05, 0.025, 0.025, 0.0125, 0.0125]
   matched = observed_lrs == expected_lrs
    
   print(f"expected: {expected_lrs}")
   print(f"observed: {observed_lrs}")
   print(f"matched: {matched}")

   Here, step_size=2 keeps the initial rate for the first scheduler step, halves it on the second step, and halves it again on the fourth step.

2. Run the smoke script.

   $ python scheduler_demo.py
   initial lr: 0.0500
   epoch 1: loss=0.7962, lr=0.0500
   epoch 2: loss=0.2397, lr=0.0250
   epoch 3: loss=0.2241, lr=0.0250
   epoch 4: loss=0.2175, lr=0.0125
   epoch 5: loss=0.2111, lr=0.0125
   expected: [0.05, 0.025, 0.025, 0.0125, 0.0125]
   observed: [0.05, 0.025, 0.025, 0.0125, 0.0125]
   matched: True

3. Confirm the scheduler reached the intended rates.

   The expected and observed lines match, so scheduler.step() changed the optimizer learning rate at the configured step boundary.

4. Add the scheduler after creating the optimizer in training code.

   optimizer = torch.optim.SGD(model.parameters(), lr=0.05)
   scheduler = torch.optim.lr_scheduler.StepLR(
       optimizer,
       step_size=2,
       gamma=0.5,
   )

   Use a scheduler whose step interval matches the training plan. StepLR is commonly stepped once per epoch, while batch-level schedulers such as OneCycleLR are stepped after each training batch.

5. Step the scheduler after each optimizer update boundary.

   for epoch in range(num_epochs):
       for batch_features, batch_targets in loader:
           optimizer.zero_grad(set_to_none=True)
           predictions = model(batch_features)
           loss = loss_fn(predictions, batch_targets)
           loss.backward()
           optimizer.step()
    
       scheduler.step()
       current_lr = scheduler.get_last_lr()[0]
       print(f"epoch {epoch + 1}: lr={current_lr:.6f}")

   Calling scheduler.step() before optimizer.step() can skip the first scheduled learning-rate value in current PyTorch behavior.

6. Save the scheduler state with training checkpoints.

   torch.save(
       {
           "model": model.state_dict(),
           "optimizer": optimizer.state_dict(),
           "scheduler": scheduler.state_dict(),
       },
       "checkpoint.pt",
   )

   Restore the scheduler state with scheduler.load_state_dict(…) after recreating the same optimizer and scheduler objects.
   Related: How to save and restore a PyTorch training checkpoint

7. Remove the temporary smoke script.

   $ rm scheduler_demo.py