Clean PyTorch: Dataset, DataLoader, Model, Trainer

The late-delivery model had been running in shadow mode for three weeks. Every evening, a script scored the next day’s orders; every morning, Arjun compared the scores with what actually happened. It was doing roughly what the test set had promised. Then the monsoon arrived.

On the first properly wet Tuesday, 31 of 52 orders arrived late. The model had flagged nine of them. Meera forwarded the numbers with a one-word subject line: “Rain??”

Arjun opened the notebook that held everything: late_model_FINAL_v3.ipynb, 47 cells, some of them commented out, two of them called # DON'T RUN. He found the scoring code. It had been copied from cell 31 into a separate script months ago, and it did this:

todays_orders = load_orders("2026-07-14")
mean, std = todays_orders.mean(dim=0), todays_orders.std(dim=0)   # <- here
x = (todays_orders - mean) / std

It standardised each day’s orders using that day’s statistics. In training, about 30% of orders were in the rain, so a rainy order had a z-score of $(1-0.31)/0.46=1.5$: “much rainier than normal, worry”. On a day when 80% of orders were wet, the script computed a new mean of 0.8 and a new standard deviation of 0.4, so each rainy order became $(1-0.8)/0.4=0.5$. The model was being told that rain on the wettest day of the year was only mildly unusual.

“The model learned the right thing,” Priya said, reading over his shoulder. “The plumbing around it lied to it. And it lied because the important numbers lived in a notebook cell, not with the model.”

Arjun tries to tidy the notebook

Arjun’s first fix was obvious. He selected all 47 cells, pasted them into train.py, deleted the commented-out bits, and ran it. It worked. Then Priya asked for three things:

1. “Retrain on October’s orders and compare with September.”
2. “Try a model with 64 hidden neurons instead of 256.”
3. “Score tomorrow’s orders using the September model, without retraining.”

The first meant changing a filename in three places. The second meant editing the model definition, which lived in the middle of the training loop, which also held the early stopping logic and the plotting code. The third was impossible without copying half the file into another script, which is exactly how the rain bug had been born.

While doing the first, he noticed something else. On October’s data, epoch 1 printed a training loss of 0.36. A freshly initialised model should start near 0.69 (the loss for guessing 50% on everything). Priya pointed at it immediately. “That’s not a new model.”

In the notebook, he had re-run the data and training cells but not the cell that created the model. Python still held the old, trained model object from earlier that afternoon, so “retraining on October” had quietly meant “keep training the September model on October data”. In a notebook, the code you can see is not necessarily the code that ran.

The kitchen that runs on stations

“Walk me through how the central kitchen makes 600 croissants before 6am,” Priya said.

Arjun had seen it. “The storeroom keeps every ingredient on labelled shelves. A runner fetches what each station needs and brings it on trays. The lamination station follows the recipe. And the head chef runs the shift: what gets baked when, tasting, deciding when a batch is done.”

“Could the head chef swap the croissant recipe for pain au chocolat without retraining the runners?”

“Sure. The runners just bring trays of whatever’s on the list.”

“Could a new runner start on Monday without reading the recipe?”

“Yes. They only need to know where the shelves are.”

“That’s the whole design,” Priya said. “Four stations, each with one job and a clear handover. In PyTorch they have names.”

Kitchen	PyTorch	Its one job
Storeroom shelves	Dataset	Given an index, return one prepared example
Runner with trays	DataLoader	Shuffle, group examples into batches, deliver them
Recipe	nn.Module (the model)	Turn a batch of inputs into outputs
Head chef	Trainer (your class)	Run the loop: train, evaluate, stop early, save

Where the analogy breaks: a real head chef improvises. A Trainer should not. Its value is that it does the same thing, in the same order, every time.

Following order 17 through the stations

Before writing any code, Priya traced one order through the design so Arjun could see each handover. Order 17 from September’s CSV:

distance_km,items,hour,raining,partner_rating,late
7.97,9,13,0,3.09,0

A 7.97 km order of 9 items at 1pm, dry, with a partner rated 3.09, which arrived on time.

Storeroom (Dataset). Asked for item 17, it standardises the row using the training statistics it computed once when it was created (means of 8.02 km, 6.46 items, 14.94 hours, 0.31 rain, 3.999 rating) and returns a pair:

features: [-0.01,  0.75, -0.45, -0.67, -1.55]
label:    [0.]

Read the z-scores: average distance, a fairly big order, early afternoon, not raining, and a much weaker partner than usual (1.55 standard deviations below average).

Runner (DataLoader). It asks the storeroom for 64 random indices, 17 among them, and stacks the 64 pairs into two tensors: features of shape (64, 5) and labels of shape (64, 1). Stacking single examples into a batch is called collating.

Recipe (model). The batch goes through the network: (64, 5) in, (64, 1) raw scores out.

Head chef (Trainer). Computes the loss for the batch, runs the five-step update from Module 27, and after each epoch scores the validation orders. When validation stops improving, it reloads the best weights and saves a checkpoint (a file containing the model’s weights plus everything needed to use them again).

Later, the scoring script loads that checkpoint and asks about order 17: 21% chance of being late. Change the same order to 7pm in the rain and it says 66%. Crucially, both answers use the September training statistics stored in the checkpoint, never the statistics of whichever day happens to be scored.

The real code, station by station

A small project layout:

late_delivery/
├── data.py        # Dataset and CSV loading
├── model.py       # the network
├── trainer.py     # the training loop
├── train.py       # wires it together: config, data, model, trainer
└── predict.py     # loads a checkpoint, scores new orders

Station 1: the Dataset

A PyTorch Dataset is any class with two methods: __len__ (how many examples) and __getitem__ (return example number idx). Python calls these when you write len(ds) and ds[17].

# data.py
import csv
import torch
from torch.utils.data import Dataset

FEATURES = ["distance_km", "items", "hour", "raining", "partner_rating"]

def load_csv(path: str) -> tuple[torch.Tensor, torch.Tensor]:
    """Read an orders CSV into a feature tensor (n, 5) and a label tensor (n, 1)."""
    with open(path, newline="") as f:
        rows = list(csv.DictReader(f))
    X = torch.tensor([[float(r[c]) for c in FEATURES] for r in rows], dtype=torch.float32)
    y = torch.tensor([[float(r["late"])] for r in rows], dtype=torch.float32)
    return X, y

class DeliveryDataset(Dataset):
    """Crumb delivery orders. Knows how to fetch ONE standardised order."""

    def __init__(self, X: torch.Tensor, y: torch.Tensor,
                 mean: torch.Tensor | None = None, std: torch.Tensor | None = None):
        self.X, self.y = X, y
        # Training data computes its own statistics; validation and test data reuse them.
        self.mean = X.mean(dim=0) if mean is None else mean
        self.std = X.std(dim=0) if std is None else std

    def __len__(self) -> int:
        return len(self.X)

    def __getitem__(self, idx: int) -> tuple[torch.Tensor, torch.Tensor]:
        x = (self.X[idx] - self.mean) / self.std
        return x, self.y[idx]

The design choice that kills the rain bug is in __init__: the dataset either computes statistics (for training data) or is given them (for everything else). There is no code path where validation or production data standardises itself.

Check it before going further. Small checks at each station save hours later:

from torch.utils.data import DataLoader
from data import DeliveryDataset, load_csv

X, y = load_csv("orders_september.csv")
train_set = DeliveryDataset(X[:-500], y[:-500])
print(len(train_set))          # 2000
print(train_set[17])           # (tensor([-0.0128,  0.7475, -0.4492, -0.6670, -1.5524]), tensor([0.]))

loader = DataLoader(train_set, batch_size=64, shuffle=True)
xb, yb = next(iter(loader))
print(xb.shape, yb.shape)      # torch.Size([64, 5]) torch.Size([64, 1])

Station 2: the DataLoader

You do not write this one; you configure it. The arguments you will use most:

• batch_size: examples per batch.
• shuffle=True for training, False for validation and test (order does not matter there, and a fixed order makes results easier to compare).
• num_workers: how many background processes prepare batches in parallel. Leave it at 0 for small tables; set it to 2 to 8 for images, where loading and augmenting from disk is the bottleneck.
• drop_last=True: discard a final, smaller batch, occasionally useful with BatchNorm (next lesson).

Station 3: the model

The same network as Module 27, now in its own file and configurable:

# model.py
from torch import nn

class LateDeliveryNet(nn.Module):
    def __init__(self, n_features: int = 5, hidden: int = 64, dropout: float = 0.3):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(n_features, hidden), nn.ReLU(), nn.Dropout(dropout),
            nn.Linear(hidden, hidden), nn.ReLU(), nn.Dropout(dropout),
            nn.Linear(hidden, 1),
        )

    def forward(self, x):          # x: (batch, n_features)
        return self.net(x)         # (batch, 1) raw scores

Priya’s second request, “try 64 hidden neurons instead of 256”, is now a single argument. With hidden=64 it has 4,609 parameters.

Station 4: the Trainer

The head chef. Every line here appeared in Module 27; the class gives it a home and a name.

# trainer.py
import copy
import torch
from torch import nn
from torch.utils.data import DataLoader

class Trainer:
    """Runs the training loop, evaluation, early stopping and checkpoints."""

    def __init__(self, model: nn.Module, train_loader: DataLoader, val_loader: DataLoader,
                 lr: float = 1e-3, weight_decay: float = 1e-2, device: str = "cpu"):
        self.model = model.to(device)
        self.train_loader = train_loader
        self.val_loader = val_loader
        self.device = device
        self.loss_fn = nn.BCEWithLogitsLoss()
        self.optimizer = torch.optim.AdamW(model.parameters(), lr=lr, weight_decay=weight_decay)

    def train_one_epoch(self) -> float:
        self.model.train()                                  # dropout on
        total, count = 0.0, 0
        for xb, yb in self.train_loader:
            xb, yb = xb.to(self.device), yb.to(self.device)
            loss = self.loss_fn(self.model(xb), yb)
            self.optimizer.zero_grad()
            loss.backward()
            self.optimizer.step()
            total += loss.item() * len(xb)
            count += len(xb)
        return total / count

    @torch.no_grad()
    def evaluate(self) -> float:
        self.model.eval()                                   # dropout off
        total, count = 0.0, 0
        for xb, yb in self.val_loader:
            xb, yb = xb.to(self.device), yb.to(self.device)
            total += self.loss_fn(self.model(xb), yb).item() * len(xb)
            count += len(xb)
        return total / count

    def fit(self, epochs: int = 100, patience: int = 10) -> float:
        best_val, best_state, bad = float("inf"), None, 0
        for epoch in range(1, epochs + 1):
            train_loss = self.train_one_epoch()
            val_loss = self.evaluate()
            print(f"epoch {epoch:3d}  train {train_loss:.3f}  val {val_loss:.3f}")
            if val_loss < best_val:
                best_val, bad = val_loss, 0
                best_state = copy.deepcopy(self.model.state_dict())
            else:
                bad += 1
                if bad >= patience:
                    print(f"early stop; best val {best_val:.3f}")
                    break
        self.model.load_state_dict(best_state)
        return best_val

    def save(self, path: str, **extra) -> None:
        torch.save({"model_state": self.model.state_dict(), **extra}, path)

Three details worth noticing. @torch.no_grad() on evaluate switches off autograd for the whole method. Every batch is moved to self.device inside the loop, so the same class runs on a laptop CPU or a cloud GPU. And save accepts any extra items, which is how the training statistics travel with the weights.

Wiring it together

train.py is the only file that knows about all four stations. Its settings live in a dataclass (a Python class that is mostly a bundle of named fields), so every run’s configuration is explicit and can be saved.

# train.py
import argparse
from dataclasses import dataclass, asdict
import torch
from torch.utils.data import DataLoader
from data import DeliveryDataset, load_csv
from model import LateDeliveryNet
from trainer import Trainer

@dataclass
class Config:
    data: str = "orders_september.csv"
    batch_size: int = 64
    hidden: int = 64
    dropout: float = 0.3
    lr: float = 1e-3
    epochs: int = 100
    patience: int = 10
    seed: int = 0

def main(cfg: Config) -> None:
    torch.manual_seed(cfg.seed)                           # same seed, same run
    device = "cuda" if torch.cuda.is_available() else "mps" if torch.backends.mps.is_available() else "cpu"

    X, y = load_csv(cfg.data)
    n_val = len(X) // 5                                   # newest 20% of orders for validation
    train_set = DeliveryDataset(X[:-n_val], y[:-n_val])
    val_set = DeliveryDataset(X[-n_val:], y[-n_val:], mean=train_set.mean, std=train_set.std)

    train_loader = DataLoader(train_set, batch_size=cfg.batch_size, shuffle=True)
    val_loader = DataLoader(val_set, batch_size=256)

    model = LateDeliveryNet(hidden=cfg.hidden, dropout=cfg.dropout)   # always a fresh model
    trainer = Trainer(model, train_loader, val_loader, lr=cfg.lr, device=device)
    trainer.fit(epochs=cfg.epochs, patience=cfg.patience)
    trainer.save("late_model.pt", config=asdict(cfg), mean=train_set.mean, std=train_set.std)

if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--data", default=Config.data)
    parser.add_argument("--lr", type=float, default=Config.lr)
    args = parser.parse_args()
    main(Config(data=args.data, lr=args.lr))

Two habits are baked in here. The validation set is the newest 20% of orders, not a random 20%, because the model will always be used on orders newer than the ones it learned from. And torch.manual_seed makes runs repeatable: the same seed gives the same initial weights and the same shuffles, so when two runs differ, it is because of something you changed. A random seed is the starting number for a random number generator; fix it and the “random” sequence is the same every time.

Run it:

uv run python train.py

epoch   1  train 0.552  val 0.448
epoch   2  train 0.428  val 0.392
epoch   3  train 0.400  val 0.370
...
epoch  22  train 0.375  val 0.355
early stop; best val 0.354

Epoch 1 starts at 0.55, from a fresh model, as it should. Priya’s first request is now a command-line flag:

uv run python train.py --data orders_october.csv

The scoring script, fixed

Priya’s third request, scoring with the September model and no retraining:

# predict.py
import torch
from model import LateDeliveryNet

checkpoint = torch.load("late_model.pt")
cfg = checkpoint["config"]
model = LateDeliveryNet(hidden=cfg["hidden"], dropout=cfg["dropout"])
model.load_state_dict(checkpoint["model_state"])
model.eval()

def late_probability(distance_km, items, hour, raining, partner_rating) -> float:
    raw = torch.tensor([[distance_km, items, hour, raining, partner_rating]], dtype=torch.float32)
    x = (raw - checkpoint["mean"]) / checkpoint["std"]      # the TRAINING statistics
    with torch.no_grad():
        return torch.sigmoid(model(x)).item()

print(f"{late_probability(7.97, 9, 13, 0, 3.09):.2f}")   # order 17: 0.21
print(f"{late_probability(7.97, 9, 19, 1, 3.09):.2f}")   # 7pm, raining: 0.66

The model is rebuilt from the saved config, so if someone trains a 256-neuron version next month, the scoring script adapts automatically. And there is no line anywhere that computes statistics from the data being scored.

Where clean code still goes wrong

num_workers on macOS and Windows. With num_workers > 0, the DataLoader starts new Python processes, which re-import your script. Without the if __name__ == "__main__": guard at the bottom of train.py, each worker starts training too, and you get a confusing crash or a runaway machine. Keep the guard, always.

Datasets that load everything in __init__. Fine for a 2,500-row CSV. For 50,000 product photos, load file paths in __init__ and open each image in __getitem__. You will do exactly that in the next lesson but one.

Statistics that sneak back in. Any new feature (say, “orders already in the rider’s bag”) needs its statistics saved too. The rain bug is not a one-off; it is a category. Code review checklists at serious ML teams include “where do the normalisation statistics come from?”.

Seeds that only fix some randomness. torch.manual_seed fixes PyTorch’s generator. If you also shuffle with Python’s random or NumPy, seed those too. And GPU kernels are not always bit-for-bit deterministic, so expect tiny differences between machines even with seeds.

Building a framework instead of a model. Four small classes are enough. If you find yourself adding plugin systems and config inheritance, look at PyTorch Lightning, which is a well-tested version of this exact Trainer idea, rather than writing your own.

Old checkpoints and new code. If you rename a layer in model.py, old checkpoints fail to load with “unexpected key” errors. Store a version number in the checkpoint and keep a changelog. Module 25’s MLOps lessons turn this into a model registry.

Try it yourself

1. Add a test method to Trainer that takes any DataLoader and returns the confusion matrix and recall at a threshold, using torchmetrics from Module 27.
2. Train on September, then test on all of October without retraining (October is your test set). Then retrain on October and test on September. Do the two models agree on which features matter? (Hint: score the same five hand-picked orders with both.)
3. Deliberately reintroduce the rain bug: in predict.py, standardise a batch of 10 rainy orders using that batch’s own mean and standard deviation. How much do their late probabilities change?

Show the answer

1. Create BinaryConfusionMatrix(threshold=t) and BinaryRecall(threshold=t), call .update(torch.sigmoid(logits).squeeze(1), yb.squeeze(1).int()) for each batch inside a torch.no_grad() loop with model.eval(), then .compute() at the end.
2. Because both months come from the same underlying rule in this course’s generator, the two models should give similar probabilities, within a few points. With real data, a large disagreement is an early warning that the world changed between months.
3. With every order raining, the batch’s rain column has a standard deviation of zero, and dividing by it produces nan or infinity. With mostly rainy orders (say 8 of 10), the rain z-scores shrink and the late probabilities of the risky orders fall by roughly 10 to 20 points: in one test, 0.62 became 0.43 and 0.44 became 0.28. Orders that were already low-risk barely move, which is exactly why the bug hid so well. This is the monsoon Tuesday, reproduced.

Common mistakes

• Leaving the real logic in a notebook. Notebooks are excellent for exploring. Once something works, move it into modules and import them from the notebook if you still want to explore.
• Computing statistics inside the scoring path. Save them at training time; load them at prediction time.
• Creating the model outside main. A model created once at the top of a script or notebook gets reused by accident. Create it fresh inside the function that trains it.
• Shuffling validation data and wondering why metrics wobble. Shuffle training data only.
• No seed. Two runs that differ by 2% might be a real improvement or might be luck. With a fixed seed, you know.

Recap

• A Dataset returns one prepared example via __len__ and __getitem__; a DataLoader shuffles and batches them.
• The model is an nn.Module in its own file, configured by arguments rather than edits.
• A Trainer class owns the loop, evaluation, early stopping and saving, and moves every batch to the right device.
• A checkpoint holds the weights plus everything needed to use them: config, normalisation statistics, threshold.
• A config dataclass, a fixed seed and a time-based validation split make runs repeatable and honest.

On Thursday it rained again. The scoring script, now loading September’s statistics from the checkpoint, flagged 24 of the 29 late orders. Meera asked what had changed.

“The model was fine,” Arjun said. “The script that fed it was re-measuring ‘normal’ every day. On a rainy day, rain looked normal, so the model stopped worrying about it. Now the model always compares against the September it learned from. And retraining is one command, so nobody copies notebook cells into scripts any more.”

“Good. Then you can start on the photo sorter.”

Arjun did, that afternoon. He built a 20-layer network for pastry photos, ran train.py, and watched the loss sit at exactly 2.303 for ten epochs without moving. Priya smiled. “Your gradients are vanishing. Tomorrow.”

Go deeper

• PyTorch tutorials: the “Datasets & DataLoaders” page in “Learn the Basics” covers custom datasets in detail, and “Saving and Loading Models” covers checkpoints.
• Python dataclasses for the config pattern.