How to Use StepLR in PyTorch for Learning Rate Scheduling

Use torch.optim.lr_scheduler.StepLR to decrease the learning rate by a factor every fixed number of epochs. Initialize it with your optimizer, step size, and decay factor, then call scheduler.step() after each epoch to update the learning rate.

Syntax

The StepLR scheduler is created by passing your optimizer, the number of epochs between each learning rate update (step_size), and the decay factor (gamma). You then call scheduler.step() after each epoch to adjust the learning rate.

• optimizer: The optimizer whose learning rate you want to schedule.
• step_size: Number of epochs before reducing the learning rate.
• gamma: Multiplicative factor to reduce the learning rate.

python

scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=10, gamma=0.1)

Example

This example shows how to use StepLR with a simple optimizer. The learning rate starts at 0.1 and is reduced by 10 times every 5 epochs. The code prints the learning rate each epoch to show the changes.

python

import torch
import torch.nn as nn
import torch.optim as optim

# Dummy model
model = nn.Linear(10, 1)

# Optimizer with initial lr=0.1
optimizer = optim.SGD(model.parameters(), lr=0.1)

# StepLR scheduler: reduce lr by 0.1 every 5 epochs
scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=5, gamma=0.1)

for epoch in range(15):
    # Training code would go here
    
    # Print current learning rate
    lr = optimizer.param_groups[0]['lr']
    print(f"Epoch {epoch+1}: Learning Rate = {lr}")
    
    # Step the scheduler at the end of epoch
    scheduler.step()

Output

Epoch 1: Learning Rate = 0.1
Epoch 2: Learning Rate = 0.1
Epoch 3: Learning Rate = 0.1
Epoch 4: Learning Rate = 0.1
Epoch 5: Learning Rate = 0.1
Epoch 6: Learning Rate = 0.010000000000000002
Epoch 7: Learning Rate = 0.010000000000000002
Epoch 8: Learning Rate = 0.010000000000000002
Epoch 9: Learning Rate = 0.010000000000000002
Epoch 10: Learning Rate = 0.010000000000000002
Epoch 11: Learning Rate = 0.0010000000000000002
Epoch 12: Learning Rate = 0.0010000000000000002
Epoch 13: Learning Rate = 0.0010000000000000002
Epoch 14: Learning Rate = 0.0010000000000000002
Epoch 15: Learning Rate = 0.0010000000000000002

Common Pitfalls

Not calling scheduler.step() at the right time: It should be called once per epoch, usually at the end. Calling it every batch will reduce the learning rate too fast.

Forgetting to pass the optimizer to the scheduler: The scheduler needs the optimizer to update its learning rate.

Misunderstanding step_size: It is the number of epochs between learning rate changes, not the total number of epochs.

python

import torch.optim as optim

# Wrong: calling scheduler.step() every batch inside training loop
for epoch in range(3):
    for batch in range(10):
        # training step
        optimizer.step()
        scheduler.step()  # WRONG: called every batch

# Right: call scheduler.step() once per epoch
for epoch in range(3):
    for batch in range(10):
        # training step
        optimizer.step()
    scheduler.step()  # Correct: called once per epoch

Quick Reference

Parameter	Description	Example
optimizer	The optimizer to adjust learning rate for	optim.SGD(model.parameters(), lr=0.1)
step_size	Epochs between learning rate decay	5
gamma	Decay factor for learning rate	0.1
scheduler.step()	Call once per epoch to update lr	scheduler.step()

Key Takeaways

Initialize StepLR with optimizer, step_size (epochs), and gamma (decay factor).

Call scheduler.step() once per epoch to update the learning rate.

StepLR reduces learning rate by multiplying it by gamma every step_size epochs.

Avoid calling scheduler.step() every batch to prevent too rapid decay.

Check learning rate during training via optimizer.param_groups[0]['lr'].