🎓 A1Learner
🔍

PyTorch Beginner Capstone

The beginner capstone is one full training program that builds a model from raw data to a saved file. It trains a house-price regressor on synthetic data with scaled features and a validation split.

The price rule

Each made-up house gets its price from a fixed rule plus noise. Bigger and newer houses cost more.

rooms, area, age = 3.0, 1450.0, 12.0
price = 60 + 18 * rooms + 0.09 * area - 1.2 * age
print(price)

Scale one feature

Subtracting the mean centres the values on zero. Dividing by the std evens out the spread.

area = torch.tensor([600.0, 1200.0, 1800.0, 2400.0])
scaled = (area - area.mean()) / area.std()
print(scaled)

Split rows with one shuffle

randperm shuffles the row numbers once. The first part trains, and the rest validates.

torch.manual_seed(0)
order = torch.randperm(8)
print(order[:6])
print(order[6:])

Example

import torch
from torch import nn
from torch.utils.data import DataLoader, TensorDataset

torch.manual_seed(0)

rows = 160
rooms = torch.randint(1, 6, (rows,)).float()
area = torch.rand(rows) * 2000 + 500
age = torch.rand(rows) * 40
noise = torch.randn(rows) * 8
price = 60 + 18 * rooms + 0.09 * area - 1.2 * age + noise

features = torch.stack([rooms, area, age], dim=1)
targets = price.reshape(-1, 1)

order = torch.randperm(rows)
train_idx, val_idx = order[:128], order[128:]
train_x, train_y = features[train_idx], targets[train_idx]
val_x, val_y = features[val_idx], targets[val_idx]

mean = train_x.mean(dim=0)
std = train_x.std(dim=0)
train_x = (train_x - mean) / std
val_x = (val_x - mean) / std

train_set = TensorDataset(train_x, train_y)
loader = DataLoader(train_set, batch_size=32, shuffle=True)

torch.manual_seed(1)
model = nn.Sequential(
    nn.Linear(3, 16),
    nn.ReLU(),
    nn.Linear(16, 1),
)
loss_fn = nn.MSELoss()
optimizer = torch.optim.Adam(model.parameters(), lr=0.05)

print("Epoch | Train loss | Val loss")
for epoch in range(1, 61):
    model.train()
    for batch_x, batch_y in loader:
        loss = loss_fn(model(batch_x), batch_y)
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
    if epoch % 15 == 0:
        model.eval()
        with torch.no_grad():
            train_loss = loss_fn(model(train_x), train_y).item()
            val_loss = loss_fn(model(val_x), val_y).item()
        print(f"{epoch:>5} | {train_loss:>10.1f} | {val_loss:>8.1f}")

torch.save(model.state_dict(), "capstone_model.pt")

fresh = nn.Sequential(
    nn.Linear(3, 16),
    nn.ReLU(),
    nn.Linear(16, 1),
)
fresh.load_state_dict(torch.load("capstone_model.pt"))
fresh.eval()
model.eval()
with torch.no_grad():
    same = torch.equal(model(val_x), fresh(val_x))
    final_val = loss_fn(fresh(val_x), val_y).item()

print("Saved and reloaded model matches:", same)
print(f"Reloaded validation loss: {final_val:.1f}")

Key points

  • The target is a house price in thousands of dollars
  • Use one seed at the top so the whole run repeats exactly
  • Scale each feature with the training slice's mean and std
  • The validation slice judges the model on rows training never touched
  • Use torch.save on the state_dict to keep the trained model in a file

In this example, one seed at the top makes the whole run repeatable. The data is 160 made-up houses, each with rooms, floor area, and age. The price comes from a fixed rule plus random noise. A shuffled split holds back 32 rows for validation. The features are scaled with the training rows' mean and std. An MLP with one hidden layer trains for 60 epochs on batches of 32. Every 15 epochs, the train and validation losses print side by side. Both losses fall steeply, ending near 140. The trained state_dict is saved to a file and loaded into a fresh model. The fresh model's predictions match exactly, and its validation loss prints last. The next course opens with PyTorch Intermediate Intro.

Try it Yourself

Edit the code, then press Run.

Running is simulated · Python Playground

Exercise

Why does the capstone scale the features before training?