PyTorch Cheatsheet: Tensors, Autograd and nn.Module Reference

# Install — the build matrix is brittle, grab the exact command from pytorch.org
# CUDA 12.x (Linux + Windows)
pip install --index-url https://download.pytorch.org/whl/cu124 \
  torch torchvision torchaudio

# CPU-only (any OS)
pip install --index-url https://download.pytorch.org/whl/cpu \
  torch torchvision torchaudio

# Apple Silicon — built into the regular wheel; uses MPS automatically
pip install torch torchvision torchaudio

# Conda (CPU)
conda install -c pytorch pytorch torchvision torchaudio cpuonly

# Confirm the accelerator the runtime sees
python -c "
import torch
print('torch', torch.__version__)
print('cuda', torch.cuda.is_available(), 'mps', torch.backends.mps.is_available())
print('device count', torch.cuda.device_count())
"

# Pick the device once, reuse everywhere
# device = torch.device('cuda' if torch.cuda.is_available()
#                       else 'mps' if torch.backends.mps.is_available()
#                       else 'cpu')

import torch
import torch.nn as nn
import torch.nn.functional as F

class MLP(nn.Module):
    def __init__(self, in_dim: int, hidden: int, out_dim: int, p: float = 0.1):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(in_dim, hidden),
            nn.GELU(),
            nn.Dropout(p),
            nn.Linear(hidden, hidden),
            nn.GELU(),
            nn.Dropout(p),
            nn.Linear(hidden, out_dim),
        )

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.net(x)

model = MLP(in_dim=28*28, hidden=256, out_dim=10).to("cuda")

# Inspect parameter count + memory footprint
n_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
print(f"trainable params: {n_params/1e6:.2f}M")

# Initialize weights — register a small init function
def init_weights(m):
    if isinstance(m, nn.Linear):
        nn.init.kaiming_normal_(m.weight, nonlinearity="relu")
        if m.bias is not None: nn.init.zeros_(m.bias)
model.apply(init_weights)

import torch
import torch.nn as nn
from torch.amp import autocast, GradScaler

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

optim = torch.optim.AdamW(model.parameters(), lr=3e-4, weight_decay=0.01)
sched = torch.optim.lr_scheduler.CosineAnnealingLR(optim, T_max=epochs)
loss_fn = nn.CrossEntropyLoss()
scaler  = GradScaler()                        # mixed-precision loss scaling

for epoch in range(epochs):
    model.train()
    for xb, yb in train_loader:
        xb = xb.to(device, non_blocking=True)
        yb = yb.to(device, non_blocking=True)

        optim.zero_grad(set_to_none=True)     # cheaper than zero_grad()
        with autocast(device_type=device, dtype=torch.float16):
            logits = model(xb)
            loss   = loss_fn(logits, yb)

        scaler.scale(loss).backward()
        scaler.unscale_(optim)                # unscale before clipping
        torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
        scaler.step(optim)
        scaler.update()

    sched.step()

    # Eval pass
    model.eval()
    correct = total = 0
    with torch.inference_mode():              # faster than no_grad
        for xb, yb in val_loader:
            xb, yb = xb.to(device), yb.to(device)
            correct += (model(xb).argmax(1) == yb).sum().item()
            total   += yb.size(0)
    print(f"epoch {epoch} val_acc={correct/total:.4f}")

# MNIST classifier — Dataset, DataLoader, nn.Module, AMP, save/load, inference.
import torch
import torch.nn as nn
from torch.utils.data import DataLoader
from torchvision import datasets, transforms
from torch.amp import autocast, GradScaler

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

tf = transforms.Compose([transforms.ToTensor(),
                         transforms.Normalize((0.1307,), (0.3081,))])
train_ds = datasets.MNIST("./data", train=True,  download=True, transform=tf)
val_ds   = datasets.MNIST("./data", train=False, download=True, transform=tf)

train_loader = DataLoader(train_ds, batch_size=128, shuffle=True,
                          num_workers=2, pin_memory=True)
val_loader   = DataLoader(val_ds,   batch_size=256, num_workers=2, pin_memory=True)

model = nn.Sequential(nn.Flatten(), nn.Linear(784, 256), nn.GELU(),
                      nn.Linear(256, 10)).to(device)
opt    = torch.optim.AdamW(model.parameters(), lr=3e-4)
scaler = GradScaler()
loss_fn = nn.CrossEntropyLoss()

for epoch in range(3):
    model.train()
    for xb, yb in train_loader:
        xb, yb = xb.to(device), yb.to(device)
        opt.zero_grad(set_to_none=True)
        with autocast(device_type=device, dtype=torch.float16):
            loss = loss_fn(model(xb), yb)
        scaler.scale(loss).backward(); scaler.step(opt); scaler.update()

# Save + reload
torch.save(model.state_dict(), "mnist.pt")
model.load_state_dict(torch.load("mnist.pt", map_location=device))

# Inference
model.eval()
with torch.inference_mode():
    sample, _ = val_ds[0]
    pred = model(sample.unsqueeze(0).to(device)).argmax(1).item()
print("pred:", pred)