Unit 03.04: Building a minimal training loop

Every training loop in this course, and almost every one you will read elsewhere, is these four lines repeated.

The loop, and a problem with a checkable answer

Rather than train on something opaque, use data generated by a rule you know: y = 2x₀ - 3x₁ + 1. If the loop works, the learned parameters must converge to those coefficients.

That gives you a test, not just a decreasing number:

import torch
from torch import nn

torch.manual_seed(0)

# A problem with a real answer: y = 2*x0 - 3*x1 + 1
X = torch.randn(200, 2)
y = (2 * X[:, 0] - 3 * X[:, 1] + 1).unsqueeze(1)

model = nn.Linear(2, 1)
optimizer = torch.optim.SGD(model.parameters(), lr=0.05)
loss_fn = nn.MSELoss()

for epoch in range(200):
    optimizer.zero_grad()           # 1. clear last step's gradients
    loss = loss_fn(model(X), y)     # 2. forward pass and loss
    loss.backward()                 # 3. gradients
    optimizer.step()                # 4. update the parameters
    if epoch % 50 == 0:
        print(f"epoch {epoch:>3}  loss {loss.item():.4f}")

print("final loss:", round(loss_fn(model(X), y).item(), 6))
print("learned weights:", [round(v, 2) for v in model.weight.detach().flatten().tolist()])
print("learned bias   :", round(model.bias.item(), 2))
# Weights converge to about [2.0, -3.0] and bias to about 1.0 -- the model
# recovered the rule that generated the data.

The loss falls to essentially zero, and the learned weights come out at approximately [2.0, -3.0] with a bias near 1.0 — the model recovered the rule that generated the data.

This is the strongest kind of check available early on. A falling loss tells you something is happening; recovering known parameters tells you it is the right thing.

The mistake this prevents

Trusting a falling loss as proof of correctness. Loss falls for buggy loops too — with a mis-shaped target, a mis-scaled input, or a leak. Test on data whose answer you already know.

Takeaway

Four steps per iteration. Validate the loop on a problem with a known answer before pointing it at real data.

Mark complete