PyTorch Cheatsheet
Optimizers
Use this PyTorch reference while you build software engineering projects, review code for technical interview prep, or polish examples for a software engineer resume.
Optimizer Basics
import torch.optim as optim optimizer = optim.Adam(model.parameters(), lr=1e-3) # Training step pattern optimizer.zero_grad() # clear accumulated gradients loss = criterion(output, target) loss.backward() # compute gradients optimizer.step() # update parameters
Built-in Optimizers
| Class | Algorithm | Key Hyperparams |
|---|---|---|
SGD | Stochastic Gradient Descent | lr, momentum, weight_decay, nesterov |
Adam | Adaptive Moment Estimation | lr, betas, eps, weight_decay, amsgrad |
AdamW | Adam + decoupled weight decay | lr, betas, eps, weight_decay, amsgrad |
Adamax | Adam with L∞ norm | lr, betas, eps, weight_decay |
NAdam | Adam + Nesterov | lr, betas, eps, weight_decay, momentum_decay |
RAdam | Rectified Adam | lr, betas, eps, weight_decay |
RMSprop | Root Mean Square Prop | lr, alpha, eps, weight_decay, momentum, centered |
Adagrad | Adaptive gradient | lr, lr_decay, eps, weight_decay |
Adadelta | Adaptive learning rate | rho, eps, weight_decay |
ASGD | Averaged SGD | lr, lambd, alpha, t0, weight_decay |
LBFGS | Limited-memory BFGS | lr, max_iter, max_eval, history_size |
SparseAdam | Adam for sparse tensors | lr, betas, eps |
SGD
optimizer = optim.SGD( model.parameters(), lr=0.01, momentum=0.9, # moving average of gradients weight_decay=1e-4, # L2 regularization nesterov=True, # Nesterov momentum (requires momentum > 0) dampening=0, # dampening for momentum maximize=False, foreach=None, # vectorized param updates (faster on GPU) differentiable=False, )
Adam
optimizer = optim.Adam( model.parameters(), lr=1e-3, betas=(0.9, 0.999), # (beta1, beta2) — 1st and 2nd moment decay eps=1e-8, # numerical stability term weight_decay=0, # L2 penalty (note: coupled — use AdamW instead) amsgrad=False, # AMSGrad variant (bounded learning rates) foreach=None, maximize=False, capturable=False, # for CUDA graph capture differentiable=False, fused=False, # fused CUDA kernel (faster, requires CUDA) )
AdamW (preferred over Adam for regularization)
optimizer = optim.AdamW( model.parameters(), lr=1e-3, betas=(0.9, 0.999), eps=1e-8, weight_decay=0.01, # decoupled weight decay (not gradient-scaled) amsgrad=False, maximize=False, foreach=None, capturable=False, differentiable=False, fused=False, )
Adam vs AdamW: Adam adds
weight_decay * paramto the gradient (scales with gradient magnitude). AdamW applies weight decay directly to parameters — mathematically cleaner and typically more effective. Use AdamW for fine-tuning transformers.
Per-Parameter Options
# Different learning rates for different parts of the model optimizer = optim.Adam([ {'params': model.backbone.parameters(), 'lr': 1e-4}, {'params': model.head.parameters(), 'lr': 1e-3}, {'params': model.embeddings.parameters(), 'lr': 5e-5, 'weight_decay': 0.0}, ], lr=1e-3) # default lr for any param group without explicit lr # Exclude bias and LayerNorm from weight decay (common in transformer training) no_decay = ['bias', 'LayerNorm.weight'] optimizer = optim.AdamW([ {'params': [p for n, p in model.named_parameters() if not any(nd in n for nd in no_decay)], 'weight_decay': 0.01}, {'params': [p for n, p in model.named_parameters() if any(nd in n for nd in no_decay)], 'weight_decay': 0.0}, ], lr=1e-4)
LBFGS (second-order, full-batch)
optimizer = optim.LBFGS( model.parameters(), lr=1.0, max_iter=20, history_size=100, line_search_fn='strong_wolfe', # or None ) # LBFGS requires a closure that evaluates the loss def closure(): optimizer.zero_grad() output = model(x) loss = criterion(output, y) loss.backward() return loss optimizer.step(closure)
Learning Rate Schedulers
from torch.optim import lr_scheduler # Step decay sched = lr_scheduler.StepLR(optimizer, step_size=10, gamma=0.1) # Multiply LR by gamma every step_size epochs # Multi-step decay sched = lr_scheduler.MultiStepLR(optimizer, milestones=[30, 60, 90], gamma=0.1) # Exponential decay sched = lr_scheduler.ExponentialLR(optimizer, gamma=0.95) # Cosine annealing sched = lr_scheduler.CosineAnnealingLR(optimizer, T_max=100, eta_min=0) # Cosine annealing with warm restarts (SGDR) sched = lr_scheduler.CosineAnnealingWarmRestarts(optimizer, T_0=10, T_mult=2, eta_min=0) # Linear warmup + cosine decay (using LambdaLR) def lr_lambda(current_step): warmup = 100 if current_step < warmup: return current_step / max(1, warmup) progress = (current_step - warmup) / max(1, total_steps - warmup) return max(0.0, 0.5 * (1.0 + math.cos(math.pi * progress))) sched = lr_scheduler.LambdaLR(optimizer, lr_lambda) # Reduce on plateau sched = lr_scheduler.ReduceLROnPlateau( optimizer, mode='min', # 'min' or 'max' factor=0.1, # multiply LR by this on plateau patience=10, # epochs without improvement before reduction threshold=1e-4, cooldown=0, min_lr=0, ) # One Cycle (Smith's 1-cycle policy) sched = lr_scheduler.OneCycleLR( optimizer, max_lr=1e-2, total_steps=total_steps, # or epochs * steps_per_epoch pct_start=0.3, # fraction of cycle for warmup anneal_strategy='cos', div_factor=25.0, # initial_lr = max_lr / div_factor final_div_factor=1e4, # min_lr = initial_lr / final_div_factor ) # Cyclic LR sched = lr_scheduler.CyclicLR(optimizer, base_lr=1e-4, max_lr=1e-2, step_size_up=2000, mode='triangular2') # Linear decay sched = lr_scheduler.LinearLR(optimizer, start_factor=1.0, end_factor=0.1, total_iters=100) # Polynomial decay sched = lr_scheduler.PolynomialLR(optimizer, total_iters=100, power=1.0) # Chain schedulers sched = lr_scheduler.SequentialLR(optimizer, schedulers=[warmup, cosine], milestones=[500]) # Combine with multiplication sched = lr_scheduler.ChainedScheduler([sched1, sched2])
Scheduler step placement
# Most schedulers: call after optimizer.step() for epoch in range(epochs): for batch in loader: optimizer.zero_grad() loss.backward() optimizer.step() sched.step() # per-step schedulers (OneCycleLR, CyclicLR) sched.step() # per-epoch schedulers (StepLR, CosineAnnealingLR) # ReduceLROnPlateau: call with metric sched.step(val_loss)
Optimizer State
optimizer.state_dict() # save state (moments, step count, etc.) optimizer.load_state_dict(state) # restore # State is per-parameter for group in optimizer.param_groups: print(group['lr']) # current learning rate per group group['lr'] = 1e-4 # modify learning rate directly # Get current LR current_lrs = [g['lr'] for g in optimizer.param_groups] sched.get_last_lr() # public API — LRs computed by the last step() # (sched.get_lr() is internal — it warns when called outside step())
Gradient Clipping (before optimizer.step)
# Clip by global L2 norm — most common torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0) # Clip by value torch.nn.utils.clip_grad_value_(model.parameters(), clip_value=0.5) # Canonical training step with clipping optimizer.zero_grad() loss.backward() torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0) optimizer.step() sched.step()
Gradient Accumulation
# Effective batch_size = batch_size * accumulate_steps accumulate_steps = 4 optimizer.zero_grad() for i, (x, y) in enumerate(loader): loss = model(x, y) / accumulate_steps loss.backward() if (i + 1) % accumulate_steps == 0: torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0) optimizer.step() optimizer.zero_grad() sched.step()
Common Hyperparameter Starting Points
| Task | Optimizer | LR | Weight Decay |
|---|---|---|---|
| Computer vision (from scratch) | SGD + momentum 0.9 | 0.01–0.1 | 1e-4 |
| Fine-tuning CNN | Adam / AdamW | 1e-4 | 0.01 |
| Transformer (from scratch) | AdamW | 1e-4 with warmup | 0.01–0.1 |
| Transformer (fine-tuning) | AdamW | 1e-5 – 5e-5 | 0.01 |
| Small MLP / tabular | Adam | 1e-3 | 1e-5 |