iwantcoding.com
🔥 Daily 👥 Rooms 🏆 Top Log in Sign up

torch.optim

By · Reviewed & last updated · Part of PyTorch Tutorial

PyTorch’s optimisers live in torch.optim. The trio that matters daily: Adam, AdamW, SGD with momentum. Pair them with a learning-rate scheduler.

Optimiser + scheduler + clipping

EXAMPLE
import torch
import torch.nn as nn
import torch.optim as optim
from torch.optim.lr_scheduler import CosineAnnealingLR, OneCycleLR, ReduceLROnPlateau

# 1) Default picks
opt = optim.AdamW(model.parameters(), lr=3e-4, weight_decay=1e-2)
opt = optim.Adam (model.parameters(), lr=1e-3)
opt = optim.SGD  (model.parameters(), lr=0.1, momentum=0.9, nesterov=True)

# 2) Per-parameter-group settings
opt = optim.AdamW([
    { 'params': model.backbone.parameters(), 'lr': 1e-4 },
    { 'params': model.head.parameters(),     'lr': 1e-3 },
], weight_decay=1e-2)

# Often: decay everything except biases + LayerNorm
no_decay = ['bias', 'LayerNorm.weight']
params_decay     = [p for n, p in model.named_parameters() if not any(nd in n for nd in no_decay)]
params_no_decay  = [p for n, p in model.named_parameters() if     any(nd in n for nd in no_decay)]

opt = optim.AdamW([
    { 'params': params_decay,    'weight_decay': 0.01 },
    { 'params': params_no_decay, 'weight_decay': 0.0  },
], lr=3e-4)

# 3) LR schedulers
# CosineAnnealing — smooth decay over a fixed number of epochs
sched = CosineAnnealingLR(opt, T_max=epochs)

# OneCycle — warm up + decay + ramp down in one schedule (fast convergence)
sched = OneCycleLR(opt,
    max_lr=3e-4,
    epochs=epochs,
    steps_per_epoch=len(train_loader),
)

# ReduceLROnPlateau — drop LR when val plateaus
sched = ReduceLROnPlateau(opt, mode='min', factor=0.5, patience=3)

# 4) Training loop boilerplate
for epoch in range(epochs):
    model.train()
    for xb, yb in train_loader:
        xb, yb = xb.to(device), yb.to(device)
        opt.zero_grad()
        loss = loss_fn(model(xb), yb)
        loss.backward()

        # Gradient clipping — stop loss spikes
        torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)

        opt.step()
        if isinstance(sched, OneCycleLR):
            sched.step()                # OneCycle steps per BATCH

    if not isinstance(sched, OneCycleLR):
        sched.step()                    # most schedulers step per EPOCH

    if isinstance(sched, ReduceLROnPlateau):
        sched.step(val_loss)            # this one steps with the metric

Why it matters

Beware: ReduceLROnPlateau wants the validation metric; CosineAnnealingLR wants nothing; OneCycleLR wants a step per batch. Mixing them up silently breaks training — lock the pattern early.

Tip: Tweak the snippet with Try it Yourself », then sit the quiz at the bottom of the page.

Example

Example
import torch
from torch import nn, optim
model = nn.Linear(10, 1)
opt = optim.AdamW(model.parameters(), lr=1e-3, weight_decay=1e-4)
sched = optim.lr_scheduler.CosineAnnealingLR(opt, T_max=10)
Try it Yourself »

Exercise

Apply an update step.

opt. ()

Discussion

Loading…