Pangram verdict · v3.3
We believe that this document is a mix of AI-generated, and human-written content
AI likelihood · overall
MixedArticle text · 1,695 words · 6 segments analyzed
A three-class spiral dataset. Shaded regions show the model's softmax confidence. The boundary sharpens as training progresses. Building a PyTorch training loop is fairly straightforward, but getting everything in the right place and in the right order can feel surprisingly fragile. There are loads of moving parts and after the most basic errors are fixed, most of the other mistakes can be pretty hard to spot. Training runs will fail to converge, produce incorrect results, or consume excessive memory if lines are misplaced. The sections below will go through each operation in sequence, explaining exactly how to write each section, and all the common mistakes to watch out for. Distributed training, FSDP, and multi-GPU setups are out of scope here, but we'll come back to that in a future essay. (The animation above was produced by running the loop on synthetic data and capturing the decision boundary at each epoch.) The complete loop Let's look, first of all, at the complete training loop. You don't need to understand or memorise it yet, just get a feel for the structure. 1import torch 2import torch.nn as nn 3from torch.utils.data import DataLoader, TensorDataset 4 5# --- data --- 6dataset = TensorDataset(X_train, y_train) 7loader = DataLoader(dataset, batch_size=64, shuffle=True) 8 9# --- model, loss, optimiser --- 10model = MLP(in_features=2, hidden=128, out_features=3) 11criterion = nn.CrossEntropyLoss() 12optimiser = torch.optim.Adam(model.parameters(), lr=1e-3) 13scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimiser, T_max=100) 14 15# --- loop --- 16for epoch in range(100): 17 model.train() 18 for X_batch, y_batch in loader: 19 optimiser.zero_grad() 20 logits = model(X_batch) 21 loss = criterion(logits, y_batch) 22 loss.backward() 23 torch.nn.utils.clip_grad_norm_(model.parameters(),
max_norm=1.0) 24 optimiser.step() 25 scheduler.step() 26 27 model.eval() 28 with torch.no_grad(): 29 val_logits = model(X_val) 30 val_loss = criterion(val_logits, y_val) Now let's go through each line and understand what it does, and how not to break it. We'll start with some of the common mistakes. TL;DR Where the order really matters Here are some of the most common failures, and how you can break the training loop by getting the placement a little bit wrong. The reason to memorise these is that none of them will raise an exception, over time you'll get a sense for what kind of errors to look for in your training runs, but for the first few times this crib sheet will help you out. LineWrong positionWhat breaksmodel.to(device)After optimiser = ...When a dtype conversion is combined (e.g. model.half()), nn.Module.to() allocates new nn.Parameter objects; the optimiser holds references to the discarded originals and applies updates to them instead.optimiser.zero_grad()After loss.backward()Gradients from multiple batches accumulate. Update uses their sum, not the current batch alone.clip_grad_norm_()Before loss.backward().grad is empty. The call is a no-op.clip_grad_norm_()After optimiser.step()Clips gradients already applied. No effect.scheduler.step()Inside batch loopLR decays len(loader) times per epoch instead of once.Omit model.train() after model.eval()—Dropout disabled, BatchNorm frozen. The model trains in eval mode without error.Omit torch.no_grad() during validation—Autograd graph builds on every validation batch. Memory grows until OOM.Log loss instead of loss.item()—Pins the computation graph in memory for the duration of the logging call. Now let's go through each of these in detail. The data There are two parts to the data pipeline in PyTorch: the Dataset and the DataLoader. The Dataset is just a Python object that implements __len__ (how many elements are in the dataset) and __getitem__ (which, unsurprisingly, gets an item). It can be a simple wrapper around tensors, or it can load data from disk on demand. The DataLoader wraps a dataset and produces batches.
Each pass through the full dataset is one epoch. With shuffle=True, examples are presented in a different order each epoch. 1dataset = TensorDataset(X_train, y_train) 2loader = DataLoader( 3 dataset, 4 batch_size=64, 5 shuffle=True, 6 num_workers=2, 7 pin_memory=True, 8 persistent_workers=True, 9) practicePyTorch DataLoaderEasyQ342Wire up a PyTorch DataLoader: batching, shuffling, and iterating. TensorDataset pairs input and label tensors by index. Indexing with dataset[i] returns (X[i], y[i]). DataLoader calls __getitem__ repeatedly, collates the results into batches, and optionally hands work to background worker processes. num_workers spawns separate processes that prefetch batches in parallel with GPU compute. The main process blocks on .next() only if a batch is not yet ready. Zero workers means the main process does all loading, which often bottlenecks GPU utilisation on data-heavy tasks. Two to four workers is practical, but the right number depends on CPU count and I/O speed. pin_memory=True allocates batch tensors in pinned host memory. The GPU DMA engine can transfer directly from pinned memory without first copying through the kernel buffer, reducing host-to-device transfer time. It only helps when num_workers > 0 and you're transferring to CUDA. persistent_workers=True keeps worker processes alive between epochs. Without it, workers are respawned at the start of each epoch, adding fork overhead that becomes measurable at large worker counts. drop_last=True discards the final batch if it is smaller than batch_size. BatchNorm statistics computed from a batch of two or three samples are noisy so dropping the remainder avoids this.
Small cost in terms of dropping the data, but it is often worth it for stability. Smaller batches produce noisier gradient estimates, which acts as implicit regularisation. Larger batches use more GPU memory but allow more parallelism. One of the most important efficiencies is knowing that powers of two align with tensor core tile sizes (typically 16×16 or 8×16 depending on dtype), so making sure batch sizes and layer dimensions are set to multiples of 8 or 16 is a good idea. .to(device) moves a tensor to the target device. For tensors, it is not in-place: it returns a new tensor and leaves the original unchanged. For example, X_batch.to('cuda') returns a new tensor on GPU; X_batch itself remains on CPU. Reproducibility Setting seeds before constructing the model and loader gives the same results on every run. This is essential in order to reproduce experiments, and make sure model behavior is deterministic. The main areas this impacts the model is in the data loader, and in the initialisation of the model weights. 1import random 2import numpy as np 3 4def set_seed(seed: int = 42): 5 torch.manual_seed(seed) 6 torch.cuda.manual_seed_all(seed) 7 np.random.seed(seed) 8 random.seed(seed) 9 torch.backends.cudnn.deterministic = True 10 torch.backends.cudnn.benchmark = False 11 12set_seed(42) torch.manual_seed seeds the CPU generator. torch.cuda.manual_seed_all seeds every GPU. NumPy and Python's random are independent RNGs that PyTorch does not touch, be careful of this, you might need another random seed for them. cudnn.deterministic = True forces cuDNN to use deterministic convolution algorithms. Some cuDNN kernels are non-deterministic by default for throughput. The deterministic alternatives are slightly slower, but practically it shouldn't matter much during development. cudnn.benchmark = False must be paired with deterministic = True. When benchmark = True, cuDNN profiles several algorithms per input shape and picks the fastest, a process that itself varies between runs. Fixing it to False makes sure you always get the same results.
When num_workers > 0, each DataLoader worker has its own RNG state, seeded by the OS when it forks the processes. To make worker randomness reproducible you should pass a generator and a worker_init_fn: 1g = torch.Generator() 2g.manual_seed(42) 3 4loader = DataLoader( 5 dataset, 6 batch_size=64, 7 shuffle=True, 8 num_workers=2, 9 generator=g, 10 worker_init_fn=lambda worker_id: np.random.seed(42 + worker_id), 11) The model nn.Module provides parameter tracking, device movement, train/eval mode switching, and serialisation. Each instance needs an __init__ (to register all the submodules) and forward (to define the computation). practiceReLU Activation + BackwardEasyQ114Implement a ReLU activation and its backward pass, in numpy, to see what nn.ReLU is actually doing. practiceCustom Linear ModuleEasyQ328Implement a custom nn.Linear as an nn.Module, registering weights and bias as parameters. 1class MLP(nn.Module): 2 def __init__(self, in_features, hidden, out_features): 3 super().__init__() 4 self.net = nn.Sequential( 5 nn.Linear(in_features, hidden), 6 nn.ReLU(), 7 nn.Linear(hidden, hidden), 8 nn.ReLU(), 9 nn.Linear(hidden, out_features), 10 ) 11 12 def forward(self, x): 13 return self.net(x) 14 15device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') 16model = MLP(in_features=2, hidden=128, out_features=3).to(device) super().__init__() is required: it initialises the module registry. Assigning submodules or parameter tensors as attributes in __init__ registers them automatically. Unregistered plain attributes are excluded from parameter iteration, serialisation, and device movement. .to(device) moves all registered parameters and buffers to the target device in-place. Call it before constructing the optimiser. For a device-only move, nn.Module.to() modifies each parameter's .data attribute in-place and the optimiser's references remain correct.
When a dtype conversion is combined (e.g. .half().to(device)), nn.Module.to() allocates new nn.Parameter objects and replaces them in the module's internal registry; the optimiser, constructed before this, retains references to the originals and applies updates to them instead of the converted parameters. ◎ thinkWhy must the model be moved to the device before constructing the optimiser? register_buffer is for tensors that should follow the module (move with .to(device), appear in state_dict) but are not trained parameters. BatchNorm's running_mean and running_var are buffers, as is the attention mask in a transformer. model.parameters() returns all leaf tensors with requires_grad=True. model.named_parameters() returns the same with names, useful for excluding biases from weight decay or assigning different learning rates to different parts of the model. Calling model(x) invokes __call__, not forward directly. This is a classic misconception for new users. That wrapper runs forward hooks, then the computation, then backward hooks. The distinction matters when using hooks for instrumentation or gradient modification. torch.compile(model) (Only available in PyTorch 2.0+) traces the forward pass and emits optimised Triton/CUDA kernels via the backend. It fuses adjacent element-wise operations, reducing memory traffic. The first forward pass is slow (needs to run compilation) but subsequent ones are 10–30% faster on GPU, sometimes higher on inference-heavy workloads. noteThe PyTorch Parameter There are two things we care about in the nn.Parameter class: the values stored in .data and the .grad attribute. The backward pass accumulates gradients additively into .grad; the optimiser reads .grad to compute the update and writes the result back to .data.You can also use the parameters without gradients by setting requires_grad=False. This is useful for freezing layers, or for inference-only models, or non-trainable parameters. Training mode 1Sets every submodule to training mode. Two common layers change behaviour: Dropout applies random masking and rescales, and BatchNorm computes statistics from the current batch rather than its stored running estimates. 1model.train() # [1] 2for X_batch, y_batch in loader: 3 ... model.train() sets a flag on the module and all sub-modules. Two common layers read it during their forward pass.