"""Train a two-layer network on MNIST and report training and validation loss per epoch. Purpose: the first training run of the course: watch the loss fall, see the validation curve, save the weights from the best epoch, and record the run in the lab notebook. Platform: all (CUDA on Tracks S and N, ROCm on Track X, MPS on Track M, or the CPU); the device is chosen automatically and printed, or forced with --device. Minimum memory: 8 GB Assumes: torch and torchvision are installed in the active environment, and about 70 MB of free disk under ./data for the MNIST download (11.6 MB compressed). Usage: python train-mnist.py [--epochs 5] [--lr 0.1] [--batch-size 128] [--hidden 256] [--train-size 50000] [--device auto|cpu|cuda|mps] [--checkpoint mnist_mlp.pt] [--label first-run] [--labbook labbook.md] [--seed 0] [--data ./data] Every epoch prints one line: mean training loss over the epoch, validation loss and accuracy at the end of the epoch, the seconds the epoch took, and "(saved)" when the validation loss is the lowest seen so far and the weights were written to --checkpoint. The test set is evaluated once, with the saved checkpoint, after the last epoch. """ import argparse import json import math import sys import time from pathlib import Path import torch from torch import nn from torch.utils.data import DataLoader, Subset, random_split from torchvision import datasets, transforms def pick_device(name: str) -> torch.device: """Return the requested device, or the best available one for "auto".""" mps = getattr(torch.backends, "mps", None) mps_ok = mps is not None and mps.is_available() if name == "auto": if torch.cuda.is_available(): return torch.device("cuda") if mps_ok: return torch.device("mps") return torch.device("cpu") if name == "cuda" and not torch.cuda.is_available(): sys.exit("train-mnist: --device cuda requested but torch.cuda.is_available() is False") if name == "mps" and not mps_ok: sys.exit("train-mnist: --device mps requested but torch.backends.mps.is_available() is False") return torch.device(name) def describe_device(device: torch.device) -> str: if device.type == "cuda": return f"cuda / {torch.cuda.get_device_name(0)}" if device.type == "mps": return "mps (Apple silicon GPU)" return "cpu" class TwoLayerNet(nn.Module): """784 inputs -> hidden units with ReLU -> 10 outputs, one per digit.""" def __init__(self, hidden: int = 256) -> None: super().__init__() self.flatten = nn.Flatten() self.layers = nn.Sequential(nn.Linear(28 * 28, hidden), nn.ReLU(), nn.Linear(hidden, 10)) def forward(self, x: torch.Tensor) -> torch.Tensor: return self.layers(self.flatten(x)) def evaluate(model: nn.Module, loader: DataLoader, device: torch.device, loss_fn: nn.Module) -> tuple[float, float]: """Average loss and accuracy over a loader, with gradients switched off.""" model.eval() total_loss, correct, seen = 0.0, 0, 0 with torch.no_grad(): for images, labels in loader: images, labels = images.to(device), labels.to(device) logits = model(images) total_loss += loss_fn(logits, labels).item() * labels.size(0) correct += (logits.argmax(dim=1) == labels).sum().item() seen += labels.size(0) return total_loss / seen, correct / seen def main() -> None: parser = argparse.ArgumentParser(description=__doc__.splitlines()[0]) parser.add_argument("--epochs", type=int, default=5) parser.add_argument("--lr", type=float, default=0.1, help="learning rate for plain SGD") parser.add_argument("--batch-size", type=int, default=128) parser.add_argument("--hidden", type=int, default=256, help="width of the hidden layer") parser.add_argument("--train-size", type=int, default=50_000, help="how many of the 50,000 training images to train on (fewer overfits sooner)") parser.add_argument("--device", default="auto", choices=["auto", "cpu", "cuda", "mps"]) parser.add_argument("--checkpoint", default="mnist_mlp.pt", help="where the best weights are saved") parser.add_argument("--label", default="", help="a short name for this run in the notebook") parser.add_argument("--labbook", default=None, help="append one JSON line per run to this file") parser.add_argument("--seed", type=int, default=0) parser.add_argument("--data", default="./data") args = parser.parse_args() if not 1 <= args.train_size <= 50_000: sys.exit("train-mnist: --train-size must be between 1 and 50000") torch.manual_seed(args.seed) device = pick_device(args.device) print(f"torch {torch.__version__}") print(f"device: {describe_device(device)}") # MNIST ships as 60,000 training and 10,000 test images. The test set is spent once, # at the end; 10,000 of the training images are held out as the validation set. to_tensor = transforms.ToTensor() # 28x28 bytes 0..255 -> float32 tensor (1, 28, 28) in 0.0..1.0 full_train = datasets.MNIST(args.data, train=True, download=True, transform=to_tensor) test_set = datasets.MNIST(args.data, train=False, download=True, transform=to_tensor) train_set, val_set = random_split(full_train, [50_000, 10_000], generator=torch.Generator().manual_seed(args.seed)) if args.train_size < 50_000: train_set = Subset(train_set, range(args.train_size)) train_loader = DataLoader(train_set, batch_size=args.batch_size, shuffle=True) val_loader = DataLoader(val_set, batch_size=1000) test_loader = DataLoader(test_set, batch_size=1000) steps = math.ceil(len(train_set) / args.batch_size) print(f"data: train {len(train_set):,} val {len(val_set):,} test {len(test_set):,} " f"batch {args.batch_size} -> {steps} steps per epoch") model = TwoLayerNet(args.hidden).to(device) n_params = sum(p.numel() for p in model.parameters()) print(f"parameters: {n_params:,}") for name, p in model.named_parameters(): print(f" {name:20s} shape {tuple(p.shape)}") loss_fn = nn.CrossEntropyLoss() optimiser = torch.optim.SGD(model.parameters(), lr=args.lr) # The untrained network: guessing among ten classes costs ln(10) = 2.3026 per image. val_loss, val_acc = evaluate(model, val_loader, device, loss_fn) print(f"epoch 0 (untrained) val loss {val_loss:.4f} val acc {val_acc:.4f}") history = [] best_val, best_epoch = float("inf"), 0 started = time.time() for epoch in range(1, args.epochs + 1): epoch_started = time.time() model.train() running, seen = 0.0, 0 for images, labels in train_loader: images, labels = images.to(device), labels.to(device) logits = model(images) # forward loss = loss_fn(logits, labels) # loss optimiser.zero_grad() loss.backward() # backward optimiser.step() # step running += loss.item() * labels.size(0) seen += labels.size(0) train_loss = running / seen val_loss, val_acc = evaluate(model, val_loader, device, loss_fn) seconds = time.time() - epoch_started history.append({"epoch": epoch, "train_loss": round(train_loss, 4), "val_loss": round(val_loss, 4), "val_acc": round(val_acc, 4), "seconds": round(seconds, 1)}) marker = "" if val_loss < best_val: # False when val_loss is nan, so a diverged epoch is never saved best_val, best_epoch = val_loss, epoch torch.save(model.state_dict(), args.checkpoint) marker = " (saved)" print(f"epoch {epoch:2d} train loss {train_loss:.4f} val loss {val_loss:.4f} " f"val acc {val_acc:.4f} {seconds:5.1f} s{marker}") elapsed = time.time() - started # The test set is looked at once, with the checkpoint from the validation minimum. if best_epoch == 0: test_loss, test_acc = float("nan"), float("nan") print(f"no epoch improved the validation loss: nothing saved to {args.checkpoint}, " f"test set not used ({elapsed:.0f} s)") else: model.load_state_dict(torch.load(args.checkpoint, map_location=device)) test_loss, test_acc = evaluate(model, test_loader, device, loss_fn) print(f"best epoch {best_epoch}: test loss {test_loss:.4f} test acc {test_acc:.4f} " f"({elapsed:.0f} s total, checkpoint {args.checkpoint})") if args.labbook: record = { "lab": "part-01/train-mnist", "label": args.label, "device": describe_device(device), "torch": torch.__version__, "epochs": args.epochs, "lr": args.lr, "batch_size": args.batch_size, "hidden": args.hidden, "train_size": len(train_set), "parameters": n_params, "best_epoch": best_epoch, "test_acc": None if math.isnan(test_acc) else round(test_acc, 4), "seconds": round(elapsed, 1), "checkpoint": args.checkpoint, "history": history, } with Path(args.labbook).open("a", encoding="utf-8") as fh: fh.write(json.dumps(record) + "\n") print(f"recorded in {args.labbook}") if __name__ == "__main__": main()