"""DL Week 3: Optimization and Generalization. An independent teaching example accompanying the Algorithms II Week 4 material. Run with: python -m pip install torch python dl-week-3.py Creates synthetic data locally; writes learning curves, a comparison summary, and an inference checkpoint in the current directory. """ import copy import csv import math import statistics import time import torch from torch import nn from torch.utils.data import DataLoader, TensorDataset torch.set_num_threads(1) SEEDS = (17, 29) MAX_EPOCHS = 50 PATIENCE = 6 BATCH_SIZE = 32 # Start without decay to study the optimizer/rate choices. Investigate # regularization separately; equal decay numbers have different meanings # for momentum SGD's coupled penalty and AdamW's decoupled shrinkage. WEIGHT_DECAY = 0.0 CONFIGS = ( ("momentum", 0.01), ("momentum", 0.05), ("adamw", 0.001), ("adamw", 0.01), ) # The data generator and split are independent of each model's seed. data_rng = torch.Generator().manual_seed(42) X = torch.randn(600, 20, generator=data_rng) signal = (1.4 * X[:, 0] - 0.9 * X[:, 1] + 0.6 * X[:, 2] * X[:, 3] + 0.5 * torch.randn(600, generator=data_rng)) y = (signal > 0).float().reshape(-1, 1) indices = torch.randperm(len(X), generator=data_rng) train_idx, val_idx, test_idx = indices[:360], indices[360:480], indices[480:] X_train, y_train = X[train_idx], y[train_idx] X_val, y_val = X[val_idx], y[val_idx] X_test, y_test = X[test_idx], y[test_idx] assert not set(train_idx.tolist()) & set(test_idx.tolist()) assert not set(val_idx.tolist()) & set(test_idx.tolist()) # Fit input statistics on training rows only. train_mean = X_train.mean(dim=0) train_std = X_train.std(dim=0, unbiased=False).clamp_min(1e-6) class Classifier(nn.Module): def __init__(self): super().__init__() # Stored with the weights so inference accepts raw inputs. self.register_buffer("mean", train_mean.clone()) self.register_buffer("std", train_std.clone()) self.layers = nn.Sequential( nn.Linear(20, 64), nn.ReLU(), nn.Dropout(p=0.25), nn.Linear(64, 1), ) def forward(self, features): return self.layers((features - self.mean) / self.std) loss_fn = nn.BCEWithLogitsLoss() def score(model, features, labels): model.eval() # Disable dropout during both train-set and validation scoring. with torch.inference_mode(): logits = model(features) loss = float(loss_fn(logits, labels)) accuracy = float(((logits >= 0) == labels.bool()).float().mean()) return loss, accuracy curves = [] def fit_run(name, rate, seed): torch.manual_seed(seed) model = Classifier() # Separate generator: same initial batch order for a given seed/config. loader = DataLoader( TensorDataset(X_train, y_train), batch_size=BATCH_SIZE, shuffle=True, generator=torch.Generator().manual_seed(seed + 1000), num_workers=0, ) if name == "momentum": optimizer = torch.optim.SGD( model.parameters(), lr=rate, momentum=0.9, weight_decay=WEIGHT_DECAY, ) else: optimizer = torch.optim.AdamW( model.parameters(), lr=rate, weight_decay=WEIGHT_DECAY, ) best_loss, best_epoch, waiting, best_state = math.inf, 0, 0, None started = time.perf_counter() for epoch in range(1, MAX_EPOCHS + 1): model.train() # score() switches to eval, so reset every epoch. for features, labels in loader: optimizer.zero_grad(set_to_none=True) loss = loss_fn(model(features), labels) if not torch.isfinite(loss): raise RuntimeError("Non-finite training loss; inspect this run.") loss.backward() optimizer.step() # Both curves use eval mode and the same loss definition. train_loss, _ = score(model, X_train, y_train) val_loss, val_accuracy = score(model, X_val, y_val) if not math.isfinite(val_loss): raise RuntimeError("Non-finite validation loss.") curves.append({ "optimizer": name, "lr": rate, "seed": seed, "epoch": epoch, "train_loss_eval_mode": train_loss, "validation_loss": val_loss, "validation_accuracy": val_accuracy, }) if val_loss < best_loss: best_loss, best_epoch, waiting = val_loss, epoch, 0 # A reference to state_dict() would keep changing during training. best_state = copy.deepcopy(model.state_dict()) else: waiting += 1 if waiting >= PATIENCE: break elapsed = time.perf_counter() - started assert best_state is not None model.load_state_dict(best_state) restored_loss, restored_accuracy = score(model, X_val, y_val) assert math.isclose(restored_loss, best_loss, abs_tol=1e-6) result = { "optimizer": name, "lr": rate, "seed": seed, "best_epoch": best_epoch, "epochs_run": epoch, "validation_loss": restored_loss, "validation_accuracy": restored_accuracy, "seconds": elapsed, } print(result) return model, result runs = {} summaries = [] for name, rate in CONFIGS: for seed in SEEDS: model, summary = fit_run(name, rate, seed) runs[(name, rate, seed)] = model summaries.append(summary) # Choose the configuration using mean validation loss across the fixed seeds. # Do not choose a favorable seed from its score. mean_validation = { (name, rate): statistics.mean( row["validation_loss"] for row in summaries if row["optimizer"] == name and row["lr"] == rate ) for name, rate in CONFIGS } selected = min(mean_validation, key=mean_validation.get) print("Mean validation loss by configuration:", mean_validation) print("Selected configuration:", selected) # Prespecify the first seed's retained model for the single final test report. # This is a small teaching study, not a robust ranking of optimizer families. final_model = runs[(*selected, SEEDS[0])] test_loss, test_accuracy = score(final_model, X_test, y_test) print("Final held-out test loss / accuracy:", test_loss, test_accuracy) for filename, rows in ( ("dl-week-3-curves.csv", curves), ("dl-week-3-summary.csv", summaries), ): with open(filename, "w", newline="", encoding="utf-8") as handle: writer = csv.DictWriter(handle, fieldnames=list(rows[0])) writer.writeheader() writer.writerows(rows) checkpoint = { "state_dict": final_model.state_dict(), "optimizer_name": selected[0], "learning_rate": selected[1], "seed": SEEDS[0], "feature_order": [f"x{i}" for i in range(20)], } torch.save(checkpoint, "dl-week-3.pt") reloaded = Classifier() reloaded.load_state_dict( torch.load("dl-week-3.pt", weights_only=True)["state_dict"] ) final_model.eval() reloaded.eval() with torch.inference_mode(): torch.testing.assert_close(final_model(X_val), reloaded(X_val)) print("Saved inference checkpoint; reloaded validation predictions match.") print("Versions: PyTorch", torch.__version__) print("Open dl-week-3-curves.csv in Excel; group by optimizer, lr, and seed.") # For exact training resumption, also save optimizer/scheduler state, epoch, # random-generator state, and any other state required by the input pipeline.