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

ClassAlgorithmKey Hyperparams
SGDStochastic Gradient Descentlr, momentum, weight_decay, nesterov
AdamAdaptive Moment Estimationlr, betas, eps, weight_decay, amsgrad
AdamWAdam + decoupled weight decaylr, betas, eps, weight_decay, amsgrad
AdamaxAdam with L∞ normlr, betas, eps, weight_decay
NAdamAdam + Nesterovlr, betas, eps, weight_decay, momentum_decay
RAdamRectified Adamlr, betas, eps, weight_decay
RMSpropRoot Mean Square Proplr, alpha, eps, weight_decay, momentum, centered
AdagradAdaptive gradientlr, lr_decay, eps, weight_decay
AdadeltaAdaptive learning raterho, eps, weight_decay
ASGDAveraged SGDlr, lambd, alpha, t0, weight_decay
LBFGSLimited-memory BFGSlr, max_iter, max_eval, history_size
SparseAdamAdam for sparse tensorslr, 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 * param to 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

TaskOptimizerLRWeight Decay
Computer vision (from scratch)SGD + momentum 0.90.01–0.11e-4
Fine-tuning CNNAdam / AdamW1e-40.01
Transformer (from scratch)AdamW1e-4 with warmup0.01–0.1
Transformer (fine-tuning)AdamW1e-5 – 5e-50.01
Small MLP / tabularAdam1e-31e-5