Cheatsheet

import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.data import DataLoader

# ===== Device & seeds =====
DEVICE = 'cuda' if torch.cuda.is_available() else 'mps' if torch.backends.mps.is_available() else 'cpu'
torch.manual_seed(42)
if DEVICE == 'cuda':
    torch.backends.cudnn.benchmark = True   # faster on fixed input shapes
    torch.set_float32_matmul_precision('high')

# ===== Model + DataLoader skeleton =====
class MLP(nn.Module):
    def __init__(self):
        super().__init__()
        self.net = nn.Sequential(nn.Linear(32, 128), nn.ReLU(),
                                  nn.Linear(128, 10))
    def forward(self, x): return self.net(x)

model = MLP().to(DEVICE)
opt   = torch.optim.AdamW(model.parameters(), lr=3e-4, weight_decay=1e-2)
sched = torch.optim.lr_scheduler.CosineAnnealingLR(opt, T_max=10)
loss_fn = nn.CrossEntropyLoss()
scaler = torch.amp.GradScaler('cuda', enabled=DEVICE == 'cuda')

# ===== Train step with mixed precision =====
def train_step(xb, yb):
    opt.zero_grad(set_to_none=True)
    with torch.amp.autocast(device_type=DEVICE.split(':')[0],
                            dtype=torch.float16, enabled=DEVICE == 'cuda'):
        logits = model(xb)
        loss   = loss_fn(logits, yb)
    scaler.scale(loss).backward()
    # Optional gradient clipping AFTER scaler.unscale_
    scaler.unscale_(opt)
    torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
    scaler.step(opt)
    scaler.update()
    return loss.item()

# ===== DataLoader defaults — the right ones =====
def make_loader(ds, batch=128, train=True):
    return DataLoader(ds, batch_size=batch,
        shuffle=train,
        num_workers=4,
        pin_memory=DEVICE == 'cuda',
        persistent_workers=True,
        drop_last=train,                # avoids the BN-on-batch-of-1 bug
    )

# ===== Inference defaults =====
@torch.inference_mode()
def predict(model, loader):
    model.eval()
    outs = []
    for xb in loader:
        xb = xb.to(DEVICE, non_blocking=True)
        outs.append(model(xb).argmax(dim=-1).cpu())
    return torch.cat(outs)

# ===== Gradient accumulation =====
ACCUM = 4
def accumulated_train(epoch_loader):
    opt.zero_grad(set_to_none=True)
    for i, (xb, yb) in enumerate(epoch_loader):
        xb, yb = xb.to(DEVICE, non_blocking=True), yb.to(DEVICE, non_blocking=True)
        with torch.amp.autocast(device_type='cuda', dtype=torch.float16, enabled=DEVICE == 'cuda'):
            loss = loss_fn(model(xb), yb) / ACCUM
        scaler.scale(loss).backward()
        if (i + 1) % ACCUM == 0:
            scaler.unscale_(opt)
            torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
            scaler.step(opt)
            scaler.update()
            opt.zero_grad(set_to_none=True)

# ===== Checkpointing — everything you need to resume =====
def save(path, epoch, best_val):
    torch.save({
        'epoch': epoch, 'best_val': best_val,
        'model': model.state_dict(),
        'opt':   opt.state_dict(),
        'sched': sched.state_dict(),
        'scaler': scaler.state_dict(),
    }, path)

def load(path):
    ck = torch.load(path, map_location=DEVICE)
    model.load_state_dict(ck['model'])
    opt.load_state_dict(ck['opt'])
    sched.load_state_dict(ck['sched'])
    scaler.load_state_dict(ck['scaler'])
    return ck['epoch'], ck['best_val']

# ===== torch.compile — Python -> graph, often 1.3-2x faster (PyTorch 2+) =====
model = torch.compile(model, mode='max-autotune')

# ===== Scheduler patterns =====
#   warmup + cosine     for transformers / big batch training
#   step                for image classification on top of a frozen backbone
#   one-cycle           for fastai-style 'find LR, ride it' training
#   reduce-on-plateau   for any val-driven schedule

# ===== Pitfalls =====
# - calling .item() / .cpu() inside the train loop kills throughput
# - drop_last=False + BatchNorm + batch-of-1 = NaN loss
# - model.eval() forgotten in inference -> dropout still active
# - opt.zero_grad() without set_to_none=True wastes memory bandwidth
# - missing pin_memory + non_blocking transfer -> CPU<->GPU bottleneck
# - tensors built on CPU then moved to GPU per step instead of building on GPU

# tensor zeros randn | nn.Module forward | optimizer zero_grad backward step | DataLoader