PyTorch · Lesson 6

Build & Train a Real Net

Everything from Lessons 1–5 in one program: a 2-layer neural network that classifies points, trained on your Apple GPU, scored on a held-out set. We go from 45% to 100% accuracy on a problem no straight line can solve.

🎯 Mission: read & write real PyTorch
The one idea

A "deep" net is just linear layers stacked with a nonlinearity between them. The nonlinearity is the whole point: without it, ten stacked layers collapse into one line. Add it, and the network can carve curved decision boundaries. The loop, the loss, the optimizer, the data pipeline — all the machinery you've already built — stays exactly the same. Only the model grows.

01 A problem no line can solve

Our data: an inner disk (class 0) sitting inside an outer ring (class 1). No straight line separates a disk from the ring around it — so nn.Linear alone (Lesson 4) is hopeless here. We need a real network.

This is the payoff lesson: it combines the tensors, autograd, loop, modules/optimizers, and data pipeline you already know into the standard shape of real PyTorch model code.

02 The nonlinearity is the magic — ReLU

Stack two linear layers with nothing between them and the math collapses: W₂(W₁x) = (W₂W₁)x is still just one linear layer. To get more power you insert a simple nonlinear function between layers. The workhorse is ReLU: relu(z) = max(0, z) — keep positives, zero out negatives.

💡 Why this unlocks everything That tiny "bend" at zero lets the network combine many simple linear pieces into arbitrarily curved boundaries (this is, loosely, the universal approximation idea). No nonlinearity → only straight-line decisions, forever. It's the single most important reason deep nets work.

Source: torch.nn.ReLU · Build the model tutorial.

03 The model: 2 → 16 → ReLU → 2

nn.Sequential chains layers — the output of each feeds the next. Two inputs (the x,y coordinates), a 16-unit hidden layer with ReLU, then 2 outputs (one score per class):

model = nn.Sequential(
    nn.Linear(2, 16),    # 2 features → 16 hidden units
    nn.ReLU(),           # the bend that makes it nonlinear
    nn.Linear(16, 2),    # 16 → 2 logits (one per class)
).to(device)

print(sum(p.numel() for p in model.parameters()))   # 82 params
input (x, y) 2 16 hidden · ReLU 16 2 logits 2 argmax → class Linear(2,16) ReLU → Linear(16,2)
An MLP (multilayer perceptron). 82 params = (2×16+16) + (16×2+2). The hidden layer + ReLU is what bends the decision boundary into a circle.

04 Classification loss: CrossEntropyLoss

For predicting a number we used MSE. For predicting a class, the standard loss is cross-entropy. It takes the raw output scores (logits) and the integer class labels:

loss_fn = nn.CrossEntropyLoss()
loss = loss_fn(logits, labels)    # logits: (N, 2) floats ; labels: (N,) ints (0 or 1)
⚠ Two gotchas that bite everyone (1) No softmax in your model. CrossEntropyLoss applies log-softmax internally — feed it the raw logits. Adding your own softmax first is a classic bug. (2) Labels are integer class indices of dtype long with shape (N,) — not one-hot vectors, not floats. The output layer has one unit per class.

Source: torch.nn.CrossEntropyLoss.

05 On the GPU, with train/eval modes

Two idioms appear in every real training script. First, move the model and every batch to the device. Second, flip the model between train() and eval() modes, and wrap evaluation in no_grad().

model.train()                 # training mode
for xb, yb in train_loader:
    xb, yb = xb.to(device), yb.to(device)   # batch → GPU (must match model)
    ...

@torch.no_grad()              # no graph needed for evaluation → faster, less memory
def accuracy(loader):
    model.eval()              # eval mode
    ...
    pred = model(xb).argmax(dim=1)   # highest-logit class
    correct += (pred == yb).sum().item()
💡 Why train() / eval() matter Some layers (dropout, batch-norm) behave differently when training vs. predicting. Our tiny net has none, so it doesn't change the numbers — but writing model.train() / model.eval() is a habit worth building now, because forgetting it on a real model silently wrecks your accuracy. Accuracy itself = "what fraction did argmax get right."

06 Train it — 45% → 100%

The loop is the same five moves over batches, plus an accuracy check each pass. Real output, on device = mps:

device = mps
model params: 82
val accuracy before training: 45.0%
epoch  0 | train loss 0.4165 | val acc 90.6%
epoch 10 | train loss 0.0085 | val acc 100.0%
epoch 20 | train loss 0.0010 | val acc 100.0%
epoch 30 | train loss 0.0011 | val acc 100.0%
45% → 100%val accuracy
82parameters
mpstrained on Apple GPU

Before training the net guesses at chance (~45–50%). Within a few epochs the hidden layer has bent a circular boundary around the disk and it nails the held-out set. An untrained random net scoring ~50% then climbing is exactly the signal that learning is happening.

Run it: ~/projects/learn/public/courses/pytorch/practice/.venv/bin/python ~/projects/learn/public/courses/pytorch/practice/lesson06_real_net.py

07 The whole idiomatic stack, in one place

This program is the shape of real PyTorch. Every block traces to a lesson:

Data → tensors

Features and integer labels as tensors (L1), wrapped in a Dataset + DataLoader with a train/val split (L5).

Model

nn.Sequential of Linear + ReLU layers (L4), moved .to(device) (L1).

Loss + optimizer

CrossEntropyLoss + Adam(model.parameters()) (L4).

Loop

zero_grad → forward → loss → backward → step over batches (L2–L5), with train()/eval().

💡 You can now read real model code Open almost any PyTorch training script and you'll find these exact pieces. Bigger models swap in fancier layers and losses, but the skeleton is what you just wrote. Lesson 7 closes the loop: saving/loading the trained model and reading a real-world training script top to bottom.

08 Check yourself

Why put a ReLU between two Linear layers?
What should you feed CrossEntropyLoss?
You moved the model with .to("mps") but pass a CPU batch. What happens?
How do you turn a model's 2 output logits into a predicted class?
Why wrap the accuracy/eval function in @torch.no_grad() and model.eval()?

09 Flashcards

Q: What makes a network "deep" / nonlinear?
A: A nonlinearity (e.g. ReLU = max(0,z)) between linear layers. Without it, any stack of Linear layers reduces to one linear layer — only straight-line decisions.
Q: MSE vs CrossEntropy — when each?
A: MSELoss for predicting a continuous number (regression). CrossEntropyLoss for predicting a class (classification): feed raw logits + integer labels.
Q: Two things CrossEntropyLoss expects?
A: (1) Raw logits shape (N, C) — no softmax in the model (it's applied internally). (2) Integer class indices shape (N,), dtype long — not one-hot.
Q: Two "device" rules for training on GPU?
A: Move the model with model.to(device) and move every batch with xb.to(device). Model and data must share a device or you get a runtime error.
Q: What do model.train() and model.eval() do?
A: Switch the model between training and inference behavior for layers like dropout/batch-norm. Always set eval() (and no_grad()) when measuring/predicting.
Q: How do you get a predicted class and accuracy from logits?
A: Predicted class = logits.argmax(dim=1). Accuracy = fraction where that equals the label: (pred == y).float().mean().
👩‍🏫 I'm your teacher — ask me anything

Try: "remove the ReLU and show accuracy stays at chance," "add a second hidden layer," or "plot the decision boundary." One lesson left: how to save this trained model, load it back, and read a real-world training script end to end. Say "go" for Lesson 7.