Your first neural network
in PyTorch

The setup

The delivery problem

Imagine Thirty Minute Co., a delivery company whose entire brand is a promise: your order arrives in half an hour or it's free. A delivery boy, working for them, was late three times last month. His manager has made it clear that a fourth time is a different kind of conversation.

An order comes in. Seven miles.

Does he take it?

He could guess. Or he could look at the delivery times already on record, notice that they fall in a fairly obvious pattern, and work out what that pattern says about seven miles. Noticing such patterns in numbers you already have, and using them to answer a question about numbers you don't, is the entire job of the model we're about to build.

It will be six lines of code. Here is the data.

import torch
import torch.nn as nn
import torch.optim as optim
torch.manual_seed(42)

distances = torch.tensor([[1.0], [2.0], [3.0], [4.0]], dtype=torch.float32)
times     = torch.tensor([[6.96], [12.11], [16.77], [22.21]], dtype=torch.float32)
distances.shape, times.shape

# (torch.Size([4, 1]), torch.Size([4, 1]))

Two things to notice about the shapes before we move on.

Each tensor is [4, 1], not [4]. Four rows, one column. The row count is how many examples you have; the column count is how many numbers describe each example. Here, one delivery is described by one distance number. If you later wanted to also feed in the time of day and the weather, that second dimension would become 3, and almost nothing else in the code would change.

Roughly five minutes of travel per mile, plus a couple of minutes of fixed overhead — parking, finding the door, waiting for someone to answer. Hold that thought. We're going to watch a neural network rediscover it.

Section 00

The 6 Stage ML Pipeline

Every model you will ever train moves through the same six stages.

1. Data ingestion. Get the raw records and into a form PyTorch can read.
2. Data preparation. Clean, transform, reshape. Time data could be messy — "22 minutes" as free text sitting next to a numeric 10, negative durations, missing values, records implying someone biked at 200 mph. Cleaning that up is the work of this step.
3. Model building. Choose an architecture. How many neurons and arranged how.
4. Training. Show the model the data repeatedly and let it adjust itself.
5. Evaluation. Check the model on data it has never seen.
6. Inference. Use the trained model to answer the actual question.

Given the data is small, we will train and evaluate on 100% of the data.

Section one · The minimal model — 01

Building a neuron

An artificial neuron is modeled after a neuron in the human brain. It accepts multiple inputs, performs a computation on them, and generates a single output based on the result.

What it computes is this:

• \(x\) — the input. The distance of the delivery, in miles. This is what you know.
• \(\hat{y}\) — the prediction. The estimated delivery time, in minutes. The hat marks it as the model's guess
• \(w\) — the weight. The slope of the line. How many extra minutes each extra mile costs you. This is one of the two numbers the neuron has to find.
• \(b\) — the bias. The intercept. The fixed overhead paid on every delivery regardless of distance. This is the other one.

\(w\) and \(b\) are the parameters. "Learning", in machine learning, means searching for the values of these parameters that will make the line pass through as close as possible to the points.

Let's do that search by hand first, so you know what the machine is doing when it does it for you.

Guess one: w = 1, b = 10. A four-mile delivery predicts 1 × 4 + 10 = 14 minutes. The real answer is 22.21. We're eight minutes short, and we're short at every distance — the line is too flat. Push the slope up.

Guess two: w = 5, b = 2. Four miles predicts 5 × 4 + 2 = 22 minutes. Real answer 22.21. We're off by two-tenths of a minute.

Two guesses got us close. That was intuition — look at the error, decide which direction to move, move. A neural network does exactly the same thing, with one difference: it computes the direction using calculus instead of eyeballing it.

Section one · The minimal model — 02

Defining the model, loss function & optimizer

model = nn.Sequential(nn.Linear(1, 1))
loss_function = nn.MSELoss()
optimizer = optim.SGD(model.parameters(), lr=0.01)

Model

• nn.Linear(in_features, out_features) — a fully connected layer. It holds a weight matrix and a bias vector, and computes \(\hat{y} = Wx + b\) for every row you pass it. nn.Linear(1, 1) means one input feature, one output value. It creates and owns the \(w\) and \(b\) we were guessing at, initialising both randomly.
• nn.Sequential(*layers) — a container that chains layers together and runs them in order, like an assembly line. Data enters the first layer, its output feeds the second, and so on. With a single layer it does nothing interesting; in section two we'll drop more layers into it and it will start earning its place.

Loss

• nn.MSELoss() — mean squared error. For every data point, take the gap between prediction and reality, square it, then average across all points. It is the scorekeeper that summarizes how wrong the model currently is.
  $$\text{MSE} = \frac{1}{n}\sum_{i=1}^{n}(\hat{y}_i - y_i)^2$$
  Two consequences of that square are worth understanding. It makes every error positive, so a five-minute overestimate and a five-minute underestimate don't cancel out. And it punishes large errors disproportionately — being off by 10 minutes is a hundred times worse than being off by 1, not ten times worse.

Optimizer

• optim.SGD(params, lr) — stochastic gradient descent. Given a direction to move each parameter, it moves them. lr is the learning rate: how big a step to take. It is 0.01 here because I tried it and it worked, which is roughly how learning rates get chosen most of the time. Section 09 is about what happens when the value you picked stops working, and why the real problem was somewhere else entirely.
  Previously

  Gradient Descent

  Build a classifier by hand from weights, biases, sigmoid, loss, and gradient descent.

  Read article →
• model.parameters() — the handle that connects the optimizer to the model. It returns an iterator over every learnable tensor the model owns, and hands the actual tensor objects over so the optimizer can modify them in place. Change them here, and the model changes.
  Previously

  Backpropagation

  Learn how gradients flow through a neural network using the chain rule, then build a tiny autograd engine from scratch.

  Read article →

Let's look at what those parameters start out as:

list(model.parameters())
# [Parameter containing:
#  tensor([[0.7645]], requires_grad=True),
#  Parameter containing:
#  tensor([0.8300], requires_grad=True)]

w = 0.7645, b = 0.8300. Random. Just as we started somewhere with w = 1, b = 10 and adjusted from there, the model needs an initial guess to learn and adjust from.

requires_grad=True is the flag that matters during training. It tells PyTorch: track every operation involving this tensor, because I will want to know how the final loss responds to changes in it. That tracking is what makes loss.backward() to compute gradients.

Section one · The minimal model — 03

What an untrained model predicts

The model is already fully functional. It just isn't any good yet. Let's ask it something.

with torch.no_grad():
    test_distance = torch.tensor([[25.0]], dtype=torch.float32)
    predicted_time = model(test_distance)
    print(f'PREDICTED TIME : {predicted_time.item(): .2f} minutes')

# PREDICTED TIME :  19.94 minutes

A twenty-five mile bike delivery in nineteen minutes and fifty-six seconds. That's a sustained 75 mph on a bicycle.

It's worth noting that the model didn't error. It didn't warn you. It returned a confident, well-formatted, completely absurd number. 0.7645 × 25 + 0.83 = 19.94 — the arithmetic is flawless. The arithmetic is always flawless. A model has no idea what a bicycle is, and it will never tell you that its answer is nonsense.

Section one · The minimal model — 04

The training loop in 6 lines

Everything so far has been setup. This is the part that learns.

for epoch in range(100):
    optimizer.zero_grad()
    outputs = model(distances)
    loss = loss_function(outputs, times)
    loss.backward()
    optimizer.step()

• for epoch in range(100) — an epoch is one full pass over the dataset. Here we make a hundred of them, so the model sees all four deliveries a hundred times over.
• optimizer.zero_grad() — wipe the gradients from the previous iteration. PyTorch accumulates gradients by default rather than replacing them. Forget this line will make every update faster and unstable. Your loss will do something strange, and nothing will tell you why. This is the single most common bug in beginner PyTorch code.
• outputs = model(distances) — the forward pass. All four distances go in, four predictions come out.
• loss = loss_function(outputs, times) — one number summarising how wrong we currently are, via mean squared error.
• loss.backward() — the backward pass. This is where requires_grad=True pays off. PyTorch walks back through every operation that produced the loss and computes, for each parameter, the partial derivative of the loss with respect to that parameter. In plain terms: if I nudge this number up slightly, does the error get better or worse, and by how much? The answers get stashed in w.grad and b.grad. .
  
  → Read more: how gradients are computed
• optimizer.step() — Before this step, the parameters are not updated. This step allows them to update. Each parameter takes a step in the direction that reduces the loss, scaled by the learning rate:
  $$w \leftarrow w - \eta \frac{\partial L}{\partial w}$$
  The minus sign is the whole trick. The gradient points uphill, toward more error. We go the other way.

Here's what a hundred of those looks like:

EPOCH : 1   | LOSS : 161.4454 | WEIGHTS : [[1.4600073]], BIAS: [1.0654308]
EPOCH : 2   | LOSS : 112.0370 | WEIGHTS : [[2.0393846]], BIAS: [1.2613717]
EPOCH : 3   | LOSS :  77.7535 | WEIGHTS : [[2.5220582]], BIAS: [1.4244250]
...
EPOCH : 98  | LOSS :   0.0359 | WEIGHTS : [[4.9558935]], BIAS: [2.1602232]
EPOCH : 99  | LOSS :   0.0359 | WEIGHTS : [[4.9561480]], BIAS: [2.1594741]
EPOCH : 100 | LOSS :   0.0358 | WEIGHTS : [[4.9564023]], BIAS: [2.1587272]

Loss from 161.45 to 0.036 — a factor of about 4,500 — in a hundred passes over four data points.

Note the pace. The first few epochs cut the loss by more than three quarters. The last few barely move it. That's not the optimizer getting tired; it's the gradient itself shrinking as we approach the bottom of the bowl. Near the minimum the surface is nearly flat, so the steps get small automatically.

(Small detail if you're reading the numbers closely: the loss printed on a given line is computed with the parameters from before that epoch's step, while the weights printed are from after it. Epoch 1's loss of 161.4454 is the error of the original random w = 0.7645, b = 0.8300.)

And now the payoff. The model landed on:

Our hand-guess a few sections ago was w = 5, b = 2. Gradient descent, working entirely from the arithmetic of the errors, arrived at the same answer we reached by squinting at a chart — and it read the underlying physics correctly on the way. Just under five minutes per mile of travel. About two minutes and ten seconds of fixed overhead per delivery, regardless of distance.

Nobody told it that deliveries have overhead. It found the overhead because the overhead was in the data.

Section one · The minimal model — 05

Watching it learn

Let's watch the model learn, epoch by epoch.

The trick is to record a snapshot every epoch while training — the loss, the current weight and bias, and the full set of predictions:

history = {'epoch': [], 'loss': [], 'predictions': [], 'weight': [], 'bias': []}

for epoch in range(100):
    optimizer.zero_grad()
    outputs = model(distances)
    loss = loss_function(outputs, times)
    loss.backward()
    optimizer.step()

    history['epoch'].append(epoch + 1)
    history['loss'].append(loss.item())
    history['predictions'].append(outputs.detach().numpy().flatten())
    history['weight'].append(list(model.parameters())[0].detach().numpy().flatten()[0])
    history['bias'].append(list(model.parameters())[1].detach().numpy().flatten()[0])

.detach() is doing quiet but essential work here. It returns a copy of the tensor cut loose from the gradient graph.

Watch the middle panel especially. The weight and bias don't move in lockstep — the weight does most of the early work while the bias trails, then the bias makes a small correction at the end. The optimizer is adjusting both simultaneously, but the loss is far more sensitive to the slope than to the intercept, so the slope moves first.

Section two · Making it non-linear — 06

When the line breaks

The company expands. Bikes for short hops, cars for anything further out. Suddenly you have nineteen more miles of data.

# Combined dataset: bikes for short distances, cars for longer ones
distances = torch.tensor([
    [1.0], [1.5], [2.0], ..., [19.0], [19.5], [20.0]
], dtype=torch.float32)   # 39 points, 0.5 mile spacing

times = torch.tensor([
    [6.96], [9.67], [12.11], ..., [90.73], [90.39], [92.98]
], dtype=torch.float32)

Plot it and the problem is immediate. The first five miles climb at the old bike rate, about five minutes a mile. Then the curve steepens sharply — those are the miles just past bike range, through the densest city traffic. And then, somewhere around twelve miles, it flattens out. Twelve miles takes 78 minutes; twenty miles takes 93. Eight extra miles for fifteen extra minutes.

That's the highway. Once a car gets clear of the city it moves fast, and each additional mile costs far less than the ones before it.

So let's train the same linear model on it and watch what happens. From here on the training loop lives inside a reusable train_model() function — same six lines, plus the history snapshots.

model = nn.Sequential(nn.Linear(1, 1))
loss_function = nn.MSELoss()
optimizer = optim.SGD(model.parameters(), lr=0.001)

history = train_model(model, loss_function, optimizer, distances, times, 200)

The model isn't broken. It converged. It found the optimal parameters for its architecture. The architecture is what's wrong.

A linear model makes exactly one assumption, and it makes it inescapably: every mile costs the same amount of time. That assumption was true when everything was a bike ride across a few miles of the same neighbourhood. It is flatly false the moment your data spans bikes, city traffic and highways. No amount of training fixes an assumption. You have to change the model.

Section two · Making it non-linear — 07

ReLU, one bend at a time

The obvious move is more neurons. One neuron learns one relationship, so several should learn several. Let's think about whether that's true.

Put two neurons in a hidden layer. Each takes the distance and computes its own weighted sum — call them \(z_1\) and \(z_2\). Then an output neuron combines them into a single prediction.

It's a straight line. Different coefficients, same shape.

This isn't a quirk of having two neurons. A linear function of a linear function is linear, always, and it doesn't matter how many layers you stack or how wide you make them. A hundred-layer network of pure nn.Linear layers is mathematically identical to a single nn.Linear layer. Bigger, slower, and exactly as incapable of drawing a curve.

You need something between the layers that isn't linear. That's an activation function, applied element-wise to each neuron's output before it goes anywhere.

The one to know is ReLU — rectified linear unit — and it is almost insultingly simple:

Negative input, output zero. Positive input, pass it through untouched. And that's it.

Here's how something that simple buys you curves. A neuron computes w × distance + b, which is a straight line. Run it through ReLU and the part of that line below zero gets flattened to zero. So you get a flat stretch, then a bend, then the line continuing at its original slope.

And the bend is not in an arbitrary place. It sits exactly where w × distance + b = 0:

The neuron's bias controls where it starts paying attention. That's a genuinely important idea. Each hidden neuron gets to choose a distance at which it switches on, and beyond that point it contributes at its own slope. One neuron gives you one bend. Three neurons give you three bends, at three different distances, with three different slopes — and the output layer gets to weight and sum them into any piecewise-linear shape it likes.

In code:

model = nn.Sequential(
    nn.Linear(1, 3),   # one input  → three hidden neurons
    nn.ReLU(),         # ← the entire difference
    nn.Linear(3, 1)    # three hidden → one output
)

nn.Sequential runs these in order: linear, then ReLU, then linear. We've gone from 2 parameters to 10 — three weights and three biases in the hidden layer, three weights and one bias in the output layer. And for the first time, the function this model can represent is not a straight line.

Section two · Making it non-linear — 08

How many neurons?

So: how many bends do we need? Let's try a few.

torch.manual_seed(27)
model = nn.Sequential(nn.Linear(1, 3), nn.ReLU(), nn.Linear(3, 1))
optimizer = optim.SGD(model.parameters(), lr=0.01)
history = train_model(model, loss_function, optimizer, distances, times, 3000)

# ... and the same with 2 neurons (lr=0.003, 10,000 epochs)
# ... and with 5 neurons (lr=0.001, 20,000 epochs)

Here's what happened to each one:

Look at that table for a moment, because it is strange in two separate ways.

First, the collapse. The three-neuron run didn't diverge — the loss didn't explode. It settled, perfectly stably, on the single most useless model possible: predict the average, always. What happened is a failure mode called dying ReLU. Early in training, the large raw inputs (up to 20) and large targets (up to 93) produce large gradients, and a few big steps pushed all three hidden neurons' outputs negative across the entire input range. ReLU turns a negative input into zero — and a neuron that outputs zero everywhere also has zero gradient everywhere, so no future update can revive it. All three neurons died within the first few epochs, the network reduced to its output bias, and the best a lone bias can do is the mean.

Second, the compensation. The runs that did work only worked because the learning rate was cut to a crawl — three to ten times smaller — and the epoch count ballooned to ten and twenty thousand. A larger model needing a smaller learning rate and vastly more epochs is backwards. Every time a run collapsed I turned the learning rate down and the epochs up, and every time that helped, so I kept doing it. It felt like tuning. It was actually a symptom I was treating instead of diagnosing.

Section two · Making it non-linear — 09

The fix wasn't more neurons

Go back to the animation in section 05 — the panel where the weight and bias moved at completely different rates. The weight did all the early work; the bias barely budged and then made a small late correction. Here is where that comes back.

Look at the actual numbers going into this model. Distance runs from 1 to 20. Time runs from 7 to 93. That's roughly a five-fold difference in scale between the input and the target — and inside the network, the gradient with respect to a weight is proportional to the input value that weight is multiplied by.

Which means gradients arriving from the 20-mile end of the dataset are about twenty times larger than gradients from the 1-mile end. A single learning rate has to work for both. Pick one large enough to make progress on the small-input examples and the large-input examples produce updates that overshoot wildly — on this data, wildly enough to kill every ReLU in the first few epochs. Pick one small enough to be safe at the top end and the bottom end crawls.

That's the trap. There is no good learning rate, so I kept picking worse ones and compensating with epochs.

The fix is not in the model at all. It's stage 2 of the pipeline — data preparation.

distances_norm = (distances - distances.mean()) / distances.std()
times_norm     = (times - times.mean()) / times.std()

This is standardisation, or z-scoring. Subtract the mean, divide by the standard deviation. Both tensors now have mean 0 and standard deviation 1. The delivery data hasn't changed shape — the curve is still exactly the same curve — it has just been recentred and rescaled so that no direction in parameter space is intrinsically steeper than any other.

Same three-neuron model, same seed, same learning rate as the run that collapsed:

torch.manual_seed(27)
model = nn.Sequential(nn.Linear(1, 3), nn.ReLU(), nn.Linear(3, 1))
loss_function = nn.MSELoss()
optimizer = optim.SGD(model.parameters(), lr=0.01)

history = train_model(model, loss_function, optimizer,
                      distances_norm, times_norm, 3000)

Three lines of arithmetic did what a week of hyperparameter fiddling wouldn't.

This is the lesson I'd most like to leave you with from this post, and it's the one that generalises furthest beyond it. When training misbehaves, the instinct is to reach for the model — more layers, more neurons, a different optimizer, a learning rate schedule. Check the data first. Enormous amounts of deep learning practice, from feature scaling through batch normalisation and layer normalisation, is fundamentally about keeping numbers inside a range where gradient descent behaves itself.

with torch.no_grad():
    test_distance = torch.tensor([[5.1]], dtype=torch.float32)
    test_distance_norm = (test_distance - distances.mean()) / distances.std()
    predictions = model(test_distance_norm)
    print(f'FINAL PREDICTION : {(predictions * times.std() + times.mean()).item():.2f} minutes')

# FINAL PREDICTION : 38.69 minutes

Note carefully which statistics are used. The input is normalised with the distance mean and standard deviation; the output is de-normalised with the time ones. Mixing those up is a silent, extremely annoying bug — you'll get numbers, they'll be plausible-looking, and they'll be wrong.

Is 38.69 any good? We never trained on 5.1 miles, so there's no ground truth. But the neighbouring points are: 5.0 miles took 37.15 minutes and 5.5 miles took 42.35. A prediction of 38.69 sits sensibly between them, close to the 5.0 end, which is where 5.1 miles belongs.

The model has learned the shape of the relationship, not just the points it was shown. That's the entire objective.

Two important caveats, so nobody walks away with the wrong idea:

• Compute the statistics on training data only. Here we normalised everything at once because everything is training data. In a real project, mean() and std() come from the training split and get applied unchanged to validation and test data. Otherwise information about the test set leaks into training and your evaluation numbers become fiction.
• Save the statistics with the model. A trained network without the two numbers used to scale its inputs is useless — you cannot feed it anything and you cannot interpret what comes out. They're part of the model. Store them together.

Section two · Making it non-linear — 10

Looking inside the network

One last thing. We've been treating the network as a box that turns distances into times. It isn't a box — it's ten numbers, and we can look at every one of them.

# inside train_model, once per snapshot:
history['params'].append([p.detach().numpy().copy() for p in model.parameters()])

Recording the full weight matrices every epoch gives us the network's actual anatomy over time.

This is worth a minute of your time. At initialisation the three hidden neurons are interchangeable — three random lines with no assigned role. Nothing in the code tells neuron 1 to handle short distances and neuron 3 to handle the highway stretch. That specialisation is emergent. It happens because it's the arrangement that minimises the loss, and gradient descent finds it.

Scale this up: more neurons, more layers, images instead of distances. The mechanism does not change. It's this, repeated. Neurons quietly dividing a problem among themselves because the loss went down when they did.

Summary

Five things worth keeping:

• A neuron is \(wx + b\). Everything else is arrangement.
• zero_grad → forward → loss → backward → step. Every training loop you ever write is this loop.
• Without a non-linearity, depth buys you nothing at all.
• When training misbehaves, look at your data's scale before you touch the architecture.
• The model will never tell you its answer is absurd.