Machine Learning from Scratch · Lesson 01 · Part 2

Your First Neuron, in PyTorch

Turn the hand calculation from Part 1 into a program, and let PyTorch compute the gradient for you.

Before you begin

Finish Part 1 first: this part reuses its dataset, its first update and its numbers. You need basic Python (variables, lists, loops, indentation) and an ordinary computer; no GPU.

Plan for about one to two hours, in one or two sessions: up to half an hour to install Python and PyTorch, and the rest for the labs and exercises. This is an estimate; take as long as the labs need.

By the end you should be able to explain every line of a training loop, check PyTorch’s gradients against your own, and diagnose a run that diverges.

Recap — what you already know from Part 1

In Part 1 you calculated the loss, the gradients and the parameter updates yourself. You already know how this model learns. Part 2 implements exactly the same operations; the mathematics does not change. PyTorch stores the numbers as tensors and calculates the gradient automatically. You still choose the model, the loss and the learning rate, and you still write the update.

The same steps in both parts. The numbers are Part 1’s first update, starting at $w = b = 0$ .
Part 1 — by hand	Part 2 — PyTorch
Calculate the predictions $\hat{y} = wx + b$	`prediction = x * w + b`
Calculate the mean squared error: $L = 4400$	`loss = torch.mean(...)`
Derive the partial derivatives: $\frac{\partial L}{\partial w} = -440$ , $\frac{\partial L}{\partial b} = -120$	`loss.backward()`
Update the parameters by hand: $w = 0.44$ , $b = 0.12$	`w -= lr * w.grad`
Repeat the calculations	`for epoch in range(...)`

Keep these numbers at hand. In Lab A you will make PyTorch reproduce every one of them; if it disagrees, the code is wrong, not the mathematics.

1. Tensors — the numbers PyTorch works with

A tensor is an array of numbers with a shape and a data type. A single number is a zero-dimensional tensor (a scalar); a row of five numbers is a one-dimensional tensor with shape (5,); a grid is two-dimensional. “Dimension” counts axes, not values.

Our inputs and targets have shape (5,); the weight and bias are scalars with shape (). Multiplying a five-element tensor by a scalar applies the scalar to every element. This is broadcasting. Predictions and errors keep shape (5,), and their mean is a scalar. We use floating-point numbers because the updates and gradients are fractions.

Shapes matter. If the targets had shape (5, 1) and the predictions (5,), broadcasting would silently compare every prediction with every target, a (5, 5) grid, and compute the wrong loss.

2. Autograd — PyTorch computes the gradient

In Part 1 you computed the predictions, then the errors, then their squares and mean. That is the forward pass. PyTorch records these steps as a computation graph. Calling backward() walks the graph in reverse and applies the chain rule at each step; this is backpropagation. The machinery that does it automatically is called automatic differentiation, and PyTorch’s version is autograd.

Parameters + inputs → predictions
Predictions − targets → errors
Square errors → mean → loss
Backward pass → gradient
Optimizer uses gradient → new parameters → repeat

The forward pass measures the mistakes, the backward pass computes the gradient, and the update changes the model. These are three separate jobs.

Autograd computes exactly the formulas you derived by hand:

\begin{aligned} \frac{\partial L}{\partial w} &= \frac{2}{5} \sum_{i=1}^{5} e_i x_i \\[6pt] \frac{\partial L}{\partial b} &= \frac{2}{5} \sum_{i=1}^{5} e_i \end{aligned}

It does not update the parameters. The code that changes them is the optimizer. PyTorch has ready-made optimizers, but in this lesson you write your own in two lines.

3. Guided laboratory

Set up a CPU environment

Use Python 3.10–3.14. PyTorch 2.14 has no packages for newer versions yet, so on Python 3.15 pip install torch fails with “No matching distribution found”. On a Mac with Homebrew, run brew install python@3.13 and create the environment with python3.13 -m venv .venv instead of the first line below. If installation fails for another reason, use the PyTorch installation selector. On Windows, activate the environment with .venv\Scripts\Activate.ps1 in PowerShell.

python3 -m venv .venv
source .venv/bin/activate
python -m pip install torch
python -c "import sys, torch; print(sys.version); print(torch.__version__)"
python train.py

PyTorch is the only package you need; there is no dataset to download. Note the Python and PyTorch versions the check prints and keep them with your results. The program has no randomness, so it needs no random seed.

Tested on CPU with Python 3.12.14 and PyTorch 2.14.1.

Start small: reproduce one state

Before writing the whole loop, check that PyTorch agrees with your hand calculation. This snippet sets w and b to the values after the first update in Part 1 and only measures the model; it does not learn.

Show the forward pass

import torch

x = torch.tensor([1., 2., 3., 4., 5.])
y = torch.tensor([20., 40., 60., 80., 100.])
w = torch.tensor(0.44)
b = torch.tensor(0.12)
prediction = x * w + b
print(prediction)
print(torch.mean((prediction - y) ** 2).item())

x and y hold the five samples; w and b hold the state after one update. There is no requires_grad or backward because nothing is being trained. Expect [0.56, 1.00, 1.44, 1.88, 2.32] and an MSE of about 4194.4608, matching the worked example in Part 1.

The complete program

Download train.py

Show the complete program

import torch

x = torch.tensor([1., 2., 3., 4., 5.])
y = torch.tensor([20., 40., 60., 80., 100.])
w = torch.tensor(0.0, requires_grad=True)
b = torch.tensor(0.0, requires_grad=True)
lr = 0.001

for epoch in range(1000):
    prediction = x * w + b
    loss = torch.mean((prediction - y) ** 2)
    w.grad = None
    b.grad = None
    loss.backward()
    if epoch % 100 == 0:
        print(f"epoch={epoch:4d} loss={loss.item():.3f} w={w.item():.3f} b={b.item():.3f}")

    with torch.no_grad():
        w -= lr * w.grad
        b -= lr * b.grad

with torch.no_grad():
    print("Final MSE:", torch.mean((x * w + b - y) ** 2).item(), "w:", w.item(), "b:", b.item())
    print("Prediction for 7:", (7 * w + b).item())

Read the program line by line

import torch: Loads PyTorch’s tensors and autograd.
x = torch.tensor(...) and y = torch.tensor(...): The five inputs and their targets, in matching order. The decimal points make them floating-point. Both have shape (5,), and neither is trained.
w = torch.tensor(0.0, requires_grad=True) and the matching b line: The two trainable parameters, as scalars. requires_grad=True tells autograd to track how they affect the loss. Starting from zero is fine for a single neuron; in networks with many layers it is not, and a later lesson explains why.
lr = 0.001: The learning rate $\eta$ . You choose it; the program never changes it.
for epoch in range(1000):: 1000 full-batch updates, numbered 0 to 999. Each new gradient is computed from the parameters the previous update produced.
prediction = x * w + b: The forward pass for all five samples at once. Broadcasting shares one w and one b across them.
loss = torch.mean((prediction - y) ** 2): Errors, squared, averaged: exactly the MSE from Part 1.
w.grad = None and b.grad = None: PyTorch adds new gradients to old ones by default. Clearing them keeps only the current gradient. This clears the gradient, not the parameter.
loss.backward(): Runs backpropagation and stores $\frac{\partial L}{\partial w}$ and $\frac{\partial L}{\partial b}$ in w.grad and b.grad. It does not change w or b.
if epoch % 100 == 0: and print(...): Reports every 100th epoch, before the update, so the loss, w and b in a line all belong to the same state; epoch 0 is the untrained model. .item() turns a scalar tensor into a Python number, and :.3f rounds only the printout.
with torch.no_grad():: Stops autograd from recording. The update is not part of the loss calculation, and PyTorch only lets you change these tensors in place outside the graph.
w -= lr * w.grad and b -= lr * b.grad: The optimizer: gradient descent, as in Part 1. Both gradients come from the same forward pass, so the update is simultaneous even though it takes two lines.
The final with torch.no_grad(): and two print calls: Measure the trained model and predict x = 7 without recording a graph. Using a model without training it is called inference. The final loss is still measured on the training samples, so it says nothing about new data.

Lab A — stop after one update

Before running, write down the initial loss, both gradients and the new parameters from Part 1’s worked example.
In a copy of the file, change range(1000) to range(1). Right after loss.backward(), add print(w.grad.item(), b.grad.item()).
Run the copy. Expect −440 and −120, then w ≈ 0.44, b ≈ 0.12 and a final MSE ≈ 4194.4608. Small differences in the last digits come from floating-point arithmetic.
If you print the original loss again after the update, it still says 4400. Why? It is a value computed earlier, not a formula that recalculates itself; you have to run the forward pass again to measure the new state.

Lab B — train, observe, explain

Run the original 1000-update program. The first report is epoch= 0 loss=4400.000 w=0.000 b=0.000.
The final result should be close to MSE 2.45, w 18.98, b 3.67, and a prediction of 136.56 at x = 7. Check $7w + b$ with the parameters your own run printed.
Repeat with learning rates 0, 0.001, 0.05 and 0.1, each time from zero. Record the rate, the initial and final loss, and whether the loss stayed finite. At 0.1 the loss overflows to inf and then nan (“not a number”): the run has failed, and its prediction means nothing.
Explain why a falling training loss does not prove the model is accurate on new data. All five points were used for training; nothing was held back to test it.

The rule the data were invented from gives 140 at seven hours, but seven lies outside the training range [1, 5]. Predicting outside the range you trained on is called extrapolation. The answer also exceeds a typical 100-point exam maximum, because a straight line has no upper bound.

Common misconceptions

“Backward trains the weights.” Print w before and after backward(), then after the update. Which step actually changes it?
“A negative gradient is a negative error.” An error belongs to one prediction; the gradient combines all five samples. What role does x play in the weight gradient?
“A larger learning rate always learns faster.” Compare 0.05 and 0.1.
“A zero gradient means a perfect model.” It means you are at a flat point of the loss. With noisy data or a model that is too simple, the lowest loss can still be above zero.
“Clearing the gradient means setting the parameter to zero.” Clear .grad, not w or b. If every epoch reset the weights, what progress would survive?
“A low MSE means accuracy on unseen data.” It measures the fit to the samples you trained on. What evidence is still missing?

4. Exercises

Start each exercise from a fresh copy of train.py. Write your prediction first, then code it.

Start at $w = 10$ , $b = 5$ and run one update. Compare autograd’s gradients and the new parameters with your answer to exercise 1 in Part 1.
Freeze b at zero by not updating it, and train only w. Why can this dataset still be fitted exactly? Now use the targets [25, 45, 65, 85, 105]: what does the frozen bias cost you?
Record the loss at updates 0–100 and plot it on paper or in a tool of your choice, with labelled axes (updates; MSE in points²). Compare two learning rates from the same starting point.
Use the trained model at x = 3.5 and x = 7. Which is interpolation and which is extrapolation? Why does neither tell you how well the model generalizes?

Hints

Freezing a parameter means leaving out its update line. To compare runs fairly, recreate the starting tensors each time. Interpolation stays inside the range of inputs the model was trained on.

5. Assessment — explain the training loop

Answer without scrolling back. Back each answer with code you ran or output you saw.

If you call backward() repeatedly but delete the update block, what happens to the predictions? What if you keep the updates but stop clearing the gradients?
A run with learning rate 0.1 prints nan. What went wrong, and why is it not a bug in PyTorch?
A run returns 140 for x = 7. What does that tell you about the implementation, and what doesn’t it tell you about real exam results?

Reasoning guide — open after writing your answers

backward() alone only fills the gradient buffers; the predictions never change. Without clearing, each backward pass adds to the previous gradients, so every step uses a sum of old gradients: a different algorithm, which here quickly becomes unstable. At 0.1 every step overshoots further than the last until the float32 numbers overflow to infinity and then NaN: the divergence you saw in Part 1’s simulation. A prediction of 140 matches the rule the data were invented from, so the code works. But the data are synthetic and x = 7 lies outside the training range, so it says nothing about real students.

Completion checklist

Keep your code, the version printout, your loss plot and the learning-rate table. The checkboxes are reminders for you, not a test.

I can explain every line of the training loop. I have checked PyTorch’s gradients against my own calculation. I can explain divergence, extrapolation, and why training loss is not a test score.

Saved only in this browser.

References and next steps

PyTorch’s tensor tutorial goes deeper into shapes, and PyTorch’s autograd tutorial into differentiation. Lesson 02 introduces richer tensor operations and nonlinear networks.

← Back to the course