Python / PyTorch Fundamentals Interview Questions
How do you implement and use learning rate schedulers in PyTorch?
A fixed learning rate throughout training is rarely optimal — too high late in training prevents fine convergence, while too low early on wastes time. PyTorch's torch.optim.lr_scheduler module adjusts the learning rate systematically as training progresses.
import torch import torch.nn as nn import torch.optim as optim model = nn.Linear(10, 1) optimizer = optim.AdamW(model.parameters(), lr=1e-3) # ââ StepLR: multiply lr by gamma every step_size epochs scheduler_step = optim.lr_scheduler.StepLR( optimizer, step_size=10, gamma=0.1 ) # ââ CosineAnnealingLR: smooth decay following a cosine curve scheduler_cos = optim.lr_scheduler.CosineAnnealingLR( optimizer, T_max=100, eta_min=1e-6 ) # ââ ReduceLROnPlateau: reduce lr when a metric stops improving scheduler_plateau = optim.lr_scheduler.ReduceLROnPlateau( optimizer, mode="min", factor=0.5, patience=5 ) # ââ OneCycleLR: warmup then decay â fast convergence ("super-convergence") n_epochs, steps_per_epoch = 10, 100 scheduler_1cycle = optim.lr_scheduler.OneCycleLR( optimizer, max_lr=1e-2, total_steps=n_epochs * steps_per_epoch, pct_start=0.3, # 30% of steps used for warmup ) # ââ Training loop with scheduler for epoch in range(100): train_one_epoch(model, train_loader, optimizer, loss_fn, device) val_loss, _ = validate(model, val_loader, loss_fn, device) scheduler_cos.step() # epoch-based scheduler â call once per epoch scheduler_plateau.step(val_loss) # metric-based â pass the metric value current_lr = optimizer.param_groups[0]["lr"] print(f"Epoch {epoch}: lr={current_lr:.6f}") # Note: OneCycleLR and some schedulers are called PER BATCH, not per epoch # for step in range(total_steps): # train_step(...) # scheduler_1cycle.step() # called inside the batch loop
| Scheduler | Behaviour | Call frequency |
|---|---|---|
| StepLR | Multiply lr by gamma every N epochs | Per epoch |
| CosineAnnealingLR | Smooth cosine decay | Per epoch |
| ReduceLROnPlateau | Reduce lr when validation metric plateaus | Per epoch, after computing metric |
| OneCycleLR | Warmup then decay in one cycle | Per batch/step |
| LinearLR / warmup schedules | Linear ramp from low to target lr | Per step, common for transformers |
More Related questions...