The Mental Model of PyTorch: Build a Neural Network in 5 Phases

Yachika Anand

9 min read

·

Jul 6, 2026

--

1

Listen

Share

I stopped copy-pasting training loops and finally understood what every line does — here’s the whole picture, taught the way I wish someone had taught me.

Estimated read: ~11 minutes · Level: you know a little Python and NumPy, and you’ve seen a neural network but it still feels like magic. By the end, it won’t.

For the longest time, my PyTorch code worked and I had no idea why.

I’d paste in a training loop from a tutorial, swap the model, change the file path, and hope. Sometimes the accuracy went up. When it didn’t, I’d poke at random lines — zero_grad() here, .to(device) there — like a person slapping the side of a TV. It worked often enough that I never had to admit I didn't understand it.

Then I forced myself to answer one question for every single line: why is this here? This post is that answer, compressed. By the end you’ll have a mental model that makes every PyTorch script — yours or anyone else’s — readable at a glance.

The one picture to hold onto: a neural network is just a function with millions of adjustable knobs (the weights). “Training” means showing it examples and gently turning those knobs until its guesses get good. Everything else is machinery for turning them efficiently.

Hold that sentence. Every concept below is just a part of that machine.

We’ll build a real thing: a network that looks at a 28×28 photo of a piece of clothing and says “that’s a sneaker.” The dataset is Fashion-MNIST. The full, runnable, heavily-commented notebook is here — read this for the why, run that for the how:

▶️ Run the full notebook →

Press enter or click to view image in full size

The map: 5 phases

Almost every PyTorch project, no matter how fancy, is these same five moves in the same order:

Phase: What it does

🧱 1 · Foundations: the two ideas everything rests on

📦 2 · The Data: get examples ready

🏗️ 3 · The Model: build the network

🎛️ 4 · Training & Tuning: make it learn well

🏁 5 · Results: measure it, save it

Once you see the phases, unfamiliar code stops being a wall of syntax and becomes “oh, this is just Phase 2.” Let’s walk them.

🧱 Phase 1 · Foundations

Two ideas sit under everything in PyTorch. Get these and the rest clicks into place.

A tensor is just a smarter array

You already know NumPy arrays. A tensor is the same thing — a grid of numbers — with two superpowers: it can live on a GPU (for speed), and it can remember the operations performed on it. In PyTorch everything is a tensor: the images, the weights, even the gradients.

Autograd: PyTorch does the calculus for you

Here’s the part that feels like magic. To turn the knobs the right way, training needs derivatives: “if I nudge this one weight, does the error go up or down?” Working that out by hand for millions of weights would be hopeless.

So PyTorch quietly watches. The moment a tensor is marked requires_grad=True, PyTorch records every operation you perform on it, secretly building a graph of the computation. Then you call .backward() on the final error, and it walks that graph backwards, applies the chain rule automatically, and fills in the gradient for every weight.

You write the forward maths. Autograd works out the backward maths. That’s the whole trick.

🧠 Keep this sentence: the forward pass builds the graph, .backward() reads it. Training and prediction are really just "recorder on" versus "recorder off" (torch.no_grad()).

What’s actually inside a network

A single neuron is embarrassingly simple: it multiplies each input by a weight, sums them, adds a bias, and outputs the result. Those weights and biases are exactly the knobs from our one sentence.

A layer (nn.Linear) is just many neurons in parallel. Stack layers and you get a deep network.

But there’s a catch, and it’s the reason depth works at all: if every layer only did weights × inputs + bias, stacking them would be pointless — a chain of straight lines is still just a straight line. So after each layer we apply a non-linear activation (the workhorse is ReLU: max(0, x), which just zeroes out negatives). That bend is what lets the network trace the curvy patterns of real data.

🧠 In one line: a network is Linear → activation → Linear → activation → … → Linear. Weights are learned; activations are what make it more than a fancy straight line.

The training loop — the heartbeat of all deep learning

If you remember one thing from this entire post, remember this rhythm. Every neural network, in any framework, learns by repeating the same four beats — once per batch of data:

for epoch in range(epochs):                # one full pass over the data
    for features, labels in train_loader:  # one batch at a time
        outputs = model(features)          # 1. FORWARD — guess
        loss = criterion(outputs, labels)  # 2. MEASURE — how wrong?
        optimizer.zero_grad()              #    wipe last round's gradients
        loss.backward()                    # 3. BACKWARD — assign blame
        optimizer.step()                   # 4. UPDATE — nudge the knobs

That’s it. That’s training. Forward, measure, backward, update — thousands of times, and the guesses slowly get good.

The one chore everyone forgets is optimizer.zero_grad(). PyTorch accumulates gradients by default, so if you don't wipe them each round, this batch's gradients pile onto the last batch's and poison the update. The mantra: clear → compute → step.

Everything in the remaining phases just prepares the ingredients for this loop.

[IMAGE: a simple 4-arrow cycle diagram — Forward → Measure → Backward → Update — looping back]

📦 Phase 2 · The Data

A model is only as good as what you feed it. Preparing data is three deliberate moves — load, split, scale.

Load. Fashion-MNIST lives in a CSV: each row is one image (column 0 is the label 0–9, the other 784 columns are pixel brightnesses of a flattened 28×28 picture).

Split. We hold back a test set the model will never train on — so later we can tell whether it genuinely learned or just memorised. Think of it as an exam you don’t get to see beforehand.

Scale. Networks learn far more smoothly when inputs are small and on a similar range. For images, dividing by 255 squashes every pixel into 0–1. Done.

✋ The one rule to burn in: scale after splitting, never before — otherwise information from your test set leaks into training and your exam is rigged.

Dataset & DataLoader — the conveyor belt

We never shove all 60,000 images through the model at once. Instead we stream small batches. PyTorch splits this across two cooperating classes, and there’s a clean way to picture it — a warehouse:

• Dataset is the shelf clerk: it knows how to fetch one item when asked (you write three tiny methods — __init__, __len__, __getitem__).
• DataLoader is the conveyor belt: it repeatedly asks the clerk for items, shuffles them, and packs them into batches. On your side it's just for batch in loader:.

⚠️ The dtype trap that bites everyone: features must be float32, but labels must be long (plain integers) — the loss function insists on integer class indices. Set both once, in the Dataset.

🏗️ Phase 3 · The Model

Now we assemble the network from Lego blocks (torch.nn gives you pre-built layers so you never write matrix calculus by hand). Every model you'll ever write follows the same skeleton: inherit from nn.Module, define your layers, write a forward() that says how data flows. Then you run it as model(x) — never model.forward(x) — because the model(x) form runs important bookkeeping around your forward pass.

The interesting part of Phase 3 isn’t the layers themselves — it’s the enemy they’re fighting: overfitting.

Overfitting is when a network is powerful enough to just memorise the training images: top marks on the practice exam, terrible marks on the real one.

We want general patterns, not memorisation. There are seven classic defences; our model uses three, and knowing why is half the battle:

• Dropout — randomly switch off some neurons during training, so the network can’t lean on any single one and builds robust, redundant features.
• Batch Normalization — re-centre the numbers flowing between layers so training stays stable and fast.
• L2 / weight decay — gently penalise large weights so the model stays humble.

So each hidden layer becomes a little recipe: Linear → BatchNorm → ReLU → Dropout. That's the whole model.

🔌 One switch to remember: Dropout and BatchNorm behave differently during training vs testing. model.train() and model.eval() flip all of them at once — which is why evaluation always starts with model.eval().

🎛️ Phase 4 · Training & Tuning

Two objects make learning actually happen, and it helps to cast them as a duo:

• The loss (criterion) is the critic. After each guess it scores how wrong the model was — one number, lower is better. For our multi-class problem that's CrossEntropyLoss (feed it raw scores, not softmax — it applies softmax internally).
• The optimizer is the coach. It takes autograd’s gradients and decides how to adjust each weight — how big a step, how much momentum. Adam is a strong default.

The fast lane: GPU

A network is a giant pile of small multiplications. A CPU does them in a short queue; a GPU does thousands side by side — hours become minutes. There’s exactly one golden rule: a tensor and the model touching it must live on the same device. So you pick the device once, move the model there, and move each batch there. That’s all .to(device) is doing everywhere you see it.

Stop guessing hyperparameters: Optuna

Some numbers the network learns on its own (the weights). Others you must choose before training: how many layers? how wide? what learning rate? These are hyperparameters, and there’s no gradient to guide them.

Instead of trying combinations by hand, Optuna searches intelligently — it learns which regions of the space look promising and concentrates there. You write an objective function (build → train → return the score) and let Optuna run it dozens of times, hunting for the best setup.

[IMAGE: optional — an Optuna optimization-history plot showing accuracy climbing across trials]

🏁 Phase 5 · Results

Training accuracy flatters you — the model has seen that data. The honest question is how it does on the held-out test set. Four habits keep the measurement trustworthy:

1. model.eval() — flip off Dropout/BatchNorm's training behaviour.
2. with torch.no_grad(): — recorder off; we're answering, not learning.
3. Turn scores into answers — torch.max picks the highest-scoring class.
4. Tally — correct predictions ÷ total = accuracy.

Then save the model so a runtime restart doesn’t erase all that compute — the community-standard way is to save the state_dict (a dictionary of the learned weights), not the whole object.

🎓 And that closes the loop: data → batches → model → loss → optimizer → fight overfitting → tune → measure → save. Read a PyTorch script with these five phases in mind and it stops being a mystery.

So, what did we actually build?

A network that turns knobs until it can recognise clothing — and, more importantly, a mental model you can reuse for any PyTorch project. Next time you meet an unfamiliar script, you won’t see a wall of syntax. You’ll see: “Phase 2 here, Phase 4 there, and there’s the four-beat loop.”

That shift — from slapping the side of the TV to actually reading the circuit — is the whole point.

▶️ See every line run, with a comment on each step:

👉 The full annotated notebook →

📚 The whole series & code on GitHub →

Press enter or click to view image in full size

📬 This is Part 1 of a 5-part series: The Mental Model of PyTorch

We start with the foundations and build up to how ChatGPT-style models work — one intuitive notebook at a time.

1. The Mental Model of PyTorch — Build a Neural Network in 5 Phases ← you are here
2. Why Flattening an Image Throws Away the Picture — CNNs
3. Don’t Train From Scratch: Standing on ImageNet’s Shoulders — transfer learning
4. Giving a Network Memory — RNNs and a Q&A bot
5. How Your Phone Keyboard Works — LSTMs, language modeling, and the door to Transformers

Next up → Part 2: we keep the picture 2-D and watch accuracy jump.

If this made one concept click, a clap or a comment tells me which one — and helps the next confused-me find it. What still feels like magic? Ask below.