Summary

# ===== The three abstractions =====
# 1. Tensor  - n-dimensional array with autograd hooks
# 2. nn.Module - layer / model with .parameters() and .forward()
# 3. DataLoader - batched iterator with parallelism and shuffling

# ===== The training loop (the one you memorise) =====
import torch
from torch import nn
from torch.utils.data import DataLoader

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

model = MyNet().to(device)
optim = torch.optim.AdamW(model.parameters(), lr=3e-4)
loss_fn = nn.CrossEntropyLoss()
train_loader = DataLoader(train_ds, batch_size=64, shuffle=True, num_workers=4, pin_memory=True)
val_loader   = DataLoader(val_ds, batch_size=128, num_workers=4, pin_memory=True)

for epoch in range(10):
    model.train()
    for x, y in train_loader:
        x = x.to(device, non_blocking=True); y = y.to(device, non_blocking=True)
        optim.zero_grad(set_to_none=True)
        out = model(x)
        loss = loss_fn(out, y)
        loss.backward()
        optim.step()

    model.eval()
    correct = total = 0
    with torch.no_grad():
        for x, y in val_loader:
            x = x.to(device); y = y.to(device)
            pred = model(x).argmax(1)
            correct += (pred == y).sum().item(); total += y.numel()
    print(f'epoch {epoch} val_acc {correct/total:.4f}')

torch.save(model.state_dict(), 'model.pt')

# ===== GPU patterns =====
# - .to(device, non_blocking=True) with pin_memory=True DataLoader
# - autocast + GradScaler for mixed precision:
from torch.cuda.amp import autocast, GradScaler
scaler = GradScaler()
for x, y in train_loader:
    x, y = x.to(device, non_blocking=True), y.to(device, non_blocking=True)
    optim.zero_grad(set_to_none=True)
    with autocast(dtype=torch.float16):
        loss = loss_fn(model(x), y)
    scaler.scale(loss).backward()
    scaler.step(optim)
    scaler.update()

# ===== Reproducibility =====
import random, numpy as np
def seed_everything(s):
    random.seed(s); np.random.seed(s)
    torch.manual_seed(s); torch.cuda.manual_seed_all(s)
seed_everything(42)
torch.backends.cudnn.deterministic = True
torch.backends.cudnn.benchmark = False

# ===== When to reach for what =====
# - DataLoader with collate_fn for ragged inputs (text, graphs)
# - Lightning if you want the loop abstracted but keep PyTorch underneath
# - torch.compile(model) on PyTorch 2 for kernel fusion
# - DDP for multi-GPU; FSDP for very large models
# - ONNX export for cross-platform serving

# ===== Patterns to internalise =====
# - Train / Eval mode toggles dropout + batchnorm behaviour
# - zero_grad(set_to_none=True) is faster than zeros
# - amp + pin_memory + non_blocking is the perf trinity
# - Save state_dict, never the whole module (brittle across refactors)

# ===== Pitfalls =====
# - Forgetting model.eval() in val -> dropout pollutes metrics
# - Logging loss.item() in the hot path -> GPU/CPU sync every step
# - DataLoader with num_workers > 0 on Windows without if __name__ == '__main__'
# - Float16 underflow without GradScaler
# - Saving the whole model across refactors -> unpickling errors later

# Next: torch.compile, FSDP, FlashAttention, custom CUDA kernels, ExecuTorch.