PyTorch in practice
By the end: You can write and debug a training pipeline.
- Tensors and devices
- Modules and the training loop
- Datasets and data loading
- Where the time actually goes
New to pytorch? PyTorch is the Python library almost everyone uses to build and train neural networks. It hands you two things you would otherwise have to write: arrays of numbers that can live on a GPU instead of the main processor, and automatic differentiation — it watches the arithmetic you perform and works out on its own how every weight should move to reduce the error. Layers, optimisers and data loading come with it. Why it matters
This is where the ideas become something that runs. PyTorch is what nearly all research and a good deal of production is written in, and the section covers the parts you touch daily: tensors and the hardware they sit on, modules, the training loop, and getting data in quickly enough that the GPU is not left idle. Expect the difficulties to be mechanical rather than conceptual, which is why shapes and devices get most of the attention.
The error you will actually see What it means
mat1 and mat2 shapes cannot be multiplied (32x768 and 512x256) inner dimensions disagree
Tensor for argument input is on cpu but expected on mps model moved, batch not (CUDA words it
as: all tensors on the same device)
expected target dtype to be Long or Byte, but got Float cross_entropy wants integer class labels
Trying to backward through the graph a second time backward called twice, or a tensor kept across steps
element 0 of tensors does not require grad you are inside no_grad, or the input was detached
Can't call numpy() on Tensor that requires grad add .detach()
CUDA out of memory batch too large, or you accumulated a tensor
that dragged its graph along with itTensors#
A tensor is an array that tracks operations so gradients can be computed, and that can live on a GPU. Most day-to-day friction is shapes and devices.
It carries five properties, and nearly every error is one of them being wrong. Print all five when something is confusing rather than guessing which one it is.
import torch
x = torch.randn(32, 768)
print(tuple(x.shape)) # (32, 768) what it holds
print(x.dtype) # torch.float32 how each number is stored
print(x.device) # cpu where the memory lives
print(x.requires_grad) # False whether autograd is watching
print(x.element_size() * x.nelement()) # 98304 bytes, = 32 * 768 * 4dtype bytes 1000x1000 Use it for
torch.float32 4 4.0 MB the default; everything, until memory hurts
torch.bfloat16 2 2.0 MB mixed-precision training; same exponent range as float32
torch.float16 2 2.0 MB inference and older GPUs; overflows above 65504
torch.int64 8 8.0 MB class labels and indices; PyTorch's default integer
torch.int32 4 4.0 MB indices, when int64 is wasting memory
torch.bool 1 1.0 MB masksThe conversion between numpy and torch has a trap in it that is worth meeting deliberately. torch.from_numpy shares the underlying memory, so editing the array changes the tensor; torch.tensor copies. One is fast and surprising, the other is safe and slower.
import numpy as np
a = np.array([1.0, 2.0, 3.0])
shared = torch.from_numpy(a) # shares memory with a
copied = torch.tensor(a) # copies
a[0] = 99.0
print(a) # [99. 2. 3.]
print(shared.numpy()) # [99. 2. 3.] followed the edit
print(copied.numpy()) # [ 1. 2. 3.] did not
g = torch.ones(3, requires_grad=True)
g.numpy() # RuntimeError: Can't call numpy() on Tensor that requires grad.
g.detach().numpy() # [1. 1. 1.]Devices#
Every tensor in an operation must live on the same device. Decide once at the top of the script and move everything to the same place, rather than sprinkling .cuda() through the code and discovering at step 300 that one buffer stayed behind.
device = ("cuda" if torch.cuda.is_available()
else "mps" if torch.backends.mps.is_available() # Apple silicon
else "cpu")
model = model.to(device) # in place for a Module
x = x.to(device) # NOT in place for a Tensor: returns a copy
for xb, yb in loader:
xb, yb = xb.to(device), yb.to(device) # the batch must be moved too, every stepx = torch.randn(4, 8).to("mps")
x @ torch.randn(8, 4)
# RuntimeError: Tensor for argument #2 'mat2' is on CPU, but expected it to be on GPU
# (while checking arguments for mm)Where the model and the data get separated
Almost always the model was moved in the setup code and the batch is moved nowhere, because the loader yields CPU tensors and nothing says otherwise. The second most common cause is a tensor you created inside forward, such as torch.arange for positions, which lands on the CPU unless you pass device=x.device. Writing that argument by reflex costs nothing and removes a whole category of bug.
A model#
import torch.nn as nn
class Classifier(nn.Module):
def __init__(self, in_dim, hidden, classes):
super().__init__()
self.net = nn.Sequential(
nn.Linear(in_dim, hidden),
nn.ReLU(),
nn.Dropout(0.1),
nn.Linear(hidden, classes),
)
def forward(self, x):
return self.net(x) # raw scores, loss applies softmaxThe training loop, line by line#
Six lines, and every one of them is load-bearing. Learn what each does and the loop stops being something you copy.
model.train() # 1
for xb, yb in train_loader: # 2
xb, yb = xb.to(device), yb.to(device) # 3
optimizer.zero_grad(set_to_none=True) # 4
logits = model(xb) # 5
loss = criterion(logits, yb) # 6
loss.backward() # 7
torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0) # 8
optimizer.step() # 91 model.train() switches dropout and batchnorm into training behaviour.
2 for xb, yb in ... the loader yields one batch; workers prepared it while the GPU worked.
3 .to(device) the loader always yields CPU tensors. This is the line people omit.
4 zero_grad gradients ACCUMULATE in PyTorch. Without this, batch 50 carries
the sum of the first 50 batches and the run diverges silently.
set_to_none=True frees the buffers rather than filling them with 0.
5 model(xb) forward pass. Every operation is recorded onto the autograd graph.
6 criterion(...) one scalar. Feed it raw logits: the loss applies softmax itself.
7 loss.backward() walks that graph in reverse, adding each parameter's gradient
into p.grad. Adds, not assigns, which is why line 4 exists.
8 clip_grad_norm_ rescales all gradients if their combined norm exceeds 1.0.
Optional for small models, standard for transformers.
9 optimizer.step() applies p -= lr * p.grad, or Adam's version of it.Why gradients accumulate at all
It looks like a design mistake until you need it. Accumulation is what lets you simulate a batch of 512 on hardware that fits 64: run eight forward-and-backward passes, then step once. Divide the loss by 8 so the gradient magnitude matches a real batch of 512. Gradient checkpointing, multi-task losses and most memory tricks lean on the same behaviour.
Train and eval mode#
Two calls matter here and both fail quietly. Dropout and batch normalisation behave differently during training and inference, so forgetting model.eval() before validation gives noisy, pessimistic numbers, and forgetting model.train() afterwards silently disables dropout for the rest of the run.
It is worth knowing exactly what changes, because the answer is narrower than most people assume. The flag does nothing to gradients, nothing to memory and nothing to speed. It sets a boolean on every submodule, and only two common layer types read it.
model = nn.Sequential(nn.Linear(10, 10), nn.Dropout(0.5), nn.BatchNorm1d(10))
x = torch.randn(8, 10)
model.train()
print([round(model(x)[0, 0].item(), 4) for _ in range(5)])
# [1.5063, -0.366, -0.4132, 1.4859, -0.4225] same input, five different answers
model.eval()
print([round(model(x)[0, 0].item(), 4) for _ in range(5)])
# [0.6156, 0.6156, 0.6156, 0.6156, 0.6156] deterministicDropout zeroes a random half of the activations during training, and to keep the average signal the same it scales the survivors up by two. In eval it does nothing at all, passing the input straight through. That scaling is why you must not compensate by hand.
d = nn.Dropout(0.5); t = torch.ones(100000)
d.train(); print(round(d(t).mean().item(), 4), round((d(t) == 0).float().mean().item(), 3))
d.eval(); print(round(d(t).mean().item(), 4), round((d(t) == 0).float().mean().item(), 3))
# train mean 0.9986 zeros 0.500 half deleted, the rest doubled
# eval mean 1.0000 zeros 0.000 untouchedBatch normalisation is the more damaging of the two to get wrong, because it does not merely add noise; in training mode it normalises using the statistics of whatever batch it is handed. Give it two examples and each is forced to the batch's own mean and range, which produces the same output regardless of the input.
torch.manual_seed(0)
bn = nn.BatchNorm1d(4)
big = torch.randn(256, 4) * 5 + 3 # mean about 3, variance about 25
bn.train()
for _ in range(50):
bn(big) # learns running statistics from real batches
print(bn.running_mean.round(decimals=2).tolist()) # [2.65, 3.80, 3.41, 2.72]
print(bn.running_var.round(decimals=1).tolist()) # [20.8, 29.3, 28.7, 25.7]
two = big[:2]
bn.train(); print(bn(two)[0].round(decimals=3).tolist())
bn.eval(); print(bn(two)[0].round(decimals=3).tolist())
# train mode, batch of 2: [-1.0, -1.0, 1.0, 1.0] every value is exactly +-1
# eval mode, batch of 2: [-1.08, -1.149, -0.291, -0.249] the real numbersbn.train(); bn(big[:1])
# ValueError: Expected more than 1 value per channel when training,
# got input size torch.Size([1, 4])Those running statistics are also why a batch-normalised model has state beyond its weights. They live in the state dict as buffers, they must be saved with the checkpoint, and a model evaluated before they have stabilised will score badly for reasons that have nothing to do with its weights.
What no_grad actually changes#
A separate mechanism from eval, controlling a separate thing, and the two are needed together rather than instead of one another. eval changes layer behaviour; no_grad stops the autograd graph being built. Inside it, operations run normally but nothing records what it did, so no intermediate activations are kept.
m = nn.Sequential(*[nn.Linear(512, 512) for _ in range(20)])
x = torch.randn(256, 512)
print(m(x).grad_fn is not None) # True the graph is being built
with torch.no_grad():
print(m(x).grad_fn is not None) # False nothing recorded
# measured on CPU, 20 layers, batch 256:
# forward with gradients tracked 4.4 ms live tensors 32.0 MB
# forward under no_grad 3.9 ms live tensors 22.1 MB
#
# difference 10.0 MB, against 20 x 256 x 512 x 4 bytes = 10.5 MB of
# stored activations. That is exactly what was not kept. model.eval() torch.no_grad()
changes dropout yes no
changes batchnorm yes no
saves memory no yes, a lot
makes it faster no slightly
prevents backward() no yes
scope until .train() the with-block onlymodel.train()
for batch in train_loader:
...
model.eval()
with torch.no_grad(): # no gradients needed, saves memory and time
for batch in val_loader:
...
model.train() # switch back, or dropout stays off for the rest of the runDataset and DataLoader#
A Dataset answers two questions: how many examples are there, and give me the i-th one. That is the entire interface, and writing one is three methods. The DataLoader wraps it and handles batching, shuffling, parallel loading and collation.
from torch.utils.data import Dataset, DataLoader
class Rows(Dataset):
def __init__(self, X, y):
self.X = torch.as_tensor(X, dtype=torch.float32)
self.y = torch.as_tensor(y, dtype=torch.long) # long: cross_entropy demands it
def __len__(self):
return len(self.y) # how many
def __getitem__(self, i):
return self.X[i], self.y[i] # the i-th, unbatched
ds = Rows(torch.randn(1000, 20), torch.randint(0, 3, (1000,)))
print(len(ds), tuple(ds[0][0].shape)) # 1000 (20,) one example, no batch dimensionloader = DataLoader(ds, batch_size=64, shuffle=True, num_workers=0)
xb, yb = next(iter(loader))
print(tuple(xb.shape), tuple(yb.shape), len(loader)) # (64, 20) (64,) 16
# 1000 / 64 = 15 full batches and a final one of 40.
print(len(DataLoader(ds, batch_size=64, drop_last=True))) # 15The default collate function stacks the examples, which requires them all to be the same shape. Variable-length data needs your own, and this is the one everybody eventually writes.
from torch.nn.utils.rnn import pad_sequence
def collate(batch):
lengths = torch.tensor([len(s) for s in batch])
return pad_sequence(batch, batch_first=True, padding_value=0), lengths
padded, lengths = collate([torch.arange(n) for n in (3, 7, 5)])
print(tuple(padded.shape), lengths.tolist()) # (3, 7) [3, 7, 5]
# [[0, 1, 2, 0, 0, 0, 0],
# [0, 1, 2, 3, 4, 5, 6],
# [0, 1, 2, 3, 4, 0, 0]]A DataLoader batches, shuffles and loads in parallel. It is frequently the bottleneck: a GPU sitting idle waiting for data is the most common performance problem in training, and it looks like the model being slow.
loader = DataLoader(
dataset,
batch_size=64,
shuffle=True, # training only. Never shuffle validation; it wastes time.
num_workers=4, # parallel loading, the usual fix for an idle GPU
pin_memory=True, # faster host to device transfer
persistent_workers=True, # do not respawn the workers every epoch
prefetch_factor=2, # batches queued per worker, ready before they are asked for
)num_workers is not free either. Each worker is a separate process holding its own copy of the dataset object, so a Dataset that loaded a large array in __init__ is now in memory several times over, and workers respawn every epoch unless persistent_workers is set, which on a small dataset can cost more than the loading saved. Raise it from zero while measuring throughput and stop where it stops helping, usually around the number of physical cores.
Almost everything here is PyTorch, which is the right default: it is what most research code and most production serving is written in. The exception worth knowing about is JAX, used heavily at some frontier labs and built around a different idea — you write pure functions and transform them, with jit for compilation, grad for derivatives, vmap for batching and pmap or shard_map for spreading work across devices. It is also the native stack for TPUs, Google's accelerators, where the unit of thinking is a mesh of chips and how your arrays are sharded across it. You do not need to learn it before an interview unless the job says so, but you should not be surprised by the word, and if you are applying somewhere whose published work is in JAX, read enough to follow their code.
Where the time goes#
- Data loading, usually. Check GPU utilisation: if it is low, the model is not the problem.
- Moving data between CPU and GPU. Do it once per batch, not per item.
- Small batches, which underuse the hardware.
- Full precision, where mixed precision would be roughly twice as fast at the same quality.
Where people go wrong
Calling .item() or .cpu() inside the training loop to log something. Each one forces synchronisation with the GPU and stalls the pipeline. Accumulate on the device and read once per epoch.
If this comes up in an interview
What does model.eval() do?
Switches layers that behave differently at inference, notably dropout and batch normalisation, into evaluation mode. Forgetting it makes validation numbers noisy and pessimistic, and forgetting to switch back disables dropout for the rest of training.
Why wrap validation in torch.no_grad()?
It stops the framework recording operations for gradient computation, which is unnecessary during evaluation. It saves memory and time, and without it a large validation pass can run out of memory.
Your GPU utilisation is 20 percent during training. What is wrong?
Almost certainly data loading, not the model. Raise num_workers, enable pin_memory, and check for expensive per-item preprocessing that should be precomputed or moved onto the GPU.
Build this: Write a full training loop by hand: dataset, loader, model, optimiser, train and eval modes, and per-epoch metrics. Then measure GPU utilisation and improve it.
Watch it explained
This is one section of PyTorch, 30 hours across 6 sections. Read it in full to get the whole topic in order.