Workflow

Goal of the lesson#

By the end of this 3-hour session you should be able to:

• recognize the five steps every PyTorch project follows,
• prepare data and split it into train and test,
• build a model by subclassing nn.Module,
• write a training loop with a loss, an optimizer and gradient descent,
• track and plot loss curves,
• save and reload a trained model,
• fit a noisy non-linear curve as a capstone.

This is the most important lesson of the series. Every chapter that follows reuses the same five-step skeleton — only the data and the model change.

Suggested timing#

The 5-step workflow#

In machine learning, the model is a tiny part of the project. Most of your time will be on data and on training/diagnostics. The shape of the workflow stays remarkably constant: the same five steps for a 50-line linear regression and for a 500-million-parameter language model.

In this lesson we work on the smallest interesting problem — a linear regression that learns the line y = 0.7 x + 0.3 — so we can focus entirely on the workflow.

Setup#

uv init --python 3.12 workflow
cd workflow
uv add torch matplotlib

Imports we will reuse:

import matplotlib.pyplot as plt
import torch
from torch import nn

device = "cuda" if torch.cuda.is_available() else "cpu"
torch.manual_seed(42)

1. Prepare data#

Real ML starts with real data. Here we generate it ourselves so we know the answer in advance and can verify whether the model finds it.

WEIGHT = 0.7
BIAS = 0.3

x = torch.arange(0, 1, 0.02).unsqueeze(dim=1)   # shape [50, 1]
y = WEIGHT * x + BIAS                            # shape [50, 1]

print(x[:5])
print(y[:5])
print(x.shape, y.shape)