Practical Build 0

Linear regression: the first complete model

Introduction

Three clean points follow one hidden line: negative 1 maps to negative 5, 0 maps to negative 2, and 1 maps to 1. This lesson recovers the slope and intercept through a complete training workflow.

Learning goal

Calculate one gradient, train and validate a linear model, diagnose failures, and preserve its learned state.

Before you start

Basic algebra, Python functions and loops, lists, averages, and the idea of a derivative.

Lesson plan

  1. Write the line and its derivatives by hand while naming every data shape.
  2. Set a baseline, build the linear layer, and run training plus validation.
  3. Investigate common failures, then save and reload only the learned state.

Download the complete CPU Python script. It uses only PyTorch and the Python standard library. It does not download data or write files unless you explicitly pass a checkpoint path.

Begin with two points and one line

Linear regression predicts a number from measured input features. With one input, the model is a straight line:

prediction = weight * x + bias

The weight controls the slope. It says how much the prediction changes when x increases by one. The bias is the prediction at x = 0. Our hidden data rule is y = 3x - 2. Therefore the correct weight is 3 and the correct bias is -2. For x = -1, the exact target is -5. For x = 1, it is 1.

Real measurements are imperfect, so the full dataset adds small random noise. The model should recover values close to 3 and -2 rather than memorizing every noisy point. This is why we need held-out validation data.

First write the forward function and derivatives

Do not begin with a model class. First make the calculation visible. Let error = prediction - target. Mean squared error, or MSE, squares every error and averages the results:

Predict first: with weight 0 and bias 0, both predictions are 0. Decide whether increasing the weight should help the point (1, 1). Then decide whether lowering the bias should help both points. Keep those directions beside the gradients below.

loss = mean((prediction - target) ** 2)

d_loss_d_weight = mean(2 * error * x)
d_loss_d_bias = mean(2 * error)

The derivatives answer: “If I move this parameter a tiny amount, which way does the loss move?” Gradient descent subtracts a small multiple of each derivative.

import torch

def forward(x, weight, bias):
    return weight * x + bias

def mse_and_gradients(x, target, weight, bias):
    prediction = forward(x, weight, bias)
    error = prediction - target
    loss = (error ** 2).mean()
    weight_gradient = (2 * error * x).mean()
    bias_gradient = (2 * error).mean()
    return loss, weight_gradient, bias_gradient

x = torch.tensor([[-1.0], [1.0]])
target = torch.tensor([[-5.0], [1.0]])
weight = torch.tensor(0.0)
bias = torch.tensor(0.0)

loss, dw, db = mse_and_gradients(x, target, weight, bias)
print(loss.item(), dw.item(), db.item())

# Expected: 13.0 -6.0 4.0

At weight 0 and bias 0, both predictions are zero. The errors are 5 and -1. Their squared average is 13. The negative weight gradient says increasing the weight should reduce loss. With learning rate 0.1, one update gives weight = 0.6 and bias = -0.4. The loss falls from 13 to 8.32. One step is not enough, but it moves in the right direction.

The prediction check and the gradient signs agree. We subtract the negative weight gradient, so the weight rises. We subtract the positive bias gradient, so the bias falls. If an update moves against this hand check, inspect the error sign or the update sign before changing the learning rate.

Know every data shape and dtype

We use 256 training examples and 128 validation examples. Each example has one input feature and one numeric target. Therefore both x and target have shape (N, 1). They use torch.float32, PyTorch's common floating-point dtype for neural-network calculations.

NThe number of examples in this tensor.
Feature axisThe final axis. Its size is one because each example has one input.
Training setData used to compute gradients and update parameters.
Validation setIndependent data used only to measure performance.
def make_dataset(count, seed, noise_std=0.15):
    generator = torch.Generator().manual_seed(seed)
    x = torch.rand((count, 1), generator=generator) * 4 - 2
    noise = torch.randn((count, 1), generator=generator)
    target = 3 * x - 2 + noise_std * noise
    return x.float(), target.float()

train_x, train_y = make_dataset(256, seed=10)
validation_x, validation_y = make_dataset(128, seed=20)

print(train_x.shape, train_x.dtype)
# Expected: torch.Size([256, 1]) torch.float32

The two seeds create separate random draws. Validation points are never included in an optimizer step. Reusing training data as “validation” would only measure how well the model fits data it already saw.

Set a baseline before training

A model is useful only if it beats a simple comparison. Our baseline ignores x and always predicts the mean training target. We then measure that fixed rule on validation data.

baseline = train_y.mean().expand_as(validation_y)
baseline_mse = torch.nn.functional.mse_loss(baseline, validation_y)
print(baseline_mse.item())

The exact value is deterministic for the supplied seeds, but the important comparison is relational: the trained line should have far lower validation MSE. The complete script asserts that model MSE is below five percent of baseline MSE and below 0.08.

Move from a function to nn.Linear

nn.Linear(1, 1) implements the same weight * x + bias rule. The first 1 means one input feature. The second 1 means one output number. PyTorch stores its weight and bias as trainable tensors with requires_grad=True.

nn.Linear is an object made from a class, but you do not need object-oriented programming knowledge to use it. Think of model(x) as a function call. model.parameters() returns the two tensors that the optimizer may update.

from torch import nn

torch.manual_seed(2026)
model = nn.Linear(in_features=1, out_features=1)

print(tuple(model.weight.shape))  # Expected: (1, 1)
print(tuple(model.bias.shape))    # Expected: (1,)

Run the complete training and validation loop

One full-batch step uses all 256 training examples. The five operations stay in this order: clear old gradients, predict, calculate loss, run backward, and update parameters.

optimizer = torch.optim.SGD(model.parameters(), lr=0.05)

model.train()
for epoch in range(600):
    optimizer.zero_grad()
    prediction = model(train_x)
    loss = nn.functional.mse_loss(prediction, train_y)
    loss.backward()
    optimizer.step()

model.eval()
with torch.no_grad():
    validation_prediction = model(validation_x)
    validation_mse = nn.functional.mse_loss(
        validation_prediction, validation_y
    )

print("validation MSE:", validation_mse.item())
print("weight:", model.weight.item())
print("bias:", model.bias.item())

# Expected pattern, with the supplied seeds:
# validation MSE is far below the baseline
# weight is close to 3.0
# bias is close to -2.0

loss.backward() computes gradients but does not update anything. optimizer.step() performs the update. model.train() selects training behavior for layers such as dropout. model.eval() selects evaluation behavior. Neither command turns gradient recording on or off. The torch.no_grad() context separately disables graph recording for validation.

Investigate four common failures

1. No optimizer step

If loss and gradients exist but the parameters never change, check that optimizer.step() runs after loss.backward(). Autograd reports directions; it does not move parameters.

2. No gradient clearing

PyTorch adds each backward result to existing gradients. Without optimizer.zero_grad(), updates depend on all earlier steps. This may make training unstable and hides the intended algorithm.

3. Silent shape broadcasting

If prediction has shape (N, 1) but target has shape (N,), subtraction can produce shape (N, N). Every prediction is compared with every target. Print both shapes and require them to match before calculating MSE.

4. Learning rate too large

A very large learning rate can jump across the minimum. Loss may oscillate, increase, or become inf or nan. Try lowering the rate by ten times and check whether the first several losses decrease.

assert prediction.shape == train_y.shape
assert torch.isfinite(loss), "loss became inf or nan"

Save and load only when requested

A state_dict maps parameter names to tensors. Save it instead of the whole Python model object. Loading requires creating the same architecture first.

# This writes a file, so choose the path explicitly.
torch.save({"model_state_dict": model.state_dict()}, "linear.pt")

restored = nn.Linear(1, 1)
checkpoint = torch.load("linear.pt", map_location="cpu", weights_only=True)
restored.load_state_dict(checkpoint["model_state_dict"])
restored.eval()

The downloadable script does not write automatically. Run it normally with:

python examples/pytorch/00_linear_regression.py

To explicitly save and verify a checkpoint, add a path:

python examples/pytorch/00_linear_regression.py \
  --checkpoint /tmp/linear-regression.pt

Exercises

1. If the true rule changes to y = -2x + 4, what parameters should training recover?

Reveal the worked answer

The expected weight is -2 and bias is 4. Change only the target formula in both dataset generators. Do not change the model architecture: one input and one output are still enough.

2. Why is a separately generated validation set better than measuring the final training loss?

Reveal the worked answer

Training loss measures examples that directly influenced parameter updates. Validation data did not influence those updates, so it gives a better test of whether the learned rule works on new samples from the same process.

3. Prediction is shaped (64, 1), but target is (64,). What should you do?

Reveal the worked answer

Make the target (64, 1), for example with target = target.reshape(-1, 1). Then assert that both shapes match. Do not rely on broadcasting because it compares the wrong pairs.

Recap and next step

You started from a plain function, derived MSE gradients, and watched one manual update reduce loss. You then represented the same line with nn.Linear, trained it with autograd and SGD, compared it with a mean baseline on independent validation data, and separated evaluation mode from gradient recording. The complete script also verifies learned parameters and supports an explicit checkpoint path.

Reference: PyTorch optimization tutorial.

Before the next build, review tensors and autograd if their shapes or gradients are still unfamiliar.