PyTorch Training Best Practices

Real-World Scenario

An ML engineer trains a text classification model on a single A100 GPU. Baseline training takes 6 hours per epoch. After enabling mixed precision (fp16), gradient accumulation, and OneCycleLR scheduling, training time drops to 2.5 hours per epoch with 1.8% better validation accuracy. These aren’t advanced tricks — they’re the standard training configuration in every modern PyTorch project.

Mixed Precision Training

import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader, TensorDataset
from torch.amp import autocast, GradScaler

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

# Synthetic dataset
rng = torch.Generator().manual_seed(42)
X = torch.randn(4000, 256, generator=rng)
y = torch.randint(0, 10, (4000,))

train_ds = TensorDataset(X[:3200], y[:3200])
val_ds   = TensorDataset(X[3200:], y[3200:])
train_loader = DataLoader(train_ds, batch_size=64, shuffle=True)
val_loader   = DataLoader(val_ds,   batch_size=128)

model = nn.Sequential(
    nn.Linear(256, 512), nn.ReLU(), nn.Dropout(0.2),
    nn.Linear(512, 512), nn.ReLU(), nn.Dropout(0.2),
    nn.Linear(512, 10),
).to(device)

optimizer = optim.AdamW(model.parameters(), lr=1e-3, weight_decay=1e-2)
criterion = nn.CrossEntropyLoss()

# GradScaler prevents underflow in fp16 gradients
scaler = GradScaler(device=device.type)

for epoch in range(5):
    model.train()
    for X_batch, y_batch in train_loader:
        X_batch, y_batch = X_batch.to(device), y_batch.to(device)

        optimizer.zero_grad()

        # autocast: runs forward pass in fp16/bf16 automatically
        with autocast(device_type=device.type):
            logits = model(X_batch)
            loss   = criterion(logits, y_batch)

        # scale loss to prevent fp16 underflow, then backward
        scaler.scale(loss).backward()
        # unscale gradients and clip before optimizer step
        scaler.unscale_(optimizer)
        torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
        scaler.step(optimizer)
        scaler.update()

    print(f"Epoch {epoch+1} complete")

Gradient Accumulation

import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader, TensorDataset
from torch.amp import autocast, GradScaler

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

X = torch.randn(2000, 128)
y = torch.randint(0, 5, (2000,))
loader = DataLoader(TensorDataset(X, y), batch_size=16, shuffle=True)

model = nn.Sequential(nn.Linear(128, 256), nn.ReLU(), nn.Linear(256, 5)).to(device)
optimizer = optim.AdamW(model.parameters(), lr=1e-3)
criterion = nn.CrossEntropyLoss()
scaler    = GradScaler(device=device.type)

ACCUMULATION_STEPS = 4   # effective batch = 16 × 4 = 64

optimizer.zero_grad()  # zero once before the loop
for step, (X_batch, y_batch) in enumerate(loader, 1):
    X_batch, y_batch = X_batch.to(device), y_batch.to(device)

    with autocast(device_type=device.type):
        logits = model(X_batch)
        # Divide loss by accumulation steps so the gradient scale is correct
        loss = criterion(logits, y_batch) / ACCUMULATION_STEPS

    scaler.scale(loss).backward()   # accumulate gradients

    if step % ACCUMULATION_STEPS == 0:
        scaler.unscale_(optimizer)
        torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
        scaler.step(optimizer)
        scaler.update()
        optimizer.zero_grad()   # clear after the effective batch

Learning Rate Scheduling

import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader, TensorDataset
import math

device = torch.device("cpu")

X = torch.randn(2000, 64)
y = torch.randint(0, 4, (2000,))
loader = DataLoader(TensorDataset(X, y), batch_size=32, shuffle=True)
n_steps = len(loader)

model = nn.Sequential(nn.Linear(64, 128), nn.ReLU(), nn.Linear(128, 4)).to(device)
optimizer = optim.AdamW(model.parameters(), lr=1e-3)

# OneCycleLR — best for single-run training: warmup + cosine annealing
scheduler_one = optim.lr_scheduler.OneCycleLR(
    optimizer,
    max_lr=1e-2,
    steps_per_epoch=n_steps,
    epochs=10,
    pct_start=0.3,       # 30% warmup, 70% decay
    anneal_strategy="cos",
)

# Cosine with warmup — popular in transformers
def cosine_with_warmup(step: int, warmup_steps: int, total_steps: int) -> float:
    if step < warmup_steps:
        return float(step) / float(max(1, warmup_steps))
    progress = float(step - warmup_steps) / float(max(1, total_steps - warmup_steps))
    return max(0.0, 0.5 * (1.0 + math.cos(math.pi * progress)))

WARMUP = n_steps * 2   # 2 epochs of warmup
TOTAL  = n_steps * 10  # 10 epochs total

scheduler_cos = optim.lr_scheduler.LambdaLR(
    optimizer,
    lr_lambda=lambda step: cosine_with_warmup(step, WARMUP, TOTAL)
)

# ReduceLROnPlateau — reduce when validation loss stops improving
scheduler_plateau = optim.lr_scheduler.ReduceLROnPlateau(
    optimizer,
    mode="min",
    patience=5,
    factor=0.5,
    verbose=True,
)

# Example training loop with plateau scheduler
criterion = nn.CrossEntropyLoss()
optimizer2 = optim.AdamW(model.parameters(), lr=1e-3)

for epoch in range(10):
    model.train()
    for X_b, y_b in loader:
        optimizer2.zero_grad()
        loss = criterion(model(X_b.to(device)), y_b.to(device))
        loss.backward()
        optimizer2.step()

    val_loss = loss.item()   # in practice: evaluate on validation set
    scheduler_plateau.step(val_loss)  # pass validation loss, not train loss

Training Checkpoint and Resume

import torch
import torch.nn as nn
import torch.optim as optim
from pathlib import Path

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

model     = nn.Sequential(nn.Linear(64, 128), nn.ReLU(), nn.Linear(128, 10)).to(device)
optimizer = optim.AdamW(model.parameters(), lr=1e-3)
scheduler = optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=50)

CHECKPOINT_DIR = Path("./checkpoints")
CHECKPOINT_DIR.mkdir(exist_ok=True)

def save_checkpoint(epoch: int, model, optimizer, scheduler, val_loss: float):
    path = CHECKPOINT_DIR / f"epoch_{epoch:03d}.pt"
    torch.save({
        "epoch":      epoch,
        "model":      model.state_dict(),
        "optimizer":  optimizer.state_dict(),
        "scheduler":  scheduler.state_dict(),
        "val_loss":   val_loss,
    }, path)
    print(f"Saved checkpoint: {path}")
    return path


def load_checkpoint(path: str, model, optimizer, scheduler):
    ckpt = torch.load(path, map_location=device)
    model.load_state_dict(ckpt["model"])
    optimizer.load_state_dict(ckpt["optimizer"])
    scheduler.load_state_dict(ckpt["scheduler"])
    return ckpt["epoch"], ckpt["val_loss"]


def get_best_checkpoint(directory: Path) -> Path | None:
    """Return the checkpoint with the lowest validation loss."""
    checkpoints = list(directory.glob("epoch_*.pt"))
    if not checkpoints:
        return None
    return min(checkpoints, key=lambda p: torch.load(p, map_location="cpu")["val_loss"])


# Training loop with best-model tracking
best_val_loss = float("inf")
criterion = nn.CrossEntropyLoss()
X_val = torch.randn(200, 64).to(device)
y_val = torch.randint(0, 10, (200,)).to(device)

for epoch in range(50):
    # ... training code ...

    with torch.no_grad():
        val_loss = criterion(model(X_val), y_val).item()

    scheduler.step()

    # Save every 5 epochs and whenever we improve
    if (epoch + 1) % 5 == 0 or val_loss < best_val_loss:
        save_checkpoint(epoch + 1, model, optimizer, scheduler, val_loss)
        if val_loss < best_val_loss:
            best_val_loss = val_loss
            print(f"New best: {val_loss:.4f}")

# Load best model for inference
best_path = get_best_checkpoint(CHECKPOINT_DIR)
if best_path:
    start_epoch, _ = load_checkpoint(str(best_path), model, optimizer, scheduler)
    print(f"Loaded best checkpoint from epoch {start_epoch}")

torch.compile — Instant Speedup in PyTorch 2.x

import torch
import torch.nn as nn

# torch.compile JIT-compiles the model using TorchDynamo + TorchInductor
# Typically 1.5–3x speedup on training and inference with zero code changes
model = nn.Sequential(
    nn.Linear(512, 1024), nn.ReLU(), nn.Dropout(0.2),
    nn.Linear(1024, 1024), nn.ReLU(),
    nn.Linear(1024, 10),
)

# mode options:
#   "default"      — good balance of speed and compile time
#   "reduce-overhead" — fastest inference, longer compile
#   "max-autotune"    — maximum optimization, slowest compile
compiled_model = torch.compile(model, mode="reduce-overhead")

# First call triggers compilation (one-time cost)
x = torch.randn(32, 512)
_ = compiled_model(x)   # compile happens here

# Subsequent calls use compiled code
import timeit
t_eager    = timeit.timeit(lambda: model(x),          number=1000) / 1000
t_compiled = timeit.timeit(lambda: compiled_model(x), number=1000) / 1000
print(f"Eager:    {t_eager*1000:.2f}ms")
print(f"Compiled: {t_compiled*1000:.2f}ms")
print(f"Speedup:  {t_eager/t_compiled:.1f}x")