PyTorch Learning Rate Scheduler

The learning rate (Learning Rate) is one of the most important hyperparameters in neural network training. If the learning rate is too large, training oscillates or even diverges; if it is too small, convergence is extremely slow and it is easy to fall into local optima.

Learning Rate SchedulerBy dynamically adjusting the learning rate during training, it balances fast convergence in the early stage and fine tuning in the late stage.

1. Basic Concepts and Usage Patterns

Standard Usage Flow

)

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

model     = nn.Linear(10, 1)
optimizer = optim.SGD(model.parameters(), lr=0.1)

# 1. Create the scheduler, passing in the optimizer
scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=30, gamma=0.1)

for epoch in range(100):
    # 2. Train
    train(model, optimizer)

    # 3. After each epoch, call scheduler.step()
    scheduler.step()

    # 4. View the current learning rate
    current_lr = scheduler.get_last_lr()[0]
    print(f"Epoch {epoch+1}, LR: {current_lr:.6f}")

When to call step()

Example

# Correct: optimizer.step() first, scheduler.step() after
optimizer.step()
scheduler.step()

# Wrong: scheduler.step() before optimizer.step()
# After PyTorch 1.1.0, a warning will be generated, and some schedulers will behave abnormally
scheduler.step()
optimizer.step()

Note:optimizer.step()Must bescheduler.step()called before.

Viewing and Saving Learning Rate State

Example

# Get the current learning rate
current_lr = optimizer.param_groups[0]['lr']
current_lr = scheduler.get_last_lr()[0]   # The LR after the previous step()

# Save checkpoint (must also save the scheduler state)
torch.save({
    'epoch':          epoch,
    'model':          model.state_dict(),
    'optimizer':      optimizer.state_dict(),
    'scheduler':      scheduler.state_dict(),   # Don't miss this
}, 'checkpoint.pth')

# Restore checkpoint
ckpt = torch.load('checkpoint.pth')
model.load_state_dict(ckpt['model'])
optimizer.load_state_dict(ckpt['optimizer'])
scheduler.load_state_dict(ckpt['scheduler'])

2. Fixed Decay Schedulers

2.1 StepLR Step Decay

Every fixedstep_sizenumber of epochs, multiply the learning rate bygamma. It is one of the simplest and most commonly used schedulers.

Formula: lr = lr_base * gamma^(floor(step / step_size))

Example

optimizer = optim.SGD(model.parameters(), lr=0.1)
scheduler = optim.lr_scheduler.StepLR(
    optimizer,
    step_size=30,   # Decay every 30 epochs
    gamma=0.1       # Multiply by 0.1 each time (i.e., reduce to 1/10 of the original)
)

# LR changes:
# Epoch  0-29:  0.1
# Epoch 30-59:  0.01
# Epoch 60-89:  0.001
# Epoch 90+:    0.0001

Applicable scenarios:Tasks with a fixed training rhythm and clear stages, such as ResNet training on ImageNet (90 epochs, decaying at the 30th and 60th epochs).

2.2 MultiStepLR Multi-Milestone Decay

Decay the learning rate at several specified epochs (milestones), more flexible than StepLR.

Example

optimizer = optim.SGD(model.parameters(), lr=0.1)
scheduler = optim.lr_scheduler.MultiStepLR(
    optimizer,
    milestones=[30, 60, 80],   # Decay at the 30th, 60th, and 80th epochs
    gamma=0.1
)

# LR changes:
# Epoch  0-29:  0.1
# Epoch 30-59:  0.01
# Epoch 60-79:  0.001
# Epoch 80+:    0.0001

Applicable scenarios:When it is known at which epochs the model needs fine adjustment, such as the fine convergence stage in the middle and late phases of classification models.

2.3 ExponentialLR Exponential Decay

The learning rate decays every epoch and continues to decrease exponentially, giving smoother decay.

Formula: lr = lr_base * gamma^epoch

Example

optimizer = optim.Adam(model.parameters(), lr=0.01)
scheduler = optim.lr_scheduler.ExponentialLR(
    optimizer,
    gamma=0.95   # Multiply by 0.95 each epoch
)

# LR changes (first 5 epochs):
# Epoch 1: 0.0100
# Epoch 2: 0.0095
# Epoch 3: 0.0090
# Epoch 4: 0.0086
# Epoch 5: 0.0081

gammaIf set too small (e.g., 0.5), the learning rate will quickly approach zero; it is usually set between 0.9 and 0.99.

3. Adaptive Schedulers

3.1 ReduceLROnPlateau Metric-Based Decay

One of the most intelligent schedulers: monitor a certain metric (e.g., validation loss), and automatically reduce the learning rate when the metric stops improving. There is no need to know in advance at which epoch to decay.

Example

optimizer = optim.Adam(model.parameters(), lr=0.01)
scheduler = optim.lr_scheduler.ReduceLROnPlateau(
    optimizer,
    mode='min',       # 'min': the smaller the monitored value, the better (loss); 'max': the larger, the better (accuracy)
    factor=0.1,       # When triggered, lr = lr × factor
    patience=10,      # Number of epochs the metric is allowed to plateau; if exceeded, decay
    threshold=1e-4,   # An improvement smaller than this value is considered no improvement
    min_lr=1e-6,      # Lower bound of learning rate; it will not be lower than this value
    verbose=True      # Print decay information
)

for epoch in range(100):
    train_loss = train(model, optimizer)
    val_loss   = evaluate(model)

    # Unlike other schedulers, pass the monitored metric here
    scheduler.step(val_loss)

Writing style for monitoring Accuracy:

Example

scheduler = optim.lr_scheduler.ReduceLROnPlateau(
    optimizer,
    mode='max',       # The larger the accuracy, the better
    factor=0.5,
    patience=5,
)

scheduler.step(val_accuracy)

Applicable scenarios:The default first choice for almost all tasks, especially when you are not sure how many epochs to train or when training is unstable.

3.2 CosineAnnealingLR Cosine Annealing

The learning rate follows acosine curvefrom the initial value smoothly down to the minimum value (eta_min), avoiding the abrupt changes of step decay.

Formula: lr_t = eta_min + 0.5 * (eta_max - eta_min) * (1 + cos(t * pi / T_max)

Example

optimizer = optim.SGD(model.parameters(), lr=0.1)
scheduler = optim.lr_scheduler.CosineAnnealingLR(
    optimizer,
    T_max=100,     # Length of half a cycle (usually set to the total number of epochs)
    eta_min=1e-6   # Minimum learning rate (default 0)
)

# LR change trajectory (schematic when T_max=10):
# 0.1 -> 0.095 -> 0.079 -> 0.055 -> 0.026 -> 0.001
# (cosine curve smooth descent)

Applicable scenarios:Scenarios with a fixed number of training epochs, widely used in papers such as Vision Transformer and ResNet, with good convergence quality.

3.3 CosineAnnealingWarmRestarts

The upgraded version of cosine annealing supportsperiodic warm restarts(Warm Restarts): after each cycle ends, the learning rate resets to the initial value and starts a new round of cosine decay. This allows the model to escape local optima.

Example

optimizer = optim.SGD(model.parameters(), lr=0.1)
scheduler = optim.lr_scheduler.CosineAnnealingWarmRestarts(
    optimizer,
    T_0=10,       # Length of the first cycle (number of epochs)
    T_mult=2,     # Multiplier for the cycle length after each restart (1=equal length, 2=gradually increasing)
    eta_min=1e-6
)

# Cycle lengths when T_mult=2: 10 -> 20 -> 40 -> 80 ...
# LR changes (T_0=10, T_mult=1):
# 0.1 -> ... -> 0 -> 0.1 -> ... -> 0 -> 0.1 (restart every 10 epochs)

Applicable scenarios:When training large models, or when you want the model to choose the best among multiple convergence points, it works effectively with Snapshot Ensemble.

4. Warmup Schedulers

Warmuprefers to gradually increasing the learning rate from a very small value to the target value within the first few steps of training. Large batch size training and Transformer-like models almost all require warmup; otherwise, the initial gradient updates are too aggressive and the model is difficult to stabilize.

4.1 LinearLR Linear Scheduling

Linearly change the learning rate within a specified number of epochs (can be used for linear warmup or linear decay).

Example

optimizer = optim.Adam(model.parameters(), lr=0.01)

# Linear warmup: within the first 5 epochs, lr linearly increases from 0.01×0.1=0.001 to 0.01
warmup_scheduler = optim.lr_scheduler.LinearLR(
    optimizer,
    start_factor=0.1,   # initial lr = base_lr × start_factor
    end_factor=1.0,     # final lr = base_lr × end_factor
    total_iters=5       # Completed after 5 epochs
)

# LR changes: 0.001 -> 0.003 -> 0.005 -> 0.007 -> 0.009 -> 0.01

4.2 ConstantLR Constant Phase

Within a specified number of epochs, fix the learning rate tobase_lr × factor, and then restore the original value.

Example

# Use base_lr × 0.5 for the first 5 epochs, then restore to normal
scheduler = optim.lr_scheduler.ConstantLR(
    optimizer,
    factor=0.5,
    total_iters=5
)

4.3 SequentialLR Combined Scheduling

Combine multiple schedulersin sequence, which is the standard approach for implementing the "warmup + decay" combination strategy.

Example

optimizer = optim.Adam(model.parameters(), lr=0.01)

# Phase 1: warmup (first 5 epochs, LR linearly increases from 0.001 to 0.01)
warmup = optim.lr_scheduler.LinearLR(
    optimizer, start_factor=0.1, end_factor=1.0, total_iters=5
)

# Phase 2: cosine annealing (remaining 95 epochs)
cosine = optim.lr_scheduler.CosineAnnealingLR(
    optimizer, T_max=95, eta_min=1e-6
)

# Combined: run warmup first, switch to cosine after the 5th epoch
scheduler = optim.lr_scheduler.SequentialLR(
    optimizer,
    schedulers=[warmup, cosine],
    milestones=[5]          # Switch at the 5th epoch
)

# Usage is exactly the same as a normal scheduler
for epoch in range(100):
    train(...)
    scheduler.step()

Applicable scenarios:Standard configuration for Transformer training (warmup + cosine annealing); pretraining of BERT, GPT, and ViT all use this strategy.

5. Cyclic Schedulers

5.1 CyclicLR Cyclic Learning Rate

The learning ratebase_lrandmax_lrbetweencycles periodically, which helps the model explore a wider parameter space and escape saddle points.

Example

optimizer = optim.SGD(model.parameters(), lr=0.01)
scheduler = optim.lr_scheduler.CyclicLR(
    optimizer,
    base_lr=0.001,          # Lower bound of learning rate
    max_lr=0.01,            # Upper bound of learning rate
    step_size_up=2000,      # Number of iterations to rise from base_lr to max_lr
    step_size_down=2000,    # Number of iterations to fall from max_lr to base_lr (defaults to step_size_up)
    mode='triangular',      # Triangular cycle (constant amplitude)
    # mode='triangular2' # amplitude halves each cycle
    # mode='exp_range' # exponential amplitude decay
)

# CyclicLR is called per step (batch), not per epoch
for epoch in range(num_epochs):
    for inputs, labels in train_loader:
        optimizer.zero_grad()
        loss = criterion(model(inputs), labels)
        loss.backward()
        optimizer.step()
        scheduler.step()    # Call after each batch

Comparison of the three modes:

modeAmplitude changeCharacteristic
triangularConstantStable exploration, suitable for early stages
triangular2Halves each cycleExplore first, then converge
exp_rangeExponential decayFinally converges smoothly

5.2 OneCycleLR One-Cycle Strategy

One of the best-performing schedulers, proposed by fastai as the 1-Cycle Policy. The entire training has only one cycle: the learning rate first rises then falls, and momentum changes in the opposite direction.

Training is faster, typically requiring only1/5~1/10the number of epochs of traditional training.

Example

optimizer = optim.SGD(model.parameters(), lr=0.01, momentum=0.9)

steps_per_epoch = len(train_loader)

scheduler = optim.lr_scheduler.OneCycleLR(
    optimizer,
    max_lr=0.1,                              # Maximum learning rate
    steps_per_epoch=steps_per_epoch,         # Number of steps per epoch
    epochs=10,                               # Total number of epochs
    pct_start=0.3,                           # First 30% used for warmup rise
    anneal_strategy='cos',                   # Decay strategy ('cos' or 'linear')
    div_factor=25,                           # Initial lr = max_lr / div_factor
    final_div_factor=1e4                     # Final lr = max_lr / final_div_factor
)

# Initial lr = 0.1 / 25 = 0.004
# Peak lr = 0.1 (at 30%)
# Final lr = 0.1 / 10000 = 0.00001

# Also called per batch
for epoch in range(10):
    for inputs, labels in train_loader:
        optimizer.zero_grad()
        loss = criterion(model(inputs), labels)
        loss.backward()
        optimizer.step()
        scheduler.step()    # Call after each batch

Applicable scenarios:First choice when computational resources are limited and you need to quickly validate ideas;max_lrIt is recommended to use the LR Finder tool to determine the optimal value.

6. Custom Schedulers

6.1 LambdaLR Functional Customization

Fully customize the scheduling strategy by passing a Lambda function (which accepts the epoch number and returns a learning rate multiplier).

Example

optimizer = optim.Adam(model.parameters(), lr=0.01)

# Example 1: Transformer classic warmup strategy
# lr ∝ min(step^-0.5, step × warmup_steps^-1.5)
def transformer_lr(epoch, warmup_epochs=10, d_model=512):
    if epoch == 0:
        return 1e-7 / 0.01          # Avoid division by zero
    step = epoch + 1
    warmup = warmup_epochs
    return (d_model ** -0.5) * min(step ** -0.5, step * warmup ** -1.5) / 0.01

scheduler = optim.lr_scheduler.LambdaLR(optimizer, lr_lambda=transformer_lr)


# Example 2: polynomial decay (lr decreases linearly/polynomially from the initial value to 0)
def polynomial_decay(epoch, total_epochs=100, power=1.0):
    return max((1 - epoch / total_epochs) ** power, 0.0)

scheduler = optim.lr_scheduler.LambdaLR(optimizer, lr_lambda=polynomial_decay)


# Example 3: grouped learning rates (different strategies for different parameter groups)
optimizer = optim.Adam([
    {'params': model.backbone.parameters(), 'lr': 1e-4},
    {'params': model.head.parameters(),     'lr': 1e-3},
])

# One lambda function per parameter group
scheduler = optim.lr_scheduler.LambdaLR(
    optimizer,
    lr_lambda=[
        lambda epoch: 0.95 ** epoch,    # backbone: slow decay
        lambda epoch: 0.85 ** epoch,    # head: fast decay
    ]
)

6.2 Inheriting LRScheduler

When more complex logic is needed, inheritLRScheduler(PyTorch >= 2.0, in older versions_LRScheduler) to implement full customization:

Example

from torch.optim.lr_scheduler import LRScheduler
import math

class WarmupCosineScheduler(LRScheduler):
    """
Warmup + cosine annealing combined scheduler (manually implemented version)
- Linear warmup for the first warmup_epochs epochs
- Then cosine annealing to min_lr
    """
    def __init__(self, optimizer, warmup_epochs, total_epochs,
                 min_lr=1e-6, last_epoch=-1):
        self.warmup_epochs = warmup_epochs
        self.total_epochs  = total_epochs
        self.min_lr        = min_lr
        super().__init__(optimizer, last_epoch)

    def get_lr(self):
        epoch = self.last_epoch

        # Warmup phase: linear increase
        if epoch < self.warmup_epochs:
            warmup_factor = (epoch + 1) / self.warmup_epochs
            return [base_lr * warmup_factor for base_lr in self.base_lrs]

        # Cosine annealing phase
        progress = (epoch - self.warmup_epochs) / (
            self.total_epochs - self.warmup_epochs
        )
        cosine_factor = 0.5 * (1 + math.cos(math.pi * progress))
        return [
            self.min_lr + (base_lr - self.min_lr) * cosine_factor
            for base_lr in self.base_lrs
        ]


optimizer = optim.Adam(model.parameters(), lr=0.01)
scheduler = WarmupCosineScheduler(
    optimizer,
    warmup_epochs=10,
    total_epochs=100,
    min_lr=1e-6
)

7. Scheduler Visualization Comparison

The following code can plot the LR change curves of each scheduler for intuitive comparison:

Example

import torch
import torch.optim as optim
import matplotlib.pyplot as plt

def simulate_lr(scheduler_fn, epochs=100, steps_per_epoch=None):
    """Simulate and record the scheduler's learning rate changes"""
    model     = torch.nn.Linear(1, 1)
    optimizer = optim.SGD(model.parameters(), lr=0.1)
    scheduler = scheduler_fn(optimizer)
    lrs = []

    if steps_per_epoch:
        # Scheduler called per step
        for _ in range(epochs):
            for _ in range(steps_per_epoch):
                optimizer.step()
                scheduler.step()
                lrs.append(optimizer.param_groups[0]['lr'])
    else:
        # Scheduler called per epoch
        for _ in range(epochs):
            optimizer.step()
            scheduler.step()
            lrs.append(optimizer.param_groups[0]['lr'])
    return lrs

schedulers = {
    'StepLR(step=30, gamma=0.1)':
        lambda opt: optim.lr_scheduler.StepLR(opt, 30, 0.1),
    'CosineAnnealingLR':
        lambda opt: optim.lr_scheduler.CosineAnnealingLR(opt, T_max=100),
    'ExponentialLR(gamma=0.95)':
        lambda opt: optim.lr_scheduler.ExponentialLR(opt, 0.95),
    'CosineWarmRestarts(T0=25)':
        lambda opt: optim.lr_scheduler.CosineAnnealingWarmRestarts(opt, T_0=25),
}

plt.figure(figsize=(12, 5))
for name, fn in schedulers.items():
    plt.plot(simulate_lr(fn), label=name)

plt.xlabel('Epoch')
plt.ylabel('Learning Rate')
plt.title('PyTorch LR Scheduler Comparison')
plt.legend()
plt.grid(True, alpha=0.3)
plt.tight_layout()
plt.savefig('lr_schedulers.png', dpi=150)
plt.show()

8. Complete Training Template

A production-grade training template integrating warmup + cosine annealing, model saving and restoring:

Example

import torch
import torch.nn as nn
import torch.optim as optim
from torch.optim.lr_scheduler import SequentialLR, LinearLR, CosineAnnealingLR

# Hyperparameters
EPOCHS       = 100
WARMUP_EPOCHS = 5
BASE_LR      = 1e-3
MIN_LR       = 1e-6
SAVE_PATH    = 'best_model.pth'

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

# Model / Optimizer / Scheduler
model     = MyModel().to(device)
optimizer = optim.AdamW(model.parameters(), lr=BASE_LR, weight_decay=1e-4)
criterion = nn.CrossEntropyLoss(label_smoothing=0.1)

# Warmup 5 epochs + cosine annealing 95 epochs
warmup_sched = LinearLR(optimizer, start_factor=0.1, total_iters=WARMUP_EPOCHS)
cosine_sched = CosineAnnealingLR(optimizer, T_max=EPOCHS - WARMUP_EPOCHS, eta_min=MIN_LR)
scheduler    = SequentialLR(optimizer, [warmup_sched, cosine_sched], milestones=[WARMUP_EPOCHS])

# Restore checkpoint
start_epoch = 0
best_acc    = 0.0

try:
    ckpt = torch.load(SAVE_PATH, map_location=device)
    model.load_state_dict(ckpt['model'])
    optimizer.load_state_dict(ckpt['optimizer'])
    scheduler.load_state_dict(ckpt['scheduler'])
    start_epoch = ckpt['epoch'] + 1
    best_acc    = ckpt['best_acc']
    print(f"Resumed from Epoch {start_epoch}, best accuracy {best_acc:.4f}")
except FileNotFoundError:
    print("Training from scratch")

# Training loop
history = {'train_loss': [], 'val_acc': [], 'lr': []}

for epoch in range(start_epoch, EPOCHS):
    # Train
    model.train()
    total_loss = 0.0
    for inputs, labels in train_loader:
        inputs, labels = inputs.to(device), labels.to(device)
        optimizer.zero_grad()
        loss = criterion(model(inputs), labels)
        loss.backward()
        torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)  # Gradient clipping
        optimizer.step()
        total_loss += loss.item()

    # Validate
    model.eval()
    correct = 0
    with torch.no_grad():
        for inputs, labels in val_loader:
            inputs, labels = inputs.to(device), labels.to(device)
            correct += (model(inputs).argmax(1) == labels).sum().item()

    avg_loss = total_loss / len(train_loader)
    val_acc  = correct / len(val_loader.dataset)
    cur_lr   = scheduler.get_last_lr()[0]

    # Record & print
    history['train_loss'].append(avg_loss)
    history['val_acc'].append(val_acc)
    history['lr'].append(cur_lr)
    print(f"Epoch {epoch+1:3d}/{EPOCHS} | "
          f"Loss: {avg_loss:.4f} | Acc: {val_acc:.4f} | LR: {cur_lr:.2e}")

    # Scheduler update
    scheduler.step()

    # Save the best model
    if val_acc > best_acc:
        best_acc = val_acc
        torch.save({
            'epoch':      epoch,
            'model':      model.state_dict(),
            'optimizer':  optimizer.state_dict(),
            'scheduler':  scheduler.state_dict(),
            'best_acc':   best_acc,
        }, SAVE_PATH)
        print(f" Saving best model, Acc: {best_acc:.4f}")

print(f"\n"Training complete, best validation accuracy: {best_acc:.4f}")

9. Scheduler Selection Guide

Recommendations by Task Type

ScenarioRecommended schedulerReason
Quick experiments / prototype validationReduceLROnPlateauNo hyperparameter tuning needed, adaptive decay
Image classification (fixed epochs)CosineAnnealingLRSmooth decay, good convergence quality
Transformer / BERTLinearLR + CosineAnnealingLRWarmup is essential
Limited resources, fast trainingOneCycleLRReach the same accuracy in 1/5 of the time
Fine-tuning large models (fine-tune)LinearLR(warmup) +ConstantLRLow learning rate for stable fine-tuning
Uncertain number of epochsReduceLROnPlateauAutomatically responds, does not rely on a fixed schedule
Object detection (YOLO/Faster-RCNN)MultiStepLR / OneCycleLRHas clear decay points or pursues speed
Want to escape local optimaCosineAnnealingWarmRestartsPeriodic restart to explore parameter space

Common Configuration Combinations

Example

# Combination 1: Warmup + cosine annealing (standard for Transformer)
SequentialLR([LinearLR(5 epochs), CosineAnnealingLR(95 epochs)])

# Combination 2: monitor validation set + automatic decay (universally preferred)
ReduceLROnPlateau(mode='min', patience=5, factor=0.5)

# Combination 3: OneCycleLR (preferred for fast training)
OneCycleLR(max_lr=0.1, total_steps=total_steps, pct_start=0.3)

# Combination 4: traditional step decay (CV classic)
MultiStepLR(milestones=[30, 60, 90], gamma=0.1)

Quick Reference of Notes

ProblemSolution
Scheduler does not take effectCheckscheduler.step()whether it isoptimizer.step() afterCall
CyclicLR / OneCycleLR doesn't changeThese two are based onbatchcalls, not by epoch
ReduceLROnPlateau doesn't triggerCheckmodewhether it's set correctly (use loss for'min', use accuracy for'max'）
LR is incorrect after resuming trainingCheck whether it was saved and loadedscheduler.state_dict()
LR settings for multiple parameter groupsUseLambdaLRPass in a list, one lambda per parameter group
Other extensions