PyTorch in 8 Slides
Theme
Palette
Download
Caption (tap to copy)
📖 Deep dive (full written explanation)
This is the code-heavy post, and the cover promises a complete, runnable training loop in about thirty lines. The value is integration: the previous posts explained tensors, autograd, and the step order in isolation; here they snap together into a working program you can copy, run, and modify.
The deliberate strategy is to build it in numbered stages — model, data, loss and optimizer, loop, evaluation, GPU — so each piece is digestible, then point at the one line beginners reliably get wrong.
Defining a model means subclassing nn.Module. Two methods matter: __init__, where you create and register the layers (calling super().__init__() first is mandatory, or registration breaks), and forward, which describes how data flows through them.
nn.Sequential is a convenience that chains layers in order, so you don't have to call each by hand in forward. This tiny network takes a 4-feature input, expands to 16 hidden units with a ReLU nonlinearity, and outputs 3 class scores (logits). Registering layers as attributes is what lets model.parameters() later find every weight for the optimizer.
Data handling in PyTorch flows through two abstractions. A Dataset knows how to fetch one sample by index; a DataLoader wraps a Dataset to yield shuffled, batched groups of samples and can parallelize loading across worker processes.
Here TensorDataset bundles features X and labels y so that indexing returns matching pairs. DataLoader then serves them in batches of 16, reshuffling each epoch so the model doesn't learn the data order. In real projects you'd subclass Dataset to load images or text from disk, but the loader interface stays identical — which is the point of the abstraction.
Two choices define the learning objective. The loss function measures how wrong the model's predictions are; CrossEntropyLoss is the standard for multi-class classification. The optimizer decides how to update weights from the gradients; Adam is a robust default that adapts the step size per parameter.
Note model.parameters() passed to the optimizer — this is how the optimizer knows which tensors to update, and it works because the layers were registered as attributes in __init__. The learning rate 1e-3 is a sensible Adam starting point; it's usually the first hyperparameter you'd tune.
This is the loop everything has been building toward, and it's only five meaningful lines per batch. zero_grad clears last step's gradients. model(xb) runs the forward pass. loss_fn compares predictions to targets. backward() computes gradients into every parameter's .grad. step() applies the update.
The outer loop over epochs repeats this over the whole dataset multiple times. Printing loss.item() each epoch lets you watch it (hopefully) decrease — your first sign that learning is happening. This rhythm is identical whether the model is this toy net or a billion-parameter transformer; only the layers and data change.
The cycle diagram captures the loop as the four-beat rhythm you'll repeat for the rest of your PyTorch life: forward to get predictions, compute the loss, backward to fill gradients, step to update (and zero for next time). Drawing it as a cycle emphasizes that it runs over and over, once per batch.
Internalizing this shape means you can read any PyTorch training script at a glance — you're just looking for these four beats and noting what's plugged into each slot.
train() and eval() are mode switches, and getting them wrong is a silent bug. Some layers behave differently between training and inference. Dropout randomly zeroes activations during training for regularization but must pass everything through at test time. BatchNorm uses batch statistics while training but switches to accumulated running statistics for inference.
model.train() and model.eval() flip every such layer in the model at once. Forgetting eval() at test time leaves dropout active, so your reported accuracy is artificially depressed and noisy. It's a bug that never crashes, which is exactly why it's so common.
Correct evaluation has two parts. model.eval() puts layers in inference mode as just discussed. torch.no_grad() tells autograd to stop building the graph, since you have no intention of calling backward during evaluation — this saves memory and speeds things up.
The accuracy computation is idiomatic: argmax(1) picks the highest-scoring class per row, compares to the true labels, and the boolean result averaged as a float gives the fraction correct. Wrapping evaluation in no_grad() is a habit worth forming immediately; without it you waste memory recording a graph you'll never use.
Moving to the GPU is a two-part discipline. First pick a device once, defaulting to CUDA if available and CPU otherwise, so the same script runs anywhere. Then send both the model and every batch of data to that device.
The reason you must move both is the device-match rule: an operation's tensors must share a device. Move the model but forget the data and you get a runtime error; move the data but forget the model, same thing. This .to(device) pattern is the standard idiom and scales unchanged from a laptop GPU to a multi-node cluster.
This snippet shows the minimal GPU changes layered onto the existing loop. Choose device with a one-line conditional. Move the model once before training. Then, inside the loop, move each batch's inputs and targets to the same device before using them.
That's the entire GPU story for a single-GPU program — three small additions. The discipline is consistency: anything that touches the model must be on the model's device. Most 'works on CPU, breaks on GPU' bugs are a forgotten .to(device) on one tensor.
This recap distills the post into a memorizable rhythm. zero_grad then forward then loss; backward then step; train() to train and eval() to test; no_grad() around evaluation; everything on one device. Five lines that, once internalized, let you write or read any PyTorch training code fluently.
If you remember nothing else from the day, remember this rhythm — it's the practical core that the conceptual posts exist to justify.
The CTA points to the final post. You can now write code that runs; the danger is code that runs but trains wrong. Post 5 catalogs the silent failures — the bugs with no error message — so you can recognize and fix them fast.