Tensors, gradients and a first training loop

Learn in your browser · Chapter 01 of 06

Everything a model reads, stores and learns is a grid of numbers. Learn to read one before trusting any result.

A tensor is a grid of numbers in any number of dimensions (NumPy says array). It has a shape, the size of each dimension, and a dtype, the kind of number stored, such as float64 or int64.

Set up a virtual environment in a new folder (PowerShell shown; on macOS and Linux use python3, then .venv/bin/python). Pinned versions keep results comparable.

python -m venv .venv
.venv\Scripts\python -m pip install numpy==2.5.3 torch==2.14.0

This follows the CPU choice in PyTorch's install selector, pinned and without the unneeded torchvision: for Windows and macOS it gives pip3 install torch torchvision. Its Linux CPU line adds an --index-url option, so on Linux install numpy==2.5.3 alone, then torch==2.14.0 with that option. Only the Windows lines were run here. Save every script in this folder; per the venv documentation, calling the environment's Python needs no activation: .venv\Scripts\python shapes.py.

# shapes.py
import numpy as np

x = np.array([1.0, 2.0, 3.0])
m = np.array([[1, 2, 3], [4, 5, 6]])
print(x.shape, x.dtype, m.shape, m.dtype)
print(m + x)                  # x is stretched across both rows
col = x.reshape(3, 1)
print((col - x).shape)        # (3, 1) and (3,) broadcast to (3, 3)
counts = np.array([1, 2, 3])
counts[0] = 0.9               # an integer array silently drops the .9
print(counts)

(3,) float64 (2, 3) int64
[[2. 4. 6.]
 [5. 7. 9.]]
(3, 3)
[0 2 3]

Broadcasting combines different shapes. Compare shapes from the right: sizes fit if they are equal or one is 1, and a missing dimension counts as 1. So x (3,) is reused for each row of m (2, 3). The danger: a column (3, 1) minus a row (3,) silently becomes a (3, 3) table of every pairing.

Learn in your browser · Chapter 02 of 06

Training needs one number that says how wrong the model is, and a way to tell which direction makes it smaller.

toy.py invents 20 points near y = 3x + 2. Its seed fixes the random numbers, so every run makes the same data. The model is a line with two parameters, weight w and bias b: prediction = w * x + b.

# toy.py: data, loss and gradient
import numpy as np

rng = np.random.default_rng(seed=0)
x = rng.uniform(-1, 1, size=20)
y = 3 * x + 2 + rng.normal(0, 0.1, size=20)   # a line plus noise


def loss(w, b):
    return np.mean((w * x + b - y) ** 2)


def gradient(w, b):
    err = w * x + b - y
    return np.mean(2 * err * x), np.mean(2 * err)

The loss scores 'how wrong' as one number. Here it is the mean squared error: the average of (prediction minus target) squared, so big misses count most.

The gradient holds the loss's slope for each parameter. It points uphill, says the Deep Learning textbook, so its negative points downhill. Differentiating gives the slopes in gradient: the mean of 2 * err * x for w, and of 2 * err for b.

Always check a hand-derived gradient, for instance by nudging each parameter slightly both ways and measuring the change in loss: a finite difference.

# check.py
from toy import loss, gradient

w, b, h = 0.5, -1.0, 1e-6          # h is a tiny nudge
print("by hand %.6f %.6f" % gradient(w, b))
print("nudged  %.6f %.6f" % ((loss(w + h, b) - loss(w - h, b)) / (2 * h),
                             (loss(w, b + h) - loss(w, b - h)) / (2 * h)))

by hand -2.194225 -6.135977
nudged  -2.194225 -6.135977

They agree to six decimals. Both slopes are negative, so raising w and b lowers the loss: the true 3 and 2 are above 0.5 and -1.

Learn in your browser · Chapter 03 of 06

Gradient descent repeats one move: a small step downhill, then look again. The step size decides everything.

Each update is: new value = old value minus learning rate (a small positive number you choose) times slope. descend.py takes 50 steps from w = 0, b = 0 (loss 7.9804) at three rates.

from toy import loss, gradient

for lr in [0.01, 0.3, 1.1]:
    w, b = 0.0, 0.0
    for step in range(50):
        gw, gb = gradient(w, b)
        w, b = w - lr * gw, b - lr * gb   # step downhill
    print(f"lr {lr}: w {w:.3f} b {b:.3f} loss {loss(w, b):.4g}")

lr 0.01: w 1.035 b 1.313 loss 2.104
lr 0.3: w 2.971 b 2.007 loss 0.004496
lr 1.1: w -893.426 b -21299.910 loss 4.55e+08

At 0.01 the loss is still 2.104: right direction, too slow. At 0.3 it reaches 0.004496, the best any line can do here, as NumPy's least-squares line fit np.polyfit(x, y, 1) confirms. At 1.1 the first five losses are 6.599, 9.495, 13.722, 19.834 and 28.668: each step overshoots further.

A rising or swinging loss means too large a rate. PyTorch's tutorial warns that small values learn slowly and large ones may behave unpredictably. Try several, spread well apart, watching the loss.

Learn in your browser · Chapter 04 of 06

Hand-derived gradients do not scale. Let PyTorch compute them, check it, then write the standard training loop.

Automatic differentiation (autograd) records the forward calculation as a graph, then applies the chain rule back through it; that backward sweep is backpropagation. Mark a tensor requires_grad=True, call loss.backward(), and its slope appears in .grad. This uses check.py's point, in float64 to match NumPy.

# autograd_check.py
import torch
from toy import x, y, gradient

xt, yt = torch.tensor(x), torch.tensor(y)
w = torch.tensor(0.5, dtype=torch.float64, requires_grad=True)
b = torch.tensor(-1.0, dtype=torch.float64, requires_grad=True)
print("by hand    %.6f %.6f" % gradient(0.5, -1.0))
for run in ["autograd  ", "not zeroed"]:
    loss = ((w * xt + b - yt) ** 2).mean()      # forward and loss
    loss.backward()                              # backward
    print(run, "%.6f %.6f" % (w.grad.item(), b.grad.item()))

by hand    -2.194225 -6.135977
autograd   -2.194225 -6.135977
not zeroed -4.388451 -12.271953

Three routes agree; in full, autograd and hand differ by under 1e-15, float64 rounding. The last line is a classic bug: the backward documentation says gradients accumulate, so a second pass without clearing doubled them.

Now the real loop, in float32, PyTorch's default. Linear(1, 1) is the same line, MSELoss the same loss, SGD the same update. A batch is the examples behind one update; an epoch is one pass through the data, here 4 batches of 5.

# train_line.py
import torch
from toy import x, y

torch.manual_seed(0)
X = torch.tensor(x, dtype=torch.float32).reshape(20, 1)
Y = torch.tensor(y, dtype=torch.float32).reshape(20, 1)
model = torch.nn.Linear(1, 1)
loss_fn = torch.nn.MSELoss()
opt = torch.optim.SGD(model.parameters(), lr=0.1)

for epoch in range(1, 31):
    order = torch.randperm(20)                 # shuffle each epoch
    for start in range(0, 20, 5):              # batches of 5
        batch = order[start:start + 5]
        opt.zero_grad()                        # 1 clear
        loss = loss_fn(model(X[batch]), Y[batch])  # 2 forward, 3 loss
        loss.backward()                        # 4 backward
        opt.step()                             # 5 step
    if epoch % 10 == 0:
        with torch.no_grad():
            print(epoch, round(loss_fn(model(X), Y).item(), 5))
print(round(model.weight.item(), 3), round(model.bias.item(), 3))

10 0.00759
20 0.00452
30 0.0045
2.971 2.007

The numbered steps are the training loop; PyTorch's documentation clears gradients first, though clearing after the step works too. torch.no_grad() stops tracking while you only measure, and .reshape(20, 1) gives targets the predictions' shape.

Learn in your browser · Chapter 05 of 06

A straight line cannot fit every pattern. Add a hidden layer, then learn to tell learning from memorising.

No straight line separates XOR-shaped data: points near (0, 1) and (1, 0) are one class, points near (0, 0) and (1, 1) the other. The Deep Learning textbook shows a linear model fitted to the four XOR points outputting 0.5 everywhere. A multilayer network adds a hidden layer and a non-linear function (Tanh).

xor.py makes 200 noisy points, holds the last 50 back as a validation set, and trains a line, then a 33-parameter network. Targets are 0 or 1, so the loss is binary cross-entropy (BCEWithLogitsLoss takes raw scores).

import torch

torch.manual_seed(0)
corners = torch.tensor([[0.0, 0.0], [0.0, 1.0], [1.0, 0.0], [1.0, 1.0]])
pick = torch.randint(0, 4, (200,))
X = corners[pick] + 0.3 * torch.randn(200, 2)
Y = torch.tensor([0.0, 1.0, 1.0, 0.0])[pick].reshape(200, 1)  # XOR labels
loss_fn = torch.nn.BCEWithLogitsLoss()


def train(model, n=150, report=(2000,)):
    opt = torch.optim.SGD(model.parameters(), lr=0.5)
    for epoch in range(1, 2001):
        opt.zero_grad()
        loss = loss_fn(model(X[:n]), Y[:n])    # learn from the first n
        loss.backward()
        opt.step()
        if epoch in report:                    # judge on the last 50
            with torch.no_grad():
                out = model(X[150:])
                acc = ((out > 0).float() == Y[150:]).float().mean().item()
            print(epoch, "train %.3f valid %.3f accuracy %.0f%%"
                  % (loss.item(), loss_fn(out, Y[150:]).item(), 100 * acc))


print("line")
train(torch.nn.Linear(2, 1))
print("network")
net = torch.nn.Sequential(torch.nn.Linear(2, 8), torch.nn.Tanh(), torch.nn.Linear(8, 1))
train(net, report=(200, 400, 1000, 2000))

line
2000 train 0.680 valid 0.746 accuracy 54%
network
200 train 0.327 valid 0.330 accuracy 90%
400 train 0.249 valid 0.269 accuracy 90%
1000 train 0.222 valid 0.306 accuracy 86%
2000 train 0.210 valid 0.364 accuracy 88%

The line's training loss, 0.680, is near 0.693 (ln 2), the score for always answering 50:50, and 54% is near a coin toss. The network reaches 90%.

Training loss keeps falling, but validation loss bottoms at 0.269 near epoch 400 (at 384, checked every epoch), then rises to 0.364: the network has started fitting noise. The textbook calls this overfitting, the gap between training and test error growing too large. Remedies: keep the weights from the best validation epoch (early stopping), add data, or shrink the model.

Seeds. torch.manual_seed(0) fixes the data and starting weights; every script here printed identical output in three fresh environments on one Windows 11 machine (macOS and Linux were not tested). The NumPy and PyTorch documentation expect identical numbers only with the same versions on the same platform, so elsewhere your numbers may differ. Compare the pattern first and, as a rough rule of thumb, the first two significant figures.

Learn in your browser · Chapter 06 of 06

Most failed training runs fail quietly. Learn the usual suspects, then see what this toy leaves out.

Scroll sideways to see every column.

Save this broken loop as buggy.py and run it. The target is w = 3, b = 2.

import torch
from toy import x, y

torch.manual_seed(0)
X = torch.tensor(x, dtype=torch.float32).reshape(20, 1)
Y = torch.tensor(y, dtype=torch.float32)
model = torch.nn.Linear(1, 1)
loss_fn = torch.nn.MSELoss()
opt = torch.optim.SGD(model.parameters(), lr=0.1)

for step in range(100):
    loss = loss_fn(model(X), Y)
    loss.backward()
    opt.step()
print(round(loss.item(), 4), round(model.weight.item(), 3), round(model.bias.item(), 3))

4.5777 -0.057 1.767

PyTorch also prints a UserWarning that the target size differs from the input size. Never scroll past warnings.

Toy training is not frontier pretraining. The loop is the same idea, but published reports of large language-model runs show how much surrounds it (GPU memory maths covers the sums). These figures describe particular research runs, not how any current product is trained.

Scroll sideways to see every column.

The scaling-laws paper found loss falls as a power law as model size, data and compute grow, some trends spanning more than seven orders of magnitude; the compute-optimal paper trained over 400 models, from 70 million to over 16 billion parameters, to work out how best to spend a compute budget. The toy teaches the idea, not the practice.