Getting Started
Installation & Tensors
PyTorch tensors are the core data structure. Use to(device) to move tensors to GPU. torch.randn returns random numbers from a normal distribution. Pick the right wheel URL for your CUDA version; CPU-only builds are smaller for development.
# install PyTorch (CUDA 11.8)
pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu118
import torch
# create a tensor
x = torch.tensor([[1, 2], [3, 4]])
y = torch.zeros(3, 3)
z = torch.randn(2, 3) # standard normal
# tensor on GPU
device = "cuda" if torch.cuda.is_available() else "cpu"
x = x.to(device)Tensor Attributes
Every tensor has shape, dtype, and device attributes. Common dtypes: float32 (default), float16, int64, bool. numel() returns the total number of elements. Choose float16 for mixed precision to halve memory and speed up GPU compute.
x = torch.randn(3, 4, 5)
print(x.shape) # torch.Size([3, 4, 5])
print(x.dtype) # torch.float32
print(x.device) # cpu or cuda:0
print(x.requires_grad) # False
# specify dtype at creation
x = torch.zeros(3, dtype=torch.int64)
x = torch.tensor([1.0, 2.0], dtype=torch.float16)
# check number of elements
print(x.numel()) # 3Tensor Creation
torch.zeros/ones/full create constant tensors. arange and linspace create range tensors. rand is uniform, randn is normal. randint excludes the high bound. Use torch.manual_seed(42) before random ops for reproducibility.
# from Python lists
a = torch.tensor([1, 2, 3])
# special tensors
zeros = torch.zeros(2, 3)
ones = torch.ones(2, 3)
eye = torch.eye(3) # identity matrix
full = torch.full((2, 3), 7.0)
# ranges
arange = torch.arange(0, 10, 2) # [0, 2, 4, 6, 8]
linspace = torch.linspace(0, 1, 5) # [0, 0.25, 0.5, 0.75, 1.0]
# random
rand = torch.rand(2, 3) # uniform [0, 1)
randn = torch.randn(2, 3) # normal N(0, 1)
randint = torch.randint(0, 10, (3, 3))Type Conversion
to(dtype) converts tensor type. Shortcuts: .float(), .long(), .int(), .bool(). numpy() shares memory with the tensor (zero-copy on CPU). item() extracts a single Python scalar from a 0-d tensor. GPU tensors must be moved to CPU before numpy().
x = torch.tensor([1.5, 2.5, 3.5])
# change dtype
y = x.to(torch.int32) # tensor([1, 2, 3])
y = x.int() # shortcut
y = x.float()
y = x.bool()
y = x.long() # int64
# to numpy and back
arr = x.numpy() # shares memory (CPU)
t = torch.from_numpy(arr)
# 0-d tensor to Python scalar
val = x[0].item() # 1.5Indexing & Slicing
PyTorch indexing mirrors NumPy. Integer indexing returns a 0-d tensor (use .item() for a Python scalar). Boolean masks flatten the result. Indexing with a LongTensor selects along the first dimension. Advanced indexing creates copies; basic slicing creates views.
x = torch.arange(12).reshape(3, 4)
# tensor([[ 0, 1, 2, 3],
# [ 4, 5, 6, 7],
# [ 8, 9, 10, 11]])
x[0] # first row
x[1, 2] # element at (1, 2) -> 6
x[:, 1] # second column
x[0:2] # first two rows
x[:, ::2] # every other column
# boolean mask
mask = x > 5
x[mask] # tensor([6, 7, 8, 9, 10, 11])
# index with LongTensor
idx = torch.tensor([0, 2])
x[idx] # rows 0 and 2Reshaping & Views
view requires contiguous memory; reshape works always (copies if needed). unsqueeze/squeeze add/remove size-1 dimensions. permute reorders all dimensions; transpose swaps two. reshape(-1) flattens. Operations like permute may make tensors non-contiguous, requiring .contiguous() before view().
x = torch.arange(12)
# reshape (returns view when possible)
y = x.reshape(3, 4)
y = x.view(3, 4) # only for contiguous tensors
# add/remove dimensions
y = x.unsqueeze(0) # shape (1, 12)
y = x.squeeze() # remove size-1 dims
# transpose / permute
a = torch.randn(3, 4, 5)
b = a.permute(2, 0, 1) # shape (5, 3, 4)
b = a.transpose(0, 1) # swap dims 0 and 1
# flatten
flat = a.flatten() # shape (60,)
flat = a.reshape(-1) # equivalentTensor Operations
Arithmetic Operations
Arithmetic operators (+, -, *, /) are element-wise. In-place operations end with _ and save memory but invalidate autograd history — avoid them on tensors that require grad. Broadcasting follows NumPy rules: trailing dimensions are matched or stretched to 1.
a = torch.tensor([1.0, 2.0, 3.0])
b = torch.tensor([4.0, 5.0, 6.0])
# element-wise ops
a + b # tensor([5., 7., 9.])
a * b # tensor([ 4., 10., 18.])
a / b # element-wise division
a ** 2 # tensor([1., 4., 9.])
# in-place (suffix _)
a.add_(1) # a = a + 1, modifies in place
a.mul_(2) # a = a * 2
# scalar ops
a + 10
b * 0.5Element-wise Math
Element-wise math mirrors torch.* and tensor.* methods. clamp is the PyTorch name for clipping. Use torch.finfo(x.dtype).eps instead of a hardcoded 1e-8 for a numerically safe log. Activation functions are also available as nn modules (nn.ReLU()) for models.
x = torch.tensor([-1.0, 0.0, 1.0, 2.0])
torch.abs(x) # absolute value
torch.sqrt(torch.abs(x))
torch.exp(x) # e^x
torch.log(x.abs() + 1e-8)
torch.sin(x)
torch.clamp(x, -0.5, 1.5) # clip values
# activation functions
torch.sigmoid(x) # 1 / (1 + e^-x)
torch.tanh(x)
torch.relu(x) # max(0, x)
# rounding
y = torch.tensor([1.4, 1.5, 2.6])
torch.round(y) # tensor([1., 2., 3.])
torch.floor(y)
torch.ceil(y)Reduction Operations
Reductions collapse one or more dimensions. Pass dim to choose the axis; keepdim=True keeps the reduced axis as size 1 (critical for broadcasting). max/min with a dim return a named tuple (values, indices). item() pulls a 0-d tensor to a Python scalar.
x = torch.arange(12).reshape(3, 4).float()
x.sum() # 66.0 (scalar tensor)
x.sum(dim=0) # sum over rows -> shape (4,)
x.sum(dim=1) # sum over columns -> shape (3,)
x.mean(dim=1)
x.max() # returns value tensor
x.max(dim=1) # returns (values, indices)
x.argmin(dim=1) # indices of min along dim
# keep dimensions
x.sum(dim=1, keepdim=True) # shape (3, 1)
# norm
x.norm() # Frobenius norm
x.norm(dim=1) # per-row normComparison Operations
Comparison operators return bool tensors. torch.equal checks exact equality and returns a Python bool. torch.where(cond, x, y) is the vectorized ternary — perfect for masking or selecting. topk returns both values and indices, useful for top-1 / top-5 accuracy.
a = torch.tensor([1, 2, 3, 4])
b = torch.tensor([3, 2, 1, 4])
a == b # tensor([False, True, False, True])
a > b # tensor([False, False, True, False])
a != b
torch.equal(a, b) # False (single bool)
# where: pick from two tensors
torch.where(a > b, a, b) # element-wise max
# boolean helpers
(a > 2).any() # True
(a > 2).all() # False
(a == b).sum() # 2 (count of True)
# top-k
torch.topk(a, 2) # returns top 2 values + indicesMatrix Operations
@ and torch.matmul do matrix multiplication and broadcast over batch dims. Use torch.bmm for strict 3D batched matmul (no broadcasting). einsum is expressive for multi-tensor contractions but can be slower than matmul on some hardware. Avoid inverting matrices when you can solve with torch.solve / linalg.solve instead.
a = torch.randn(2, 3)
b = torch.randn(3, 4)
# matrix multiplication
c = a @ b # shape (2, 4)
c = torch.matmul(a, b) # equivalent
# batched matmul
A = torch.randn(10, 2, 3)
B = torch.randn(10, 3, 4)
C = A @ B # shape (10, 2, 4)
# element-wise product
e = a * a # shape (2, 3)
# other ops
torch.inverse(a[:, :3]) # inverse of square slice
torch.det(a[:, :3])
torch.svd(a) # singular value decomposition
torch.einsum('bi,i->b', torch.randn(5,3), torch.randn(3)) # batched dotConcatenation & Stacking
cat joins tensors along an existing axis; stack adds a new axis. split divides by chunk size, chunk divides by number of pieces. expand creates a view with broadcasted data (no memory copy) — use it instead of repeat when possible. repeat physically duplicates data.
a = torch.randn(2, 3)
b = torch.randn(2, 3)
# concat along existing dim
c = torch.cat([a, b], dim=0) # shape (4, 3)
c = torch.cat([a, b], dim=1) # shape (2, 6)
# stack along new dim
s = torch.stack([a, b], dim=0) # shape (2, 2, 3)
# split / chunk
x = torch.arange(6).reshape(2, 3)
parts = torch.split(x, 2, dim=0) # split into chunks of size 2
chunks = torch.chunk(x, 3, dim=1) # split into 3 chunks
# repeat / expand
r = a.repeat(2, 3) # repeat whole tensor
e = a.unsqueeze(0).expand(4, 2, 3) # broadcast without copyAutograd
Basic Autograd
Set requires_grad=True to track operations for autograd. Calling .backward() on a scalar tensor computes gradients of that tensor w.r.t. all leaf tensors with requires_grad=True. Gradients accumulate into .grad — call optimizer.zero_grad() or x.grad=None between iterations.
import torch
# tensors that require gradient tracking
x = torch.tensor([2.0], requires_grad=True)
y = torch.tensor([3.0], requires_grad=True)
# build a computation graph
z = x * y + x ** 2 # z = xy + x^2
print(z.grad_fn) # <AddBackward0 object>
# backprop
z.backward()
# dz/dx = y + 2x = 3 + 4 = 7
print(x.grad) # tensor([7.])
# dz/dy = x = 2
print(y.grad) # tensor([2.])Computing Gradients
backward() only works on scalar outputs; for tensors, sum or mean first. torch.autograd.grad returns gradients as a tuple without storing them in .grad. This is preferred inside libraries or when you need gradients of non-leaf tensors.
import torch
x = torch.linspace(-3, 3, steps=10, requires_grad=True)
y = torch.sin(x)
# sum needed for backward on non-scalar
y.sum().backward()
print(x.grad) # cos(x)
# gradient w.r.t. intermediate tensor
a = torch.tensor(1.0, requires_grad=True)
b = a * 2
c = b ** 2
grads = torch.autograd.grad(c, a)
print(grads) # (tensor(8.),) dc/da = 4b * 2 = 8Disabling Gradient Tracking
Use torch.no_grad() for inference to skip graph construction — saves memory and time. .detach() returns a tensor that shares data but is cut from the graph; great for logging or feeding tensors to non-PyTorch code. model.eval() switches layers (dropout, BN) to inference mode but does NOT disable autograd — combine with torch.no_grad().
x = torch.tensor([1.0], requires_grad=True)
# option 1: context manager (preferred)
with torch.no_grad():
y = x * 2 # y.requires_grad == False
# option 2: decorator
@torch.no_grad()
def inference(x):
return model(x)
# option 3: detach from graph
y = (x * 2).detach() # y is a new tensor, no grad
# enable grad inside no_grad
with torch.no_grad():
with torch.enable_grad():
z = x * 3 # z requires gradHigher-Order Gradients
create_graph=True keeps the graph of the gradient itself so you can differentiate again. Useful for gradient penalties (WGAN-GP), meta-learning, and physics-informed losses. It costs extra memory — free intermediate graphs with retain_graph=False when done.
x = torch.tensor(0.5, requires_grad=True)
# first derivative
y = torch.sin(x)
y.backward(create_graph=True)
print(x.grad) # cos(0.5)
# second derivative (gradient of gradient)
x.grad.zero_()
g = torch.autograd.grad(torch.sin(x), x, create_graph=True)[0]
g2 = torch.autograd.grad(g, x)[0]
print(g2) # -sin(0.5)
# practical use: penalize gradient (e.g. gradient penalty)
loss = g.norm()Gradient Hooks
register_hook on a tensor lets you inspect or modify its gradient during backward — handy for debugging NaNs or implementing gradient clipping on specific tensors. register_full_backward_hook on a module sees both input and output gradients. Always return a modified gradient if you want changes to propagate.
x = torch.tensor([1.0, 2.0], requires_grad=True)
y = (x ** 2).sum()
# full backward hook on a tensor
def print_grad(grad):
print("grad:", grad)
return grad * 2 # can modify gradient
y.register_hook(print_grad)
y.backward()
print(x.grad) # 2x but doubled by hook
# module hook
def hook_fn(module, grad_input, grad_output):
print(module.__class__.__name__, grad_output)
model.layer.register_full_backward_hook(hook_fn)Computational Graph Details
Leaf tensors are user-created (requires_grad=True, no grad_fn). Intermediate results are non-leaf. By default the graph is freed after backward(); pass retain_graph=True to call backward multiple times. create_graph=True retains the gradient graph for higher-order derivatives. Memory grows quickly — clear with x.grad=None between iterations.
x = torch.tensor(2.0, requires_grad=True)
w = torch.tensor(3.0, requires_grad=True)
b = torch.tensor(1.0, requires_grad=True)
y = w * x + b
print(y.is_leaf) # False
print(x.is_leaf) # True
# retain graph for multiple backward passes
loss = (y - 5) ** 2
loss.backward(retain_graph=True)
# can call backward again
loss.backward()
# double backward
loss.backward(create_graph=True)
print(x.grad)Data Loading
Dataset & DataLoader
DataLoader batches, shuffles, and parallel-loads a Dataset. num_workers>0 uses subprocesses — set 2-8 per GPU. pin_memory=True speeds up CPU→GPU transfer. For small datasets that fit in RAM, set num_workers=0 to avoid overhead.
from torch.utils.data import Dataset, DataLoader
# built-in datasets
from torchvision import datasets
mnist = datasets.MNIST("./data", train=True, download=True)
# wrap in DataLoader
loader = DataLoader(
mnist,
batch_size=64,
shuffle=True,
num_workers=4,
pin_memory=True,
)
for x, y in loader:
print(x.shape, y.shape) # (64, 1, 28, 28), (64,)
breakBatch Sampling
Pass a batch_sampler to control exactly which indices form each batch. WeightedRandomSampler rebalances imbalanced classes by oversampling rare ones. drop_last=True drops the final short batch — important for BatchNorm stability. Samplers iterate over indices; the DataLoader fetches the actual data.
from torch.utils.data import DataLoader, BatchSampler, RandomSampler
dataset = list(range(100))
sampler = RandomSampler(dataset)
batch_sampler = BatchSampler(sampler, batch_size=10, drop_last=True)
loader = DataLoader(dataset, batch_sampler=batch_sampler)
# custom sampler
class EvenSampler:
def __init__(self, data): self.data = data
def __iter__(self): return iter(range(0, len(self.data), 2))
def __len__(self): return len(self.data) // 2
# weighted sampling for imbalanced classes
from torch.utils.data import WeightedRandomSampler
weights = [0.1] * 50 + [1.0] * 50
sampler = WeightedRandomSampler(weights, num_samples=100)Shuffling & Sampling
shuffle=True randomizes order each epoch (good for training); keep validation shuffle=False for reproducible metrics. random_split is convenient but its RNG is not seedable directly — use Subset with torch.randperm and a Generator for full control. Subset is just a thin index view, no data copy.
from torch.utils.data import DataLoader, random_split
# train/val split
train, val = random_split(dataset, [80, 20])
train_loader = DataLoader(train, batch_size=16, shuffle=True)
val_loader = DataLoader(val, batch_size=32, shuffle=False)
# reproducible split
from torch.utils.data import Subset
import numpy as np
gen = torch.Generator().manual_seed(42)
idx = torch.randperm(len(dataset), generator=gen).tolist()
train_idx, val_idx = idx[:80], idx[80:]
train, val = Subset(dataset, train_idx), Subset(dataset, val_idx)Custom Collate
collate_fn turns a list of samples into a batch. The default stacks tensors along a new dim 0. Override it for variable-length sequences (pad then sort by length for pack_padded_sequence), dict-based samples, or mixed tensor/non-tensor data. Keep the function cheap — it runs on the data-loading workers.
from torch.utils.data import default_collate
# variable-length sequences: pad in collate
def pad_collate(batch):
# batch is list of (tensor, label)
seqs, labels = zip(*batch)
lens = torch.tensor([len(s) for s in seqs])
max_len = lens.max()
padded = torch.zeros(len(seqs), max_len)
for i, s in enumerate(seqs):
padded[i, :len(s)] = s
return padded, lens, torch.tensor(labels)
loader = DataLoader(dataset, batch_size=4, collate_fn=pad_collate)
# stack batch as dict
def dict_collate(batch):
return default_collate(batch)Parallel Loading
num_workers parallelizes data loading so the GPU never waits. persistent_workers=True avoids the per-epoch worker startup cost. prefetch_factor (default 2) controls how many batches each worker prepares ahead. On Windows and Jupyter, multiprocessing can be flaky — fall back to num_workers=0 to debug. Watch out: each worker inherits the full process state, including open file handles and CUDA contexts.
loader = DataLoader(
dataset,
batch_size=128,
num_workers=8, # one process per worker
pin_memory=True, # page-locked memory for fast H2D copy
persistent_workers=True, # keep workers alive across epochs
prefetch_factor=4, # batches prefetched per worker
)
# avoid CPU bottleneck
import os
os.cpu_count() # check available cores
# in Jupyter / Windows: set start method
import torch.multiprocessing as mp
# mp.set_start_method('spawn') # uncomment if needednn.Module
Defining a Module
Subclass nn.Module and implement forward(). Declare sub-modules and parameters as attributes so they are registered automatically (this is what enables .to(device), .parameters(), and state_dict). Always call super().__init__(). Avoid arbitrary Python logic in forward — keep it differentiable and tensor-only.
import torch.nn as nn
class MLP(nn.Module):
def __init__(self, in_dim, hidden, out_dim):
super().__init__()
self.fc1 = nn.Linear(in_dim, hidden)
self.fc2 = nn.Linear(hidden, hidden)
self.fc3 = nn.Linear(hidden, out_dim)
self.act = nn.ReLU()
def forward(self, x):
x = self.act(self.fc1(x))
x = self.act(self.fc2(x))
return self.fc3(x) # logits
model = MLP(784, 128, 10)
print(model) # prints structureLinear & Activation Layers
nn.Linear applies y = xW^T + b. Weight shape is (out, in). Prefer activation modules when you want them registered (state_dict, hooks); use functional F.* for one-off ops. GELU and SwiGLU-style activations dominate transformer architectures. Softmax should usually be omitted from the model and handled by CrossEntropyLoss.
import torch.nn as nn
linear = nn.Linear(in_features=20, out_features=10, bias=True)
print(linear.weight.shape) # (10, 20)
print(linear.bias.shape) # (10,)
# common activations
nn.ReLU() # max(0, x)
nn.LeakyReLU(0.01)
nn.GELU() # smooth, used in transformers
nn.Sigmoid() # (0, 1)
nn.Tanh() # (-1, 1)
nn.Softmax(dim=1)
nn.ELU()
# functional API (no parameters)
import torch.nn.functional as F
out = F.relu(linear(x)) # same as nn.ReLU()(linear(x))Parameters & Submodules
named_parameters() yields (name, tensor) for the whole tree — great for freezing, weight decay, or debugging. .modules() recurses into children; .children() is one level only. Buffers (e.g. BatchNorm running stats) show up in state_dict but not parameters(); use named_buffers() to see them.
model = MLP(784, 128, 10)
# iterate parameters
for name, p in model.named_parameters():
print(name, p.shape, p.requires_grad)
# only trainable params
trainable = filter(lambda p: p.requires_grad, model.parameters())
# count parameters
n_params = sum(p.numel() for p in model.parameters())
print(f"{n_params:,} parameters")
# access submodules by name
model.fc1 # the Linear layer
model['fc1'] # equivalent if using nn.ModuleDict
# children vs modules
list(model.children()) # direct children
list(model.modules()) # recursively all sub-modulesSequential Models
nn.Sequential chains modules where each output feeds the next input — perfect for linear pipelines. Use OrderedDict to name stages so you can index by attribute (model.fc1). For models with branches, skip connections, or multiple inputs/outputs, subclass nn.Module and write forward() explicitly.
import torch.nn as nn
# simple stack
model = nn.Sequential(
nn.Linear(784, 256),
nn.ReLU(),
nn.Dropout(0.5),
nn.Linear(256, 10),
)
# named layers via OrderedDict
from collections import OrderedDict
model = nn.Sequential(OrderedDict([
('fc1', nn.Linear(784, 256)),
('relu', nn.ReLU()),
('dropout', nn.Dropout(0.5)),
('fc2', nn.Linear(256, 10)),
]))
print(model.fc1) # access by name
# append at runtime
model.append(nn.Softmax(dim=1))Module Methods
to(device) moves parameters and buffers; .double()/.half()/.float() cast dtypes. train()/eval() toggle the self.training flag — Dropout and BatchNorm branch on it. state_dict() returns an OrderedDict of all parameters and buffers; load_state_dict restores them. .apply(fn) recurses into all submodules, useful for custom initialization.
model = MLP(784, 128, 10)
# move to device / dtype
model = model.to('cuda')
model = model.to(torch.float16)
# train vs eval mode
model.train() # enable dropout, update BN stats
model.eval() # disable dropout, freeze BN stats
# state dict
sd = model.state_dict()
model.load_state_dict(sd)
# apply a function to all submodules
model.apply(lambda m: print(type(m).__name__))
# set requires_grad on all parameters
for p in model.parameters():
p.requires_grad_(False)Initialization
PyTorch uses Kaiming He init for Linear/Conv by default, which works for most cases. Override with torch.nn.init functions (suffix _ means in-place). Kaiming suits ReLU-family; Xavier/Glorot suits tanh/sigmoid. Initializing biases to zero is standard. .apply(fn) runs fn on every submodule — perfect for bulk initialization or weight reset.
import torch.nn as nn
import torch.nn.init as init
class Net(nn.Module):
def __init__(self):
super().__init__()
self.fc1 = nn.Linear(100, 50)
self.fc2 = nn.Linear(50, 10)
self._init_weights()
def _init_weights(self):
init.kaiming_normal_(self.fc1.weight, nonlinearity='relu')
init.zeros_(self.fc1.bias)
init.xavier_uniform_(self.fc2.weight)
init.zeros_(self.fc2.bias)
# apply to existing model
def init_all(m):
if isinstance(m, nn.Linear):
init.kaiming_normal_(m.weight)
if m.bias is not None:
init.zeros_(m.bias)
model.apply(init_all)Loss Functions
CrossEntropyLoss
CrossEntropyLoss applies log-softmax + NLLLoss internally — pass raw logits, never softmax outputs. Targets are integer class indices (not one-hot). Use weight for imbalanced classes and ignore_index=-100 to mask padding tokens. label_smoothing=0.1 is a popular regularizer for classification and NLP.
import torch
import torch.nn as nn
# logits shape (N, C), targets shape (N,) with class indices
logits = torch.randn(4, 3) # 4 samples, 3 classes
targets = torch.tensor([0, 2, 1, 0])
criterion = nn.CrossEntropyLoss()
loss = criterion(logits, targets)
# equivalently, combine LogSoftmax + NLLLoss
log_softmax = nn.LogSoftmax(dim=1)
nll = nn.NLLLoss()
loss2 = nll(log_softmax(logits), targets)
# weight classes (handle imbalance)
weights = torch.tensor([1.0, 2.0, 1.0])
criterion = nn.CrossEntropyLoss(weight=weights)
# ignore padding index (NLP)
criterion = nn.CrossEntropyLoss(ignore_index=-100)MSELoss & L1Loss
MSELoss penalizes large errors quadratically (sensitive to outliers); L1 is linear and more robust. SmoothL1Loss behaves like L2 near zero and L1 far away — used in object detection (Fast R-CNN). reduction='none' returns a per-element tensor so you can mask or weight it yourself before calling .mean().
import torch.nn as nn
pred = torch.randn(4, 1)
target = torch.randn(4, 1)
# mean squared error (L2)
mse = nn.MSELoss()
loss = mse(pred, target) # mean over all elements
# L1 (mean absolute error)
l1 = nn.L1Loss()
loss = l1(pred, target)
# smooth L1 (Huber) — robust to outliers
smooth = nn.SmoothL1Loss(beta=0.1)
# reduction options
mse_sum = nn.MSELoss(reduction='sum')
mse_none = nn.MSELoss(reduction='none') # per-elementBCELoss & BCEWithLogitsLoss
Always prefer BCEWithLogitsLoss over BCELoss+sigmoid — it uses the log-sum-exp trick for numerical stability. pos_weight multiplies the positive term to counter class imbalance (like focal loss lite). For multi-label classification (each class independent), use BCEWithLogitsLoss with shape (N, C) and sigmoid outputs.
import torch.nn as nn
# binary classification, multi-label, or sigmoid outputs
logits = torch.randn(4, 1)
target = torch.tensor([[1.0], [0.0], [1.0], [0.0]])
# preferred: combines sigmoid + BCE, numerically stable
criterion = nn.BCEWithLogitsLoss()
loss = criterion(logits, target)
# BCELoss expects probabilities (apply sigmoid first)
criterion = nn.BCELoss()
probs = torch.sigmoid(logits)
loss = criterion(probs, target)
# positive weight for imbalanced binary tasks
pos_weight = torch.tensor([5.0])
criterion = nn.BCEWithLogitsLoss(pos_weight=pos_weight)Custom Loss Function
A custom loss is just a function (or nn.Module) that returns a scalar tensor — autograd handles the rest. Make every operation differentiable (torch.where with differentiable branches is fine). Module form is preferred when you need hyperparameters, serialization, or to register it in a config. Keep reduction explicit (mean/sum/none) so behavior is predictable.
import torch
import torch.nn as nn
import torch.nn.functional as F
# as a function
def huber_loss(pred, target, delta=1.0):
error = pred - target
abs_err = error.abs()
quad = torch.where(abs_err <= delta,
0.5 * error ** 2,
delta * (abs_err - 0.5 * delta))
return quad.mean()
# as a module (so it can have parameters / state)
class FocalLoss(nn.Module):
def __init__(self, alpha=0.25, gamma=2.0):
super().__init__()
self.alpha = alpha
self.gamma = gamma
def forward(self, logits, targets):
bce = F.binary_cross_entropy_with_logits(logits, targets, reduction='none')
p = torch.sigmoid(logits)
pt = p * targets + (1 - p) * (1 - targets)
loss = self.alpha * (1 - pt) ** self.gamma * bce
return loss.mean()
criterion = FocalLoss()Loss Reduction & Combining
reduction='none' gives per-sample losses so you can mask, reweight, or aggregate manually. Combine multi-task losses by weighted sum — the weights balance gradients, not just magnitudes (tune them). Each sub-loss must stay differentiable and on the same device/dtype. Detach a loss term with .detach() if you want to log it without backpropagating through it.
import torch.nn as nn
import torch
criterion = nn.CrossEntropyLoss(reduction='none')
loss_per = criterion(logits, targets) # shape (N,)
loss = loss_per.mean()
# class-wise mask
mask = targets != -100
loss = criterion(logits[mask], targets[mask])
# multi-task loss
loss_cls = nn.CrossEntropyLoss()(logits, labels)
loss_box = nn.L1Loss()(boxes_pred, boxes_gt)
total = 1.0 * loss_cls + 0.5 * loss_box
# gradient-weighted
total = loss_cls + 0.5 * loss_box
total.backward()NLLLoss & Other Losses
NLLLoss + LogSoftmax is mathematically identical to CrossEntropyLoss but lets you reuse the log-probs. TripletMarginLoss trains metric-learning embeddings. KLDivLoss expects log-probs as input and probs as target. CTCLoss aligns unsegmented sequences to targets — pass input_lengths and target_lengths to avoid runtime errors.
import torch.nn as nn
# NLLLoss expects log-probabilities (use LogSoftmax first)
log_probs = torch.log_softmax(logits, dim=1)
loss = nn.NLLLoss()(log_probs, targets)
# triplet loss for embeddings
anchor = torch.randn(4, 128)
positive = torch.randn(4, 128)
negative = torch.randn(4, 128)
loss = nn.TripletMarginLoss(margin=1.0)(anchor, positive, negative)
# KL divergence (distributions)
loss = nn.KLDivLoss(reduction='batchmean')(log_probs, target_dist)
# cosine embedding loss
loss = nn.CosineEmbeddingLoss()(x1, x2, torch.tensor([1, -1, 1, -1]))
# CTCLoss for sequence alignment (ASR/OCR)
loss = nn.CTCLoss(blank=0, zero_infinity=True)(log_probs, targets, in_lens, out_lens)Optimizers
SGD
SGD + momentum is still a top choice for large-batch CNN training and generalizes better than Adam on some tasks. Momentum 0.9 is standard; Nesterov gives a small extra edge. weight_decay adds L2 regularization. SGD needs careful LR scheduling and warmup; it is not auto-tuning like Adam.
import torch.optim as optim
# basic SGD
optimizer = optim.SGD(model.parameters(), lr=0.01)
# SGD with momentum (classic: 0.9)
optimizer = optim.SGD(model.parameters(), lr=0.01, momentum=0.9)
# Nesterov momentum
optimizer = optim.SGD(
model.parameters(),
lr=0.01,
momentum=0.9,
nesterov=True,
)
# weight decay (L2 regularization)
optimizer = optim.SGD(model.parameters(), lr=0.01, momentum=0.9, weight_decay=1e-4)Adam & AdamW
Adam stores per-parameter moments, so it adapts the step size automatically — great default for NLP and small datasets. AdamW decouples weight decay from the gradient update, which matters at scale (transformers). betas=(0.9, 0.999) is standard; some vision models use (0.9, 0.95). Use per-parameter groups for differential learning rates in transfer learning.
import torch.optim as optim
# Adam — adaptive, default-friendly
optimizer = optim.Adam(model.parameters(), lr=1e-3)
# AdamW — decoupled weight decay (preferred for transformers)
optimizer = optim.AdamW(
model.parameters(),
lr=3e-4,
betas=(0.9, 0.999),
weight_decay=0.01,
)
# per-parameter group settings
optimizer = optim.Adam([
{'params': model.encoder.parameters(), 'lr': 1e-4},
{'params': model.head.parameters(), 'lr': 1e-3},
], lr=1e-3)Optimizer Step Pattern
The 3-step pattern — zero_grad, backward, step — is the heart of PyTorch training. set_to_none=True (default in newer versions) is faster than filling with zeros because it skips a memset and lets the optimizer skip the param. Manually edit p.grad.data between backward() and step() for clipping or custom updates.
optimizer = optim.Adam(model.parameters(), lr=1e-3)
# standard training step
optimizer.zero_grad() # clear old grads
loss = criterion(model(x), y)
loss.backward() # compute gradients
optimizer.step() # update parameters
# set gradients to None (faster, recommended)
optimizer.zero_grad(set_to_none=True)
# manual gradient modification
for p in model.parameters():
if p.grad is not None:
p.grad.data.clamp_(-1, 1) # gradient clipping
optimizer.step()Gradient Clipping
clip_grad_norm_ rescales gradients if their global L2 norm exceeds max_norm — a stable default for RNNs and transformers. clip_grad_value_ hard-clamps each element (less common). Clipping prevents exploding gradients without changing direction. Log the pre-clip grad norm to detect instability early (sudden spikes > 100 indicate problems).
# gradient norm clipping (most common)
torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
# gradient value clipping
torch.nn.utils.clip_grad_value_(model.parameters(), clip_value=0.5)
# in a training loop
optimizer.zero_grad()
loss = criterion(model(x), y)
loss.backward()
torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
optimizer.step()
# inspect gradient norm
total_norm = torch.norm(torch.stack([
p.grad.norm(2) for p in model.parameters() if p.grad is not None
]), 2)Learning Rate & Schedules
Each param_group has its own lr — iterate to set all. Manual warmup linearly ramps LR for the first N steps to stabilize early training (critical for transformers without bias correction). Use a scheduler object (next section) for cosine decay, OneCycle, etc. Always step the scheduler after optimizer.step() unless the docs say otherwise.
import torch.optim as optim
# get / set learning rate
for g in optimizer.param_groups:
print(g['lr'])
g['lr'] = 1e-4
# manual warmup
for step in range(warmup_steps):
lr = base_lr * (step + 1) / warmup_steps
for g in optimizer.param_groups:
g['lr'] = lr
# optimizer also accepts schedulers (see lr-scheduler section)
from torch.optim.lr_scheduler import StepLR
scheduler = StepLR(optimizer, step_size=10, gamma=0.1)Other Optimizers
RMSprop is the classic RNN optimizer and still works well. Adagrad accumulates squared gradients and can stall early — fine for sparse features but not deep nets. LBFGS needs a closure and full-batch evaluation — great for small convex problems or fine-tuning, too slow for large deep nets. For most new projects, start with AdamW.
import torch.optim as optim
# RMSprop — good for RNNs
optim.RMSprop(model.parameters(), lr=1e-3, alpha=0.99)
# Adagrad — adapts per-parameter, accumulates squared grads
optim.Adagrad(model.parameters(), lr=1e-2)
# Adadelta — Adagrad variant with running average
optim.Adadelta(model.parameters(), rho=0.9)
# LBFGS — full-batch quasi-Newton, second-order
optim.LBFGS(model.parameters(), lr=1.0, max_iter=20)
# NAdam — Nesterov-accelerated Adam
optim.NAdam(model.parameters(), lr=1e-3)Training Loop
Basic Training Loop
The minimal training loop: zero_grad, forward, loss, backward, step. Move data to GPU inside the loop (async overlap with num_workers). loss.item() pulls a scalar to Python — do this for logging, not for the loss tensor itself (you would break the graph). Wrap the loop in a function so you can reuse it.
model = MLP(784, 128, 10).to(device)
optimizer = optim.Adam(model.parameters(), lr=1e-3)
criterion = nn.CrossEntropyLoss()
for epoch in range(epochs):
model.train()
for x, y in train_loader:
x, y = x.to(device), y.to(device)
optimizer.zero_grad()
logits = model(x)
loss = criterion(logits, y)
loss.backward()
optimizer.step()
print(f"epoch {epoch}, loss={loss.item():.4f}")Validation in Loop
model.eval() + torch.no_grad() are both required: eval switches dropout/BN, no_grad skips the graph. Accumulate weighted loss (multiply by batch size) so the average is correct even for the last short batch. argmax(dim=1) gives predicted class indices. Save the best model when val metric improves (see save-load section).
for epoch in range(epochs):
model.train()
for x, y in train_loader:
# ... training step ...
# validation
model.eval()
val_loss, correct, total = 0.0, 0, 0
with torch.no_grad():
for x, y in val_loader:
x, y = x.to(device), y.to(device)
logits = model(x)
val_loss += criterion(logits, y).item() * x.size(0)
correct += (logits.argmax(1) == y).sum().item()
total += x.size(0)
print(f"val_loss={val_loss/total:.4f}, acc={correct/total:.4f}")Gradient Accumulation
Gradient accumulation simulates a larger batch by summing gradients over several mini-batches before stepping. Divide the loss by accum_steps so the effective gradient matches a single large batch. Make sure the dataset length is divisible by accum_steps, or step on the leftover gradients at epoch end. Especially useful for transformer fine-tuning on a single GPU.
accum_steps = 4
optimizer.zero_grad()
for i, (x, y) in enumerate(train_loader):
x, y = x.to(device), y.to(device)
loss = criterion(model(x), y) / accum_steps
loss.backward()
if (i + 1) % accum_steps == 0:
optimizer.step()
optimizer.zero_grad()
# equivalent effective batch = accum_steps * batch_size
# useful when GPU memory limits batch sizeMixed Precision Training
Mixed precision runs forward ops in float16 (or bfloat16 on Ampere+) for 1.5-3x speedup and halved memory. GradScaler multiplies the loss to prevent float16 gradient underflow, then unscales before optimizer.step(). On H100/A100 prefer bfloat16 — it has the same range as float32 and does not need a scaler. The autocast context manager only affects the forward pass.
from torch.amp import autocast, GradScaler
scaler = GradScaler('cuda')
for x, y in train_loader:
x, y = x.to(device), y.to(device)
optimizer.zero_grad()
with autocast('cuda'):
logits = model(x)
loss = criterion(logits, y)
scaler.scale(loss).backward()
scaler.step(optimizer)
scaler.update()
# autocast casts select ops to float16/bf16 for speed
# GradScaler prevents underflow by scaling the lossProgress & Logging
tqdm shows a live progress bar — use set_postfix to display the running loss. An exponential moving average (running = 0.9*running + 0.1*loss) smooths noise better than the raw last-batch value. For serious runs, log to TensorBoard or Weights & Biases instead of stdout so you can compare runs. Avoid logging every step if it slows training.
from tqdm import tqdm
for epoch in range(epochs):
model.train()
pbar = tqdm(train_loader, desc=f"epoch {epoch}")
running = 0.0
for i, (x, y) in enumerate(pbar):
x, y = x.to(device), y.to(device)
optimizer.zero_grad()
loss = criterion(model(x), y)
loss.backward()
optimizer.step()
running = 0.9 * running + 0.1 * loss.item()
pbar.set_postfix(loss=f"{running:.4f}")
# TensorBoard logging (see visualization section)
writer.add_scalar('train/loss', loss.item(), global_step)Reproducibility
Full reproducibility needs seeds for Python, NumPy, PyTorch CPU and CUDA. cudnn.deterministic=True disables non-deterministic algorithms (some conv backward ops). DataLoader workers each need seeding via worker_init_fn because they fork with their own RNG state. Even then, multi-GPU and some cuDNN ops may not be bit-reproducible — aim for statistical reproducibility.
import torch, numpy as np, random
def set_seed(seed=42):
random.seed(seed)
np.random.seed(seed)
torch.manual_seed(seed)
torch.cuda.manual_seed_all(seed)
# cudnn determinism (may slow down)
torch.backends.cudnn.deterministic = True
torch.backends.cudnn.benchmark = False
set_seed(42)
# DataLoader worker seeding
def seed_worker(worker_id):
worker_seed = torch.initial_seed() % 2**32
np.random.seed(worker_seed)
random.seed(worker_seed)
g = torch.Generator().manual_seed(42)
loader = DataLoader(ds, batch_size=32, worker_init_fn=seed_worker, generator=g)Evaluation & Metrics
Eval Mode & No Grad
Always pair model.eval() with torch.no_grad() for evaluation. eval() flips dropout and BatchNorm to inference mode (BN uses running stats); no_grad() disables autograd to save memory and time. Multiply loss by batch size when accumulating so the final average is unbiased even if the last batch is short.
model.eval()
test_loss, correct, total = 0.0, 0, 0
with torch.no_grad():
for x, y in test_loader:
x, y = x.to(device), y.to(device)
logits = model(x)
test_loss += criterion(logits, y).item() * x.size(0)
preds = logits.argmax(dim=1)
correct += (preds == y).sum().item()
total += x.size(0)
print(f"test_loss={test_loss/total:.4f}, acc={correct/total:.4f}")Accuracy & Top-k
argmax(dim=1) returns predicted class per row. topk returns the k largest values and their indices — equating them to targets expanded to (N,1) and using .any(dim=1) gives top-k correctness. Per-class accuracy exposes imbalance problems that overall accuracy hides. Always move tensors to CPU and call .item() before passing to print or log.
logits = torch.randn(8, 10)
targets = torch.randint(0, 10, (8,))
# top-1 accuracy
preds = logits.argmax(dim=1)
acc = (preds == targets).float().mean().item()
# top-k accuracy
topk = 3
_, idx = logits.topk(topk, dim=1) # shape (N, k)
correct = idx.eq(targets.view(-1, 1)).any(dim=1)
topk_acc = correct.float().mean().item()
# per-class accuracy
for c in range(10):
mask = targets == c
print(c, (preds[mask] == targets[mask]).float().mean().item())Confusion Matrix
Accumulate predictions on GPU, then move to CPU once at the end to avoid per-batch syncs. confusion_matrix shows where the model confuses classes — perfect for debugging. classification_report gives precision, recall, and F1 per class. For very large datasets, stream predictions to disk instead of holding them in memory.
from sklearn.metrics import confusion_matrix, classification_report
import numpy as np
all_preds, all_labels = [], []
model.eval()
with torch.no_grad():
for x, y in test_loader:
logits = model(x.to(device))
all_preds.append(logits.argmax(1).cpu())
all_labels.append(y)
preds = torch.cat(all_preds).numpy()
labels = torch.cat(all_labels).numpy()
cm = confusion_matrix(labels, preds)
print(cm)
print(classification_report(labels, preds))Precision, Recall & F1
torchmetrics accumulates state across batches and computes the metric correctly even with class imbalance or changing batch sizes. average='macro' treats all classes equally; 'weighted' weights by support; 'micro' equals accuracy for single-label tasks. Call .reset() between epochs or evaluations to clear the accumulator.
from torchmetrics import Accuracy, Precision, Recall, F1Score
acc = Accuracy(task='multiclass', num_classes=10).to(device)
prec = Precision(task='multiclass', num_classes=10, average='macro').to(device)
rec = Recall(task='multiclass', num_classes=10, average='macro').to(device)
f1 = F1Score(task='multiclass', num_classes=10, average='macro').to(device)
model.eval()
with torch.no_grad():
for x, y in test_loader:
logits = model(x.to(device))
preds = logits.argmax(1)
acc.update(preds, y.to(device))
prec.update(preds, y.to(device))
print(acc.compute(), prec.compute(), f1.compute())Cross-Validation
K-fold cross-validation gives a robust estimate of model performance and its variance. Re-instantiate the model per fold so weights don't leak across folds. For large datasets, 5 folds is usually enough; for small ones, use stratified splits (StratifiedKFold) to preserve class balance. Report mean +/- std so readers see stability.
from sklearn.model_selection import KFold
kfold = KFold(n_splits=5, shuffle=True, random_state=42)
fold_accs = []
for fold, (train_idx, val_idx) in enumerate(kfold.split(dataset)):
train = torch.utils.data.Subset(dataset, train_idx)
val = torch.utils.data.Subset(dataset, val_idx)
model = build_model() # fresh model per fold
train_model(model, train) # your training function
acc = evaluate(model, val)
fold_accs.append(acc)
print(f"fold {fold}: acc={acc:.4f}")
print(f"mean={np.mean(fold_accs):.4f} +/- {np.std(fold_accs):.4f}")Model Comparison
Always compare models on the same held-out test set with the same seed and metric. Beyond accuracy, look at calibration, inference latency, and memory. Use bootstrap confidence intervals when test-set differences are small. Keep a results table (CSV or W&B) so you can compare against future experiments instead of just the last run.
results = {}
for name, model in models.items():
model.eval()
accs, losses = [], []
with torch.no_grad():
for x, y in test_loader:
x, y = x.to(device), y.to(device)
logits = model(x)
losses.append(criterion(logits, y).item())
accs.append((logits.argmax(1) == y).float().mean().item())
results[name] = {
'acc': np.mean(accs),
'loss': np.mean(losses),
}
for name, r in sorted(results.items(), key=lambda kv: -kv[1]['acc']):
print(f"{name:20s} acc={r['acc']:.4f} loss={r['loss']:.4f}")Save & Load
Save & Load State Dict
state_dict() is an OrderedDict of parameter tensors — saving it (not the whole model) is the recommended pattern because it is portable and decoupled from the exact class/file layout. map_location lets you load a CUDA checkpoint on CPU. strict=False loads only matching keys, handy when the architecture changed slightly.
import torch
# save (recommended: only state_dict)
torch.save(model.state_dict(), 'model.pth')
# load
model = MLP(784, 128, 10) # must match the saved architecture
model.load_state_dict(torch.load('model.pth', map_location='cpu'))
model.eval()
# load on GPU that was trained on GPU
model.load_state_dict(torch.load('model.pth'))
model.to(device)
# strict=False allows partial loading (transfer learning)
model.load_state_dict(torch.load('model.pth'), strict=False)Save Entire Model
Saving the whole model pickles the Python object, binding it to the exact class and file structure at save time. This is brittle — any refactor breaks loading. Use it only for quick experiments. For anything you might deploy or share, save state_dict and keep the architecture in code or config.
# save entire model (uses pickle) — NOT recommended
torch.save(model, 'model_full.pth')
# load
model = torch.load('model_full.pth', map_location='cpu')
model.eval()
# pickle limitations:
# - class definition must be importable at load time
# - breaks if you refactor or rename the class
# - not portable across PyTorch versions
# prefer state_dict for productionTraining Checkpoints
A full checkpoint stores everything needed to resume training exactly: model, optimizer, scheduler, epoch, and RNG states. Optimizer state (momentum buffers, Adam moments) is essential — resuming without it effectively reinitializes the optimizer and can spike the loss. Save RNG state only if you need bit-exact reproducibility.
checkpoint = {
'epoch': epoch,
'model_state': model.state_dict(),
'optimizer_state': optimizer.state_dict(),
'scheduler_state': scheduler.state_dict() if scheduler else None,
'loss': loss.item(),
'rng_state': torch.get_rng_state(),
'cuda_rng_state': torch.cuda.get_rng_state_all(),
}
torch.save(checkpoint, f'ckpt_epoch{epoch}.pth')
# resume
ckpt = torch.load('ckpt_epoch10.pth', map_location='cpu')
model.load_state_dict(ckpt['model_state'])
optimizer.load_state_dict(ckpt['optimizer_state'])
start_epoch = ckpt['epoch'] + 1
torch.set_rng_state(ckpt['rng_state'])Save Best Model
Save the best checkpoint by validation metric so you can recover the best model even if later epochs overfit. For most tasks keep only the best checkpoint to save disk; for research, keep per-epoch checkpoints so you can roll back. Combine with a CSV log of (epoch, train_loss, val_loss, val_acc, path) to know which checkpoint to pick.
best_acc = 0.0
for epoch in range(epochs):
train_one_epoch(model, ...)
acc = evaluate(model, val_loader)
if acc > best_acc:
best_acc = acc
torch.save(model.state_dict(), 'best.pth')
print(f"new best {acc:.4f} saved")
# load best at the end
model.load_state_dict(torch.load('best.pth'))
final_acc = evaluate(model, test_loader)Loading Partial Models
strict=False returns lists of missing and unexpected keys — inspect them to confirm only the expected layers were skipped. Renaming keys lets you adapt a checkpoint to a renamed architecture. After loading a backbone, freeze its weights for pure feature extraction or train with a much smaller LR than the new head.
# transfer learning: load a backbone, drop the classifier
pretrained = torch.load('resnet_backbone.pth')
model = ResNet(num_classes=100)
# drop mismatched keys (e.g. final fc)
missing, unexpected = model.load_state_dict(pretrained, strict=False)
print('missing:', missing) # e.g. ['fc.weight', 'fc.bias']
print('unexpected:', unexpected)
# rename keys before loading
pretrained = {k.replace('encoder.', 'backbone.'): v
for k, v in pretrained.items()}
# freeze loaded weights
for name, p in model.named_parameters():
if 'fc' not in name:
p.requires_grad_(False)Export to TorchScript
TorchScript produces a self-contained graph you can load in LibTorch (C++) or Python without the original class — ideal for deployment. trace captures one control-flow path; script handles data-dependent control flow but requires type annotations. ONNX exports for inference in ONNX Runtime, TensorRT, or OpenVINO. Use dynamic_axes for variable batch size.
# scripted (production-ready, runs without Python)
scripted = torch.jit.script(model)
scripted.save('model_scripted.pt')
# traced (records a static graph from sample input)
example = torch.randn(1, 3, 224, 224)
traced = torch.jit.trace(model, example)
traced.save('model_traced.pt')
# load TorchScript in C++ / Python
loaded = torch.jit.load('model_scripted.pt')
loaded.eval()
out = loaded(example)
# export to ONNX (cross-framework)
torch.onnx.export(model, example, 'model.onnx',
input_names=['input'],
output_names=['output'],
dynamic_axes={'input': {0: 'batch'}})CNN
Conv2d Layer
Conv2d weight shape is (out_channels, in_channels/groups, kH, kW). NCHW is the default memory layout. padding='same' (PyTorch 1.9+) keeps spatial size for stride 1. groups=in_channels makes it depthwise (used in MobileNet). Compute output size: (H + 2P - K) / S + 1.
import torch.nn as nn
conv = nn.Conv2d(
in_channels=3, # RGB input
out_channels=16, # number of filters
kernel_size=3, # 3x3 filter
stride=1,
padding=1, # 'same' for kernel 3, stride 1
bias=True,
)
print(conv.weight.shape) # (16, 3, 3, 3)
x = torch.randn(4, 3, 32, 32) # NCHW
out = conv(x)
print(out.shape) # (4, 16, 32, 32) with padding=1
# depthwise separable conv
dw = nn.Conv2d(16, 16, 3, padding=1, groups=16)Pooling Layers
Pooling downsamples spatial dimensions, adding invariance to small translations. MaxPool keeps the strongest response; AvgPool is smoother. AdaptiveAvgPool2d(1) collapses spatial dims to 1x1 — the standard 'global average pool' used to feed conv features into a Linear head. Prefer strided convs over pooling for learnable downsampling in modern architectures.
import torch.nn as nn
x = torch.randn(4, 16, 32, 32)
# max pooling
pool = nn.MaxPool2d(kernel_size=2, stride=2)
out = pool(x) # shape (4, 16, 16, 16)
# average pooling
avg = nn.AvgPool2d(2, stride=2)
# adaptive pool — output size fixed regardless of input
gap = nn.AdaptiveAvgPool2d(1) # global average pool
out = gap(x) # shape (4, 16, 1, 1)
# fractional max pool (slightly improves accuracy)
fmp = nn.FractionalMaxPool2d(2, output_ratio=0.5)CNN Architecture
A canonical CNN block is Conv -> BatchNorm -> ReLU -> Pool. AdaptiveAvgPool2d(1) makes the network input-size agnostic so you can train at one resolution and test at another. BatchNorm2d goes after the conv, before the activation. Keep conv kernels at 3x3 with padding=1 (VGG-style) — they are parameter-efficient and fuse well.
import torch.nn as nn
class TinyCNN(nn.Module):
def __init__(self, num_classes=10):
super().__init__()
self.features = nn.Sequential(
nn.Conv2d(3, 32, 3, padding=1),
nn.BatchNorm2d(32),
nn.ReLU(inplace=True),
nn.MaxPool2d(2), # 32 -> 16
nn.Conv2d(32, 64, 3, padding=1),
nn.BatchNorm2d(64),
nn.ReLU(inplace=True),
nn.MaxPool2d(2), # 16 -> 8
)
self.classifier = nn.Sequential(
nn.AdaptiveAvgPool2d(1),
nn.Flatten(),
nn.Linear(64, num_classes),
)
def forward(self, x):
return self.classifier(self.features(x))Feature Maps & Receptive Field
Feature maps show what each channel detects — early layers find edges and colors, deeper layers find textures and parts. register_forward_hook lets you capture activations without modifying forward. Receptive field grows with depth and pooling; roughly, output pixel (i,j) depends on a patch of input centered there with size = sum of strides and kernels up to that layer.
# visualize feature maps
model.eval()
x = images[0:1].to(device)
with torch.no_grad():
feat = model.features[0](x) # first conv output
print(feat.shape) # (1, 32, H, W)
import matplotlib.pyplot as plt
fig, axes = plt.subplots(4, 8, figsize=(12, 6))
for i, ax in enumerate(axes.flat):
if i < feat.shape[1]:
ax.imshow(feat[0, i].cpu(), cmap='viridis')
ax.axis('off')
# hook to capture intermediate activations
acts = {}
def hook(module, inp, out): acts['layer'] = out
model.features[3].register_forward_hook(hook)Batch Normalization
BatchNorm2d normalizes per channel over (N, H, W), keeping running EMA stats for inference. train()/eval() matters — eval uses running stats so output is deterministic. BatchNorm couples examples within a batch, which breaks with very small batches; switch to GroupNorm or LayerNorm then. SyncBatchNorm is needed for correct multi-GPU BN.
import torch.nn as nn
# 2D BN for conv outputs
bn = nn.BatchNorm2d(64)
bn.weight.shape # (64,) — learnable scale gamma
bn.bias.shape # (64,) — learnable shift beta
bn.running_mean.shape # (64,) — EMA of batch means
bn.running_var.shape # (64,) — EMA of batch vars
# 1D BN for linear / time-series
bn1d = nn.BatchNorm1d(128)
# in training mode: use batch stats, update running stats
model.train()
# in eval mode: use running stats (deterministic)
model.eval()
# alternatives
nn.LayerNorm(128) # normalize over features (transformers)
nn.GroupNorm(8, 64) # split channels into 8 groups
nn.InstanceNorm2d(64) # per-sample, style transferPopular CNN Architectures
torchvision.models exposes ResNet, EfficientNet, MobileNet, ViT, and more with pretrained ImageNet weights. Replacing the final fc layer is the standard transfer-learning entry point. For detection/segmentation, drop the head and use the conv features. Check the weights API (WeightsEnum) — the old pretrained=True flag is deprecated. Each model has a known receptive field and FLOPs budget; pick one that fits your latency target.
import torchvision.models as M
# pretrained backbones
resnet = M.resnet50(weights=M.ResNet50_Weights.DEFAULT)
effnet = M.efficientnet_b0(weights=M.EfficientNet_B0_Weights.DEFAULT)
vgg = M.vgg16(weights=M.VGG16_Weights.DEFAULT)
mobnet = M.mobilenet_v3_small(weights=M.MobileNet_V3_Small_Weights.DEFAULT)
# replace classifier head
num_classes = 100
resnet.fc = nn.Linear(resnet.fc.in_features, num_classes)
# feature extractor (drop head)
backbone = nn.Sequential(*list(resnet.children())[:-1]) # outputs (N, 2048, 1, 1)
features = backbone(images).flatten(1) # (N, 2048)RNN & LSTM
RNN Layer
batch_first=True makes the input (batch, seq, feature) which matches most NLP conventions; without it PyTorch uses (seq, batch, feature). Vanilla RNN suffers from vanishing gradients on long sequences — prefer LSTM or GRU in practice. out contains per-timestep hidden states from the last layer; h_n contains the final state of each layer.
import torch.nn as nn
rnn = nn.RNN(
input_size=64,
hidden_size=128,
num_layers=2,
batch_first=True, # input shape (batch, seq, feature)
bidirectional=False,
nonlinearity='tanh',
)
x = torch.randn(32, 10, 64) # (batch, seq, input)
out, h_n = rnn(x)
print(out.shape) # (32, 10, 128) — per-step hidden states
print(h_n.shape) # (2, 32, 128) — final hidden state per layer
# last-step output for classification
last = out[:, -1, :] # shape (32, 128)LSTM Layer
LSTM returns both hidden state h_n and cell state c_n. Bidirectional doubles the output feature size. For classification, use the final hidden state of each direction (h_n[-2] is the last forward layer, h_n[-1] is the last backward layer). dropout in nn.LSTM only applies between stacked layers, not on the input — wrap the input with nn.Dropout if you need that.
import torch.nn as nn
lstm = nn.LSTM(
input_size=64,
hidden_size=128,
num_layers=2,
batch_first=True,
bidirectional=True,
dropout=0.5, # between stacked layers
)
x = torch.randn(32, 10, 64)
out, (h_n, c_n) = lstm(x)
print(out.shape) # (32, 10, 256) — 128 * 2 directions
print(h_n.shape) # (4, 32, 128) — 2 layers * 2 directions
print(c_n.shape) # (4, 32, 128) — cell state
# concatenate final forward + backward for classification
last = torch.cat([h_n[-2], h_n[-1]], dim=1) # shape (32, 256)GRU Layer
GRU has only a hidden state (no cell state) and two gates instead of three, so it has fewer parameters and trains slightly faster than LSTM. Empirically GRU and LSTM perform similarly on most tasks; GRU is a good default for smaller datasets. Like LSTM, dropout only applies between stacked layers.
import torch.nn as nn
gru = nn.GRU(
input_size=64,
hidden_size=128,
num_layers=2,
batch_first=True,
)
x = torch.randn(32, 10, 64)
out, h_n = gru(x) # GRU has no cell state
print(out.shape) # (32, 10, 128)
print(h_n.shape) # (2, 32, 128)
# GRU vs LSTM: fewer parameters, slightly faster,
# often similar performance on many tasksPacked Sequences
Packing skips the padding tokens in the RNN computation, which is faster and avoids biasing the hidden state with pad zeros. Sequences must be sorted by length descending for older PyTorch; enforce_sorted=False sorts internally (with a small overhead). pad_packed_sequence converts back to a padded tensor. Combine with a custom collate_fn that pads batches dynamically.
from torch.nn.utils.rnn import pack_padded_sequence, pad_packed_sequence
# variable-length sequences in a batch
seqs = [torch.randn(5, 64), torch.randn(3, 64), torch.randn(8, 64)]
lens = torch.tensor([5, 3, 8])
# pad to (batch, max_len, feature)
padded = torch.nn.utils.rnn.pad_sequence(seqs, batch_first=True)
# sort by length descending (required by older versions)
padded = padded[torch.argsort(lens, descending=True)]
lens_sorted = lens[torch.argsort(lens, descending=True)]
packed = pack_padded_sequence(padded, lens_sorted,
batch_first=True, enforce_sorted=True)
out_packed, h_n = lstm(packed)
out, _ = pad_packed_sequence(out_packed, batch_first=True)Sequence Classification
A standard text classifier: embedding -> biLSTM -> dropout -> linear. Use padding_idx in nn.Embedding so the pad token has zero gradient. Pass lengths.cpu() to pack_padded_sequence because it runs on CPU. The final hidden state of both directions is concatenated to capture the whole sequence. For long sequences, transformers outperform LSTMs but need more data and compute.
import torch.nn as nn
class TextClassifier(nn.Module):
def __init__(self, vocab, emb_dim=128, hidden=256, classes=10):
super().__init__()
self.emb = nn.Embedding(vocab, emb_dim, padding_idx=0)
self.lstm = nn.LSTM(emb_dim, hidden, batch_first=True,
bidirectional=True, num_layers=2)
self.dropout = nn.Dropout(0.5)
self.fc = nn.Linear(hidden * 2, classes)
def forward(self, x, lengths):
x = self.emb(x)
packed = pack_padded_sequence(x, lengths.cpu(),
batch_first=True, enforce_sorted=False)
_, (h_n, _) = self.lstm(packed)
last = torch.cat([h_n[-2], h_n[-1]], dim=1)
return self.fc(self.dropout(last))
model = TextClassifier(vocab=10000)Sequence Generation (Decoder)
A decoder runs one step at a time, conditioning on the previous output token. LSTMCell is a single-step LSTM — use it for fine-grained control over the loop. Teacher forcing feeds the ground-truth token at the next step during training, which converges much faster than feeding the model's own predictions. At inference, switch to feeding argmax (or sample) from the previous step's logits.
import torch.nn as nn
class Decoder(nn.Module):
def __init__(self, vocab, emb_dim=128, hidden=256):
super().__init__()
self.emb = nn.Embedding(vocab, emb_dim)
self.lstm = nn.LSTMCell(emb_dim, hidden)
self.fc = nn.Linear(hidden, vocab)
def forward(self, tgt, hidden):
# tgt shape (batch, T)
outputs = []
h, c = hidden
for t in range(tgt.size(1)):
emb = self.emb(tgt[:, t])
h, c = self.lstm(emb, (h, c))
outputs.append(self.fc(h))
return torch.stack(outputs, dim=1) # (batch, T, vocab)
# teacher forcing: feed ground truth next token
# at inference: feed argmax of previous outputTransformer
nn.MultiheadAttention
nn.MultiheadAttention computes scaled dot-product attention over heads in parallel. batch_first matches the NLP convention (batch, seq, dim). Use attn_mask for causal decoding (lower-triangular mask) and key_padding_mask to ignore padding tokens. need_weights=True returns the attention matrix, useful for visualization and probing.
import torch.nn as nn
attn = nn.MultiheadAttention(
embed_dim=512,
num_heads=8,
dropout=0.1,
batch_first=True, # input (batch, seq, dim)
)
q = k = v = torch.randn(32, 10, 512)
out, weights = attn(q, k, v, need_weights=True)
print(out.shape) # (32, 10, 512)
print(weights.shape) # (32, 10, 10) — attention scores
# causal / self-attention mask
mask = nn.Transformer.generate_square_subsequent_mask(10)
out, _ = attn(q, k, v, attn_mask=mask) # masked attention
# key padding mask (ignore pad positions)
key_padding_mask = (tokens == 0) # True = ignore
out, _ = attn(q, k, v, key_padding_mask=key_padding_mask)Transformer Encoder
TransformerEncoderLayer is a self-attention + FFN block with residuals and LayerNorm. norm_first=True (pre-LN) is more stable for training from scratch; post-LN (default) needs a warmup. Stack N layers with nn.TransformerEncoder. Pass src_key_padding_mask so attention ignores pad positions. For classification, mean-pool the output sequence and feed a Linear head.
import torch.nn as nn
enc_layer = nn.TransformerEncoderLayer(
d_model=512,
nhead=8,
dim_feedforward=2048,
dropout=0.1,
activation='relu',
batch_first=True,
norm_first=True, # pre-LN, more stable
)
encoder = nn.TransformerEncoder(enc_layer, num_layers=6)
x = torch.randn(32, 10, 512)
out = encoder(x)
print(out.shape) # (32, 10, 512)
# with padding mask
pad_mask = (tokens == 0)
out = encoder(x, src_key_padding_mask=pad_mask)Positional Encoding
Transformers are permutation-invariant, so they need positional information. The sinusoidal encoding generalizes to longer sequences than seen in training. register_buffer stores pe without making it a parameter — it moves with .to(device) and is saved in state_dict. Learned position embeddings (nn.Parameter) work too but do not extrapolate beyond the training length.
import torch
import torch.nn as nn
import math
class PositionalEncoding(nn.Module):
def __init__(self, d_model=512, max_len=5000, dropout=0.1):
super().__init__()
self.dropout = nn.Dropout(dropout)
pe = torch.zeros(max_len, d_model)
pos = torch.arange(0, max_len).unsqueeze(1).float()
div = torch.exp(torch.arange(0, d_model, 2).float()
* (-math.log(10000.0) / d_model))
pe[:, 0::2] = torch.sin(pos * div)
pe[:, 1::2] = torch.cos(pos * div)
self.register_buffer('pe', pe.unsqueeze(0)) # (1, max_len, d_model)
def forward(self, x):
x = x + self.pe[:, :x.size(1)]
return self.dropout(x)
# usage
emb = nn.Embedding(vocab, 512)
pos = PositionalEncoding(512)
x = pos(emb(tokens)) # (batch, seq, 512)Full Transformer (Seq2Seq)
A full encoder-decoder transformer for translation/summarization. The decoder uses causal masking on the target (so position t can only attend to <= t) plus cross-attention to the encoder output (memory). Pass key_padding_masks for both source and target so attention ignores pad tokens. The head projects decoder outputs to vocabulary logits; loss is CrossEntropyLoss with ignore_index=0.
import torch.nn as nn
class Seq2SeqTransformer(nn.Module):
def __init__(self, vocab_src, vocab_tgt, d=512, h=8, layers=6):
super().__init__()
self.src_emb = nn.Embedding(vocab_src, d, padding_idx=0)
self.tgt_emb = nn.Embedding(vocab_tgt, d, padding_idx=0)
self.pos = PositionalEncoding(d)
enc_layer = nn.TransformerEncoderLayer(d, h, d*4, batch_first=True, norm_first=True)
dec_layer = nn.TransformerDecoderLayer(d, h, d*4, batch_first=True, norm_first=True)
self.encoder = nn.TransformerEncoder(enc_layer, layers)
self.decoder = nn.TransformerDecoder(dec_layer, layers)
self.head = nn.Linear(d, vocab_tgt)
def forward(self, src, tgt, src_pad, tgt_pad):
src = self.pos(self.src_emb(src))
tgt = self.pos(self.tgt_emb(tgt))
memory = self.encoder(src, src_key_padding_mask=src_pad)
causal = nn.Transformer.generate_square_subsequent_mask(tgt.size(1))
out = self.decoder(tgt, memory,
tgt_mask=causal,
tgt_key_padding_mask=tgt_pad,
memory_key_padding_mask=src_pad)
return self.head(out)Vision Transformer (ViT) Block
A ViT block is pre-LN self-attention + pre-LN MLP with residuals — essentially the same as a transformer encoder layer. Vision Transformers patchify the image with a strided Conv2d (kernel=stride=patch_size) then treat patches as tokens. Add a CLS token and positional embeddings. nn.MultiheadAttention returns (output, weights) — unpack with need_weights=False to skip the weights.
import torch.nn as nn
class ViTBlock(nn.Module):
def __init__(self, dim=768, heads=12, mlp=3072, dropout=0.1):
super().__init__()
self.norm1 = nn.LayerNorm(dim)
self.attn = nn.MultiheadAttention(dim, heads, dropout=dropout, batch_first=True)
self.norm2 = nn.LayerNorm(dim)
self.mlp = nn.Sequential(
nn.Linear(dim, mlp),
nn.GELU(),
nn.Dropout(dropout),
nn.Linear(mlp, dim),
nn.Dropout(dropout),
)
def forward(self, x):
h = self.norm1(x)
a, _ = self.attn(h, h, h, need_weights=False)
x = x + a
x = x + self.mlp(self.norm2(x))
return x
# patchify: (B, 3, 224, 224) -> (B, 196, 768) via Conv2d(3, 768, 16, 16)Attention from Scratch
Scaled dot-product attention: scores = QK^T / sqrt(d_k), softmax, multiply by V. The 1/sqrt(d_k) keeps variance stable so softmax does not saturate. masked_fill with -inf before softmax zeroes out those positions. Multi-head attention splits the embedding into H heads, attends in parallel, and concatenates. This is exactly what nn.MultiheadAttention does — write it from scratch only to learn or to add variants like linear/flash attention.
import torch
import torch.nn.functional as F
import math
def attention(q, k, v, mask=None, dropout=None):
# q, k, v shape (batch, heads, seq, dim)
d_k = q.size(-1)
scores = (q @ k.transpose(-2, -1)) / math.sqrt(d_k)
if mask is not None:
scores = scores.masked_fill(mask == 0, float('-inf'))
attn = F.softmax(scores, dim=-1)
if dropout is not None:
attn = dropout(attn)
out = attn @ v
return out, attn
# multi-head: reshape (batch, seq, dim) -> (batch, heads, seq, dim//heads)
def split_heads(x, heads):
b, s, d = x.shape
return x.view(b, s, heads, d // heads).transpose(1, 2)GPU Training
Device Selection
Always gate on torch.cuda.is_available() so CPU-only machines do not crash. You can create tensors directly on a device with the device= argument (faster than creating on CPU and moving). torch.device context (PyTorch 2.x) sets a default device for new tensors. MPS supports Apple Silicon GPUs but lacks some CUDA features — check operator support before relying on it.
import torch
# pick device
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
print(device) # cuda or cpu
print(torch.cuda.device_count()) # number of GPUs
print(torch.cuda.get_device_name(0)) # GPU name
# use a specific GPU
torch.cuda.set_device(0)
x = torch.randn(3, device='cuda:1') # directly on GPU 1
# default device context (PyTorch 2.x)
with torch.device('cuda'):
y = torch.randn(4) # on GPU
# Apple Silicon
device = 'mps' if torch.backends.mps.is_available() else 'cpu'Moving Models & Data
model.to(device) moves parameters and buffers in place. Move input tensors inside the training loop so the DataLoader can stream from disk. non_blocking=True overlaps the CPU→GPU copy with compute when the source tensor is in pinned memory (set pin_memory=True on the DataLoader). .cpu() / .numpy() before any NumPy op — GPU tensors cannot be converted directly.
model = MLP(784, 128, 10).to(device)
for x, y in train_loader:
x, y = x.to(device), y.to(device)
# non-blocking copy if pin_memory=True
x = x.to(device, non_blocking=True)
...
# check placement
print(next(model.parameters()).device) # cuda:0
print(x.device)
# move a tensor to CPU for numpy / saving
arr = x.detach().cpu().numpy()
# half precision
model = model.half() # float16
x = x.half()Multiple GPUs (DataParallel)
DataParallel (DP) replicates the model on each GPU, splits the batch, and gathers outputs on GPU 0 — easy but slow and unbalanced (GPU 0 does the gather and loss compute). It is single-process so the GIL limits scaling. For any serious multi-GPU work, use DistributedDataParallel instead. DP is fine for quick prototyping on 2-4 GPUs.
import torch.nn as nn
# simple, single-process data parallelism
model = MLP(784, 512, 10).to(device)
model = nn.DataParallel(model) # wrap AFTER .to(device)
# now model behaves normally
out = model(x)
loss = criterion(out, y)
loss.backward() # gradients are summed across GPUs
# access underlying module
inner = model.module
# specify which GPUs
model = nn.DataParallel(model, device_ids=[0, 1, 2])DistributedDataParallel
DistributedDataParallel (DDP) spawns one process per GPU, each with its own model replica; gradients are all-reduced so every GPU sees the same gradient. It is much faster and more scalable than DataParallel. Use torchrun (or torch.distributed.launch) to start processes. Always call sampler.set_epoch(epoch) so each epoch reshuffles. See the distributed-training section for more.
import os
import torch.distributed as dist
import torch.nn as nn
# launch with: torchrun --nproc_per_node=4 train.py
dist.init_process_group(backend='nccl')
local_rank = int(os.environ['LOCAL_RANK'])
torch.cuda.set_device(local_rank)
model = MLP(784, 512, 10).to(local_rank)
model = nn.parallel.DistributedDataParallel(
model, device_ids=[local_rank], output_device=local_rank,
)
sampler = torch.utils.data.distributed.DistributedSampler(dataset)
loader = DataLoader(dataset, batch_size=64, sampler=sampler)
for epoch in range(epochs):
sampler.set_epoch(epoch) # shuffle differently each epoch
for x, y in loader:
x, y = x.to(local_rank), y.to(local_rank)
loss = criterion(model(x), y)
loss.backward()
optimizer.step()
optimizer.zero_grad()Memory Management
empty_cache() returns cached memory to the OS but does not free tensors still in use — call it after a big spike, not in the hot loop. memory_allocated is tensors in use; memory_reserved is the caching allocator's pool. To reduce OOM: lower batch size, use mixed precision, gradient checkpointing, or gradient accumulation. expandable_segments:True reduces fragmentation for variable-shape workloads. Avoid retaining references to tensors you no longer need (e.g. logged losses in a list).
import torch
# release unused memory
torch.cuda.empty_cache()
# current memory usage
print(torch.cuda.memory_allocated() / 1e9, 'GB allocated')
print(torch.cuda.memory_reserved() / 1e9, 'GB reserved')
# per-tensor memory
x = torch.randn(1000, 1000, device='cuda')
print(x.element_size() * x.nelement() / 1e6, 'MB')
# peak memory tracking
torch.cuda.reset_peak_memory_stats()
# ... run training ...
print(torch.cuda.max_memory_allocated() / 1e9, 'GB peak')
# out-of-memory debugging
import os
os.environ['PYTORCH_CUDA_ALLOC_CONF'] = 'expandable_segments:True'CuDNN & Performance
cudnn.benchmark=True lets cuDNN profile algorithms on the first forward and cache the fastest — great when input shapes are fixed (typical CNN), harmful if shapes change every step (RNNs with variable length). TF32 uses 19-bit mantissa on Ampere+ for ~3x matmul speedup with minimal accuracy loss; enable it unless you need strict fp32 reproducibility. Deterministic mode disables non-deterministic algorithms.
import torch.backends.cudnn as cudnn
# benchmark finds the fastest conv algorithm (good for fixed input sizes)
cudnn.benchmark = True
# deterministic mode (slower, reproducible)
cudnn.deterministic = True
cudnn.benchmark = False
# allow TF32 on Ampere+ (big speedup, slight precision loss)
torch.backends.cuda.matmul.allow_tf32 = True
cudnn.allow_tf32 = True
# disable cudnn (debugging)
cudnn.enabled = FalseTransforms & Preprocessing
Common Transforms
ToTensor converts a PIL image (0-255) to a float tensor in [0,1] and rearranges HWC -> CHW. Normalize then standardizes per channel using ImageNet statistics (the de-facto standard for pretrained models — always match the stats the backbone was trained with). Apply transforms in the Dataset so they run on worker processes; the GPU never waits on CPU preprocessing.
from torchvision import transforms
transform = transforms.Compose([
transforms.Resize(256),
transforms.CenterCrop(224),
transforms.ToTensor(), # PIL -> Tensor, scales to [0,1]
transforms.Normalize( # per-channel standardize
mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225],
),
])
# apply to a PIL image
img = transform(pil_image)
print(img.shape) # (3, 224, 224)
print(img.min(), img.max()) # standardized
dataset = datasets.ImageFolder('data/', transform=transform)Data Augmentation
Augmentation regularizes the model by showing it varied inputs. RandomResizedCrop + flip is the classic ImageNet augmentation. ColorJitter changes color without affecting geometry. TrivialAugmentWide and RandAugment sample random ops with random magnitudes — strong, simple, and surprisingly effective. Use a separate non-augmented transform for validation so metrics are stable.
from torchvision import transforms
train_tf = transforms.Compose([
transforms.RandomResizedCrop(224),
transforms.RandomHorizontalFlip(),
transforms.ColorJitter(brightness=0.2, contrast=0.2,
saturation=0.2, hue=0.1),
transforms.RandomRotation(15),
transforms.ToTensor(),
transforms.Normalize(mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225]),
])
# stronger augmentation: RandAugment / TrivialAugment
train_tf = transforms.Compose([
transforms.TrivialAugmentWide(),
transforms.Resize(256),
transforms.CenterCrop(224),
transforms.ToTensor(),
transforms.Normalize(mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225]),
])Transforms v2 (New API)
Transforms v2 (torchvision 0.17+) is the new API — faster, JIT-scriptable, and crucially works jointly on image/mask/bbox so detection and segmentation augmentations stay consistent. ToDtype + scale=True replaces ToTensor. Always pass antialias=True to resize ops for cleaner downsample. The legacy transforms module still works but won't get new features.
from torchvision.transforms import v2
transform = v2.Compose([
v2.RandomResizedCrop(224, antialias=True),
v2.RandomHorizontalFlip(),
v2.ToDtype(torch.float32, scale=True), # scale to [0,1]
v2.Normalize(mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225]),
])
# v2 works on images, videos, masks, and bounding boxes together
imgs, masks = transform(imgs, masks)
imgs, bboxes, labels = transform(imgs, bboxes, labels)
# supports more ops and JIT scripting
transform = v2.RandomChoice([
v2.RandomHorizontalFlip(),
v2.RandomVerticalFlip(),
])Custom Transform
A custom transform is any callable that takes a tensor (or PIL image) and returns one. Compose them with transforms.Compose. Apply noise after ToTensor so it operates on floats. For per-sample randomness use torch.rand inside __call__ so each call gets a fresh decision. Transforms v2 lets you write transform classes that handle joint (image, label) inputs by overriding forward.
import torch
from torchvision import transforms
class RandomGaussianNoise:
def __init__(self, std=0.1):
self.std = std
def __call__(self, x):
if self.std > 0:
return x + torch.randn_like(x) * self.std
return x
# compose with built-ins
transform = transforms.Compose([
transforms.ToTensor(),
RandomGaussianNoise(0.05),
transforms.Normalize([0.5], [0.5]),
])
# callable with probability
class RandomApply:
def __init__(self, fn, p=0.5):
self.fn, self.p = fn, p
def __call__(self, x):
return self.fn(x) if torch.rand(1) < self.p else xText & Tabular Preprocessing
For text, build a vocabulary from the training set (never val/test) and convert tokens to integer ids; <unk> handles unseen words and <pad> is used for batching variable-length sequences. For tabular data, standardize features using statistics computed on the training set only — leaking val/test statistics inflates metrics. Save the scaler/vocab alongside the model checkpoint so inference preprocesses identically.
import torch
from torchtext.vocab import build_vocab_from_iterator
# tokenize and numericalize text
tokens = ['hello', 'world', 'this', 'is', 'nlp']
vocab = build_vocab_from_iterator([tokens], specials=['<unk>', '<pad>'])
vocab.set_default_index(vocab['<unk>'])
# text -> token ids
ids = vocab(['hello', 'world', 'unknownword'])
print(ids) # e.g. [3, 4, 0]
# tabular: standardize features
from sklearn.preprocessing import StandardScaler
scaler = StandardScaler().fit(X_train)
X_train_t = torch.tensor(scaler.transform(X_train), dtype=torch.float32)
X_test_t = torch.tensor(scaler.transform(X_test), dtype=torch.float32)
# fit scaler ONLY on training data, apply to val/testCustom Dataset
Dataset Class
A Dataset just needs __len__ and __getitem__. Open files lazily inside __getitem__ so you can hold millions of paths in memory cheaply. Return (input, target) tuples and let the DataLoader's default collate stack them into batches. Keep __getitem__ cheap and side-effect-free because it runs on worker processes — do not mutate shared state.
from torch.utils.data import Dataset
class ImageDataset(Dataset):
def __init__(self, paths, labels, transform=None):
self.paths = paths
self.labels = labels
self.transform = transform
def __len__(self):
return len(self.paths)
def __getitem__(self, idx):
from PIL import Image
img = Image.open(self.paths[idx]).convert('RGB')
label = self.labels[idx]
if self.transform:
img = self.transform(img)
return img, label
dataset = ImageDataset(paths, labels, transform=train_tf)
img, label = dataset[0]Image Folder Dataset
ImageFolder expects one subdirectory per class and infers labels from folder names — the fastest way to start with image classification. class_to_idx maps names to integer labels; keep it consistent across train/val/test by reusing the same mapping. random_split returns Subsets that share the underlying data, so memory is not duplicated.
from torchvision import datasets
# convention: data/train/cat/*.jpg, data/train/dog/*.jpg
dataset = datasets.ImageFolder('data/train', transform=train_tf)
print(dataset.class_to_idx) # {'cat': 0, 'dog': 1}
# train / val split
from torch.utils.data import random_split
train, val = random_split(dataset, [0.8, 0.2])
# custom split with same class mapping
val_dataset = datasets.ImageFolder('data/val', transform=val_tf)
val_dataset.class_to_idx = dataset.class_to_idx # keep consistent
# use with DataLoader
loader = DataLoader(dataset, batch_size=32, shuffle=True, num_workers=4)Text Dataset
Pre-tokenize text once and store integer ids to avoid re-tokenizing every epoch. For variable-length sequences, return raw lists and pad in collate_fn so each batch pads only to its own max length. pack_padded_sequence then skips pad positions in the RNN. For transformer training, return attention masks so padding tokens are ignored by the attention.
from torch.utils.data import Dataset
import torch
class TextDataset(Dataset):
def __init__(self, texts, labels, vocab, max_len=128):
self.texts = texts
self.labels = labels
self.vocab = vocab
self.max_len = max_len
def __len__(self):
return len(self.texts)
def __getitem__(self, idx):
tokens = self.texts[idx].split()[:self.max_len]
ids = [self.vocab[t] for t in tokens]
ids += [0] * (self.max_len - len(ids)) # pad
return torch.tensor(ids), self.labels[idx]
# variable length: use custom collate
def text_collate(batch):
ids, labels = zip(*batch)
lens = torch.tensor([len(x) for x in ids])
padded = torch.nn.utils.rnn.pad_sequence(ids, batch_first=True)
return padded, lens, torch.tensor(labels)Iterable Dataset
IterableDataset streams data instead of random access — essential for huge or streaming sources (logs, databases, S3). Workers must explicitly split the stream (itertools.islice) to avoid every worker reading the same data. shuffle is not supported directly; instead, buffer samples and shuffle within the buffer. Use this for datasets too large to index by length.
from torch.utils.data import IterableDataset, DataLoader
import itertools
class StreamDataset(IterableDataset):
def __init__(self, file_path):
self.file_path = file_path
def __iter__(self):
# split work across workers
worker_info = torch.utils.data.get_worker_info()
with open(self.file_path) as f:
if worker_info is None:
lines = f
else:
lines = itertools.islice(f, worker_info.id,
None, worker_info.num_workers)
for line in lines:
yield self._parse(line)
def _parse(self, line):
# parse line into tensor + label
return torch.tensor([...]), label
loader = DataLoader(StreamDataset('big.csv'), batch_size=32)TensorDataset & ConcatDataset
TensorDataset is the fastest path when data fits in RAM — it returns tuples of tensors that the default collate stacks into batches. ConcatDataset chains datasets end-to-end; useful for combining labeled and pseudo-labeled data. Subset wraps a dataset with an index list, enabling random_split and custom slicing without copying data.
from torch.utils.data import TensorDataset, ConcatDataset, DataLoader
# wrap in-memory tensors
X = torch.randn(1000, 784)
y = torch.randint(0, 10, (1000,))
dataset = TensorDataset(X, y)
x, label = dataset[0]
# combine multiple datasets
full = ConcatDataset([dataset_a, dataset_b])
print(len(full)) # len(a) + len(b)
# Subset for indexing
from torch.utils.data import Subset
half = Subset(dataset, range(500))
# random splits
from torch.utils.data import random_split
train, val = random_split(dataset, [0.8, 0.2])LR Scheduler
StepLR & MultiStepLR
StepLR and MultiStepLR are the simplest schedulers — multiply LR by gamma at fixed epochs. They are predictable and work well for CNNs trained from scratch. Always call scheduler.step() after the epoch's training, not inside the batch loop. get_last_lr() shows the current LR for logging. Watch out: step_size counts epochs, not steps.
from torch.optim.lr_scheduler import StepLR, MultiStepLR
# decay by gamma every step_size epochs
scheduler = StepLR(optimizer, step_size=10, gamma=0.1)
# LR: 1e-3 -> 1e-4 (epoch 10) -> 1e-5 (epoch 20)
# decay at specific epochs
scheduler = MultiStepLR(optimizer, milestones=[30, 60, 80], gamma=0.1)
# in the training loop
for epoch in range(epochs):
train(...)
val(...)
scheduler.step() # call once per epoch
# inspect current LR
print(scheduler.get_last_lr())CosineAnnealingLR
Cosine annealing smoothly decays the LR — often a small accuracy bump over step decay. Warm restarts (SGDR) reset the LR to the base value at the start of each cycle, which can help the model escape sharp minima; T_mult=2 makes each cycle twice as long. Cosine schedulers are the modern default for image classification and pretraining.
from torch.optim.lr_scheduler import CosineAnnealingLR
# cosine decay to eta_min over T_max epochs
scheduler = CosineAnnealingLR(optimizer, T_max=50, eta_min=1e-6)
# LR follows cosine from base_lr down to eta_min
# with warm restarts
from torch.optim.lr_scheduler import CosineAnnealingWarmRestarts
scheduler = CosineAnnealingWarmRestarts(
optimizer,
T_0=10, # length of first cycle
T_mult=2, # each cycle is T_mult longer
eta_min=1e-6,
)ReduceLROnPlateau
ReduceLROnPlateau watches a metric and reduces LR when it stops improving — perfect when you don't know the right schedule in advance. mode='min' for losses, 'max' for accuracy. patience=N waits N epochs of no improvement before reducing. Unlike other schedulers, you must pass the metric to step(metric). Easy to overfit with too much patience; combine with early stopping.
from torch.optim.lr_scheduler import ReduceLROnPlateau
scheduler = ReduceLROnPlateau(
optimizer,
mode='min', # 'min' for loss, 'max' for accuracy
factor=0.5, # multiply LR by factor on plateau
patience=3, # epochs to wait before reducing
min_lr=1e-7,
threshold=1e-4,
)
for epoch in range(epochs):
train(...)
val_loss = validate(...)
scheduler.step(val_loss) # pass the metric!OneCycleLR & CyclicLR
OneCycleLR (super-convergence) warms up, peaks, then anneals — often the fastest way to train a model to a given accuracy. It steps per batch, so set total_steps = epochs * batches_per_epoch. pct_start=0.3 spends 30% of steps warming up. CyclicLR oscillates between base and max LR, useful for finding good LRs (LR range test).
from torch.optim.lr_scheduler import OneCycleLR
# one cycle: warm up, peak, anneal down
scheduler = OneCycleLR(
optimizer,
max_lr=1e-3,
total_steps=epochs * len(loader), # or steps_per_epoch + epochs
pct_start=0.3, # fraction spent warming up
anneal_strategy='cos',
)
# step per BATCH (not per epoch)
for epoch in range(epochs):
for x, y in loader:
train_step(...)
scheduler.step() # call once per batch
# simpler cyclic LR
from torch.optim.lr_scheduler import CyclicLR
scheduler = CyclicLR(optimizer, base_lr=1e-5, max_lr=1e-3, step_size_up=500)Warmup + Decay
Warmup stabilizes the first steps when Adam's second-moment estimate is poorly conditioned — essential for transformer training. LinearLR + SequentialLR cleanly combines warmup with any decay schedule. LambdaLR gives full control via a function of the step count — the most flexible scheduler. For HuggingFace-style schedules, look at transformers.get_cosine_schedule_with_warmup.
from torch.optim.lr_scheduler import LinearLR, SequentialLR, CosineAnnealingLR
# linear warmup for first N steps
warmup = LinearLR(optimizer, start_factor=0.1, total_iters=500)
# cosine decay after warmup
cosine = CosineAnnealingLR(optimizer, T_max=total_steps - 500, eta_min=1e-6)
# combine
scheduler = SequentialLR(
optimizer,
schedulers=[warmup, cosine],
milestones=[500],
)
# manual warmup (transformers-style)
def lr_lambda(step):
if step < 500:
return step / 500
return 0.5 * (1 + math.cos(math.pi * (step - 500) / (total - 500)))
from torch.optim.lr_scheduler import LambdaLR
scheduler = LambdaLR(optimizer, lr_lambda)Custom Scheduler
LambdaLR multiplies the base LR by lambda(step) — the most flexible way to implement custom schedules. Pass a list of lambdas (one per param_group) for per-group schedules, e.g. decay weights but not biases, or use different schedules for backbone vs head in transfer learning. Lambdas must be cheap because they run every step. Avoid Python closures that capture mutable state — keep them pure.
from torch.optim.lr_scheduler import LambdaLR
import math
# any function of step / epoch
def warmup_cosine(step, warmup=500, total=10000):
if step < warmup:
return step / warmup
progress = (step - warmup) / (total - warmup)
return 0.5 * (1 + math.cos(math.pi * progress))
scheduler = LambdaLR(optimizer, lr_lambda=warmup_cosine)
# per-parameter-group schedule (e.g. decay only weights, not bias)
scheduler = LambdaLR(
optimizer,
lr_lambda=[lambda s: 0.9 ** s, lambda s: 1.0], # one per group
)
# inspect
for step in range(0, total, 1000):
print(step, scheduler.get_last_lr())
scheduler.step()Transfer Learning
Pretrained Models
Always use the WeightsEnum API (DEFAULT picks the best available) — the old pretrained=True is deprecated. Each weights object exposes .transforms() returning the exact preprocessing the model was trained with — use it so your inputs match. The classifier head outputs ImageNet's 1000 classes; replace it for your own task.
import torchvision.models as M
# load a model with pretrained ImageNet weights
weights = M.ResNet50_Weights.DEFAULT
model = M.resnet50(weights=weights)
# get the matching preprocessing transform
preprocess = weights.transforms()
print(preprocess) # resize, crop, normalize with ImageNet stats
# model is ready for 1000-class ImageNet inference
model.eval()
with torch.no_grad():
out = model(preprocess(img).unsqueeze(0))
pred = out.argmax(1)
print(weights.meta['categories'][pred])Feature Extraction
Feature extraction freezes the backbone and trains only a new head — the fastest transfer-learning baseline, often reaching high accuracy with minutes of training on a small dataset. Because the backbone is frozen, you can precompute features once and train a cheap classifier on them, saving huge amounts of compute. Set model.eval() so BatchNorm uses running stats even while training the head.
import torch.nn as nn
import torchvision.models as M
# load backbone, freeze it, replace head
model = M.resnet50(weights=M.ResNet50_Weights.DEFAULT)
# freeze all parameters
for p in model.parameters():
p.requires_grad_(False)
# replace classifier for new task (10 classes)
model.fc = nn.Linear(model.fc.in_features, 10)
# only the new fc has requires_grad=True
# train only the head
optimizer = torch.optim.Adam(model.fc.parameters(), lr=1e-3)
# inference-friendly: extract features once
features = torch.nn.Sequential(*list(model.children())[:-1])
feat = features(images).flatten(1) # (N, 2048)Fine-tuning
Fine-tuning trains the whole model (or part of it) with a small LR so the pretrained features adapt to your task. Use differential LRs: smaller for early layers (general features), larger for the head (task-specific). Gradual unfreezing — train the head first, then unfreeze deeper layers — is more stable than training everything at once, especially on small datasets.
import torch.nn as nn
import torch.optim as optim
import torchvision.models as M
model = M.resnet50(weights=M.ResNet50_Weights.DEFAULT)
model.fc = nn.Linear(model.fc.in_features, 10)
# differential learning rates: lower for backbone, higher for head
optimizer = optim.Adam([
{'params': [p for n, p in model.named_parameters() if 'fc' not in n], 'lr': 1e-4},
{'params': model.fc.parameters(), 'lr': 1e-3},
])
# gradual unfreezing: train head first, then unfreeze later layers
def unfreeze(layer_idx):
for n, p in model.named_parameters():
if f'layer{layer_idx}' in n:
p.requires_grad_(True)
# warmup head only for a few epochs, then unfreeze layer4, etc.Freezing & Unfreezing Layers
requires_grad_(False) stops autograd from tracking and updating a parameter — but gradients still flow through frozen layers to trainable ones downstream. Use named_parameters() to filter by layer name for selective freezing. Always log the trainable parameter count so you don't accidentally freeze everything (a common silent bug). After unfreezing, re-create the optimizer so it picks up the newly trainable params.
model = M.resnet50(weights=M.ResNet50_Weights.DEFAULT)
model.fc = nn.Linear(model.fc.in_features, 10)
# freeze everything except fc
for name, p in model.named_parameters():
p.requires_grad = ('fc' in name)
# freeze specific submodule
for p in model.layer1.parameters():
p.requires_grad_(False)
# count trainable params
n = sum(p.numel() for p in model.parameters() if p.requires_grad)
print(f"trainable: {n:,}")
# unfreeze after warmup
for p in model.layer4.parameters():
p.requires_grad_(True)
# trainable params still need grad to flow: ensure no torch.no_grad
# and that frozen layers don't break autogradReplacing the Classifier
Each architecture has its classifier in a different attribute — ResNet uses .fc, EfficientNet uses .classifier[1], ViT uses .heads.head. Setting the new head with requires_grad=True means only it trains initially. For multi-task learning, replace the head with multiple Linear layers sharing the backbone. nn.Identity() is a no-op module handy for dropping the original head cleanly.
import torch.nn as nn
import torchvision.models as M
model = M.resnet50(weights=M.ResNet50_Weights.DEFAULT)
# ResNet: model.fc
model.fc = nn.Linear(model.fc.in_features, num_classes)
# EfficientNet: model.classifier[1]
model = M.efficientnet_b0(weights=M.EfficientNet_B0_Weights.DEFAULT)
model.classifier[1] = nn.Linear(model.classifier[1].in_features, num_classes)
# ViT: model.heads.head
model = M.vit_b_16(weights=M.ViT_B_16_Weights.DEFAULT)
model.heads.head = nn.Linear(model.heads.head.in_features, num_classes)
# multi-task head
class MultiTask(nn.Module):
def __init__(self, backbone, n_cls, n_reg):
super().__init__()
self.backbone = backbone
in_f = backbone.fc.in_features
backbone.fc = nn.Identity()
self.cls = nn.Linear(in_f, n_cls)
self.reg = nn.Linear(in_f, n_reg)
def forward(self, x):
f = self.backbone(x)
return self.cls(f), self.reg(f)Visualization
TensorBoard Setup
SummaryWriter streams events to a log directory; launch TensorBoard with `tensorboard --logdir=runs` to view. add_scalar plots a single metric vs step; add_scalars overlays multiple curves on one chart. Use separate subdirectories per experiment (runs/exp1, runs/exp2) to compare runs in the same UI. Flush with writer.close() at the end so the last events are written.
from torch.utils.tensorboard import SummaryWriter
writer = SummaryWriter('runs/experiment_1')
# log scalars (loss, accuracy, LR)
writer.add_scalar('train/loss', loss.item(), step)
writer.add_scalar('val/accuracy', acc, step)
writer.add_scalar('lr', optimizer.param_groups[0]['lr'], step)
# log multiple metrics together
writer.add_scalars('losses', {
'train': train_loss,
'val': val_loss,
}, step)
# close at the end
writer.close()
# launch: tensorboard --logdir=runsLogging Graphs & Histograms
add_graph visualizes the model architecture in the Graphs tab — useful for debugging forward flow. add_histogram shows the distribution of weights or gradients over time; healthy training keeps gradients nonzero and weights roughly bell-shaped. Sudden spikes in grad histograms hint at instability. Logging every step is slow — sample every N steps.
writer = SummaryWriter('runs/exp')
# log model graph (needs a sample input)
model = MLP(784, 128, 10)
writer.add_graph(model, torch.randn(1, 784))
# weight histograms (per layer, over training)
for name, p in model.named_parameters():
writer.add_histogram(f'weights/{name}', p, step)
if p.grad is not None:
writer.add_histogram(f'grads/{name}', p.grad, step)
# text
writer.add_text('config', 'lr=1e-3, batch=32', 0)
# close at the end
writer.close()Logging Images
add_image takes a CHW tensor; use make_grid to tile a batch into one image with normalize=True so it scales to a displayable range. add_figure logs a matplotlib Figure directly. draw_segmentation_masks and draw_bounding_boxes (torchvision.utils) overlay predictions for quick visual QA. Log a few validation samples per epoch to monitor qualitative progress, not just scalar metrics.
writer = SummaryWriter('runs/exp')
# single image (CHW tensor in [0,1] or [0,255])
writer.add_image('val/sample', img_tensor, step)
# grid of images
from torchvision.utils import make_grid
grid = make_grid(images, nrow=8, normalize=True)
writer.add_image('val/batch', grid, step)
# figure from matplotlib
import matplotlib.pyplot as plt
fig, ax = plt.subplots()
ax.plot([1,2,3], [1,4,9])
writer.add_figure('curves/loss', fig, step)
# overlay segmentation mask on image
from torchvision.utils import draw_segmentation_masks
overlay = draw_segmentation_masks(img, masks, alpha=0.5)
writer.add_image('val/masks', overlay, step)torchviz Graph
torchviz draws the autograd graph of a single forward pass — great for debugging shape mismatches, unintended branches, or detached tensors. Each node is an op; edges show tensor flow. show_saved=True also displays intermediate tensors kept for backward, which is what determines memory usage. install with `pip install torchviz` and Graphviz on your system PATH.
from torchviz import make_dot
x = torch.randn(1, 784, requires_grad=True)
y = model(x)
loss = criterion(y, torch.tensor([1]))
loss.backward()
# render the autograd graph
dot = make_dot(loss, params=dict(model.named_parameters()))
dot.render('model_graph', format='png') # saves model_graph.png
# also show saved tensors (memory)
dot = make_dot(loss, params=dict(model.named_parameters()),
show_attrs=True, show_saved=True)Weight & Activation Probing
Visualizing first-layer filters shows what the model 'looks for' — edges, colors, textures. Forward hooks capture activations without modifying forward, perfect for probing intermediate layers or building feature extractors. Track weight and gradient statistics (mean, std, norm) per step to catch dead neurons (zero grads) or exploding activations. torchsummary or torchinfo prints a clean parameter table per layer.
import matplotlib.pyplot as plt
# visualize conv filters
filters = model.features[0].weight.data.clone().cpu()
fig, axes = plt.subplots(4, 8, figsize=(10, 5))
for i, ax in enumerate(axes.flat):
if i < filters.size(0):
ax.imshow(filters[i, 0], cmap='gray')
ax.axis('off')
# hook to capture activations
acts = {}
def hook(name):
def fn(mod, inp, out): acts[name] = out.detach().cpu()
return fn
model.features[3].register_forward_hook(hook('l1'))
# activation statistics over training
for name, p in model.named_parameters():
print(name, p.data.mean().item(), p.data.std().item(),
p.grad.norm().item() if p.grad is not None else None)Distributed Training
DDP Setup
DistributedDataParallel (DDP) is the standard multi-GPU training pattern: one process per GPU, gradients all-reduced so all replicas stay in sync. nccl is the fastest backend on NVIDIA GPUs. torchrun sets the env vars (RANK, WORLD_SIZE, LOCAL_RANK) for you. Always call dist.destroy_process_group() at the end so the job exits cleanly. See gpu-training section for the full loop.
import os, torch
import torch.distributed as dist
import torch.nn as nn
def setup():
dist.init_process_group(backend='nccl')
local_rank = int(os.environ['LOCAL_RANK'])
torch.cuda.set_device(local_rank)
return local_rank
def cleanup():
dist.destroy_process_group()
# launch with torchrun
# torchrun --nproc_per_node=4 train.py
# rank, world_size, LOCAL_RANK are set by torchrun
local_rank = setup()
model = model.to(local_rank)
model = nn.parallel.DistributedDataParallel(
model, device_ids=[local_rank], output_device=local_rank,
)DistributedSampler
DistributedSampler shards the dataset across processes so each GPU sees a unique subset of each batch (effective batch = batch_size * world_size). You MUST call sampler.set_epoch(epoch) every epoch — without it, every epoch uses the same shuffle and the model sees the same ordering. Do not pass shuffle=True to DataLoader when using a sampler; the sampler handles ordering.
from torch.utils.data import DataLoader
from torch.utils.data.distributed import DistributedSampler
sampler = DistributedSampler(
dataset,
num_replicas=dist.get_world_size(),
rank=dist.get_rank(),
shuffle=True,
drop_last=True,
)
loader = DataLoader(dataset, batch_size=64, sampler=sampler,
num_workers=4, pin_memory=True)
for epoch in range(epochs):
sampler.set_epoch(epoch) # reshuffle per epoch
for x, y in loader:
# each GPU sees a non-overlapping subset
...Multi-Node Launch
torchrun launches one process per GPU and connects them via TCP. For multi-node, run torchrun on each node with the same master_addr and a unique node_rank. The rendezvous endpoint (--rdzv_endpoint) lets nodes find each other and supports elastic restarts. Make sure master_port is open on the firewall and that all nodes can reach master_addr. NCCL also needs high-bandwidth interconnect (InfiniBand/NVLink) for good scaling.
# single-node, 4 GPUs
torchrun --nproc_per_node=4 train.py
# multi-node (2 nodes, 4 GPUs each = 8 total)
# node 0 (master):
torchrun --nproc_per_node=4 --nnodes=2 --node_rank=0 \
--master_addr=10.0.0.1 --master_port=29500 train.py
# node 1:
torchrun --nproc_per_node=4 --nnodes=2 --node_rank=1 \
--master_addr=10.0.0.1 --master_port=29500 train.py
# in code, world_size = nnodes * nproc_per_node
print(dist.get_world_size()) # 8
# also supports elastic restart on node failure
torchrun --nproc_per_node=4 --rdzv_backend=c10d \
--rdzv_endpoint=10.0.0.1:29500 train.pyGradient Synchronization
DDP overlaps gradient all-reduce with backward computation, which is why it scales so well. find_unused_parameters=True is needed when some parameters don't get gradients every step (e.g. conditional branches) but costs overhead — avoid it if you can. model.no_sync() skips the sync for a step, useful with gradient accumulation so you only sync once per effective batch instead of every micro-batch.
# DDP all-reduces gradients after backward by default
loss.backward() # gradients synchronized across GPUs
# skip sync for unused parameters (small speedup)
model = nn.parallel.DistributedDataParallel(
model, device_ids=[local_rank],
find_unused_parameters=False, # set True if some params unused
)
# manual gradient sync (advanced)
for p in model.parameters():
if p.grad is not None:
dist.all_reduce(p.grad, op=dist.ReduceOp.SUM)
p.grad /= dist.get_world_size()
# only sync on the last micro-batch (gradient accumulation)
# use model.no_sync() context for the non-final steps
with model.no_sync():
loss.backward()Logging & Checkpointing
Always gate logging and checkpointing on rank==0 to avoid duplicates and race conditions. When saving a DDP model, use model.module.state_dict() — DDP wraps the underlying module, and saving the wrapper makes the checkpoint DDP-specific and hard to load in non-distributed code. Load the checkpoint on every rank before wrapping with DDP so all replicas start from the same weights.
# only log from rank 0 to avoid duplicates
if dist.get_rank() == 0:
writer.add_scalar('train/loss', loss.item(), step)
print(f'step {step}: loss={loss.item():.4f}')
# only save checkpoint from rank 0
if dist.get_rank() == 0:
torch.save({
'model': model.module.state_dict(), # unwrap DDP
'optimizer': optimizer.state_dict(),
'epoch': epoch,
}, 'ckpt.pt')
# load on all ranks, then wrap with DDP
state = torch.load('ckpt.pt', map_location='cpu')
model.load_state_dict(state['model'])
model = model.to(local_rank)
model = nn.parallel.DistributedDataParallel(model, device_ids=[local_rank])FullyShardedDataParallel (FSDP)
FSDP shards parameters, gradients, and optimizer states across GPUs — so a model too big for one GPU can train across several. Unlike DDP (which keeps full replicas), FSDP gathers shards on demand during forward/backward. Use it when the model doesn't fit in a single GPU's memory. Saving is trickier: wrap state_dict collection in FULL_STATE_DICT mode so rank 0 materializes the full model. Mixed precision + FSDP + activation checkpointing lets you train very large models.
from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
from torch.distributed.fsdp import ShardingStrategy
model = LargeModel().to(local_rank)
model = FSDP(
model,
sharding_strategy=ShardingStrategy.FULL_SHARD,
device_id=local_rank,
use_orig_params=True, # needed for save/load compatibility
)
# forward / backward as usual
loss = criterion(model(x), y)
loss.backward()
optimizer.step()
# save: only the full state on rank 0
from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
import torch.distributed.fsdp.api as fsdp_api
with FSDP.state_dict_type(model, fsdp_api.StateDictType.FULL_STATE_DICT):
state = model.state_dict()
if dist.get_rank() == 0:
torch.save(state, 'fsdp_ckpt.pt')Snippets PyTorch associés
Copy-paste ready code for common tasks.
Tensor Basics
Create, index, and operate on tensors.
Autograd
Compute gradients automatically with backward().
Dataset and DataLoader
Build custom datasets and batch them with DataLoader.
Model Definition
Define models with nn.Module and Sequential.
Training Loop
Run a full train-eval loop with loss and optimizer.
GPU and CUDA
Move models and tensors to GPU and handle availability.
Save and Load
Checkpoint models, optimizer state, and weights.
Transfer Learning
Fine-tune a pretrained torchvision model.
Was this helpful?