"""Why almost every modern model trains with Adam: plain SGD, momentum, Adam and AdamW, first on a long narrow valley in two parameters, then on the bouncing-ball network. Run: python3 site/ml/foundations/examples/pytorch_optimizers.py """ import math import sys from pathlib import Path import numpy as np import torch from torch import nn # ---- the valley: a loss over two parameters, 100 times steeper across than along --------------------------- START = (-4.0, 1.0) def valley_loss(w): return 0.5 * (w[0] ** 2 + 100 * w[1] ** 2) OPTIMIZERS = { "SGD": lambda params, lr: torch.optim.SGD(params, lr=lr), "momentum": lambda params, lr: torch.optim.SGD(params, lr=lr, momentum=0.9), "Adam": lambda params, lr: torch.optim.Adam(params, lr=lr), "AdamW": lambda params, lr: torch.optim.AdamW(params, lr=lr, weight_decay=0.01), } def descend(name, lr, steps=300, target=1e-3): """Run one optimizer from START. Returns the path and the first step whose loss is below `target` (or None).""" w = torch.tensor(START, dtype=torch.float64, requires_grad=True) optimizer = OPTIMIZERS[name]([w], lr) path, reached = [w.tolist()], None for step in range(steps): loss = valley_loss(w) if not math.isfinite(loss.item()) or loss.item() > 1e6: break # the steps grew instead of shrinking if reached is None and loss.item() < target: reached = step optimizer.zero_grad() loss.backward() optimizer.step() path.append(w.tolist()) return path, reached # ---- the same updates written by hand, one parameter list at a time --------------------------------------- def momentum_step(w, g, velocity, lr, beta=0.9): """velocity = beta × velocity + gradient; then step along the velocity (PyTorch's form).""" velocity = [beta * v + gi for v, gi in zip(velocity, g)] if velocity else list(g) return [wi - lr * v for wi, v in zip(w, velocity)], velocity def adam_step(w, g, m, v, t, lr, beta1=0.9, beta2=0.999, eps=1e-8, decoupled_decay=0.0): """m: running average of gradients. v: running average of squared gradients. t: the step number, from 1.""" m = [beta1 * mi + (1 - beta1) * gi for mi, gi in zip(m, g)] v = [beta2 * vi + (1 - beta2) * gi * gi for vi, gi in zip(v, g)] m_hat = [mi / (1 - beta1 ** t) for mi in m] # undo the pull toward the zero start v_hat = [vi / (1 - beta2 ** t) for vi in v] w = [wi * (1 - lr * decoupled_decay) for wi in w] # AdamW: shrink every weight by the same fraction return [wi - lr * mh / (math.sqrt(vh) + eps) for wi, mh, vh in zip(w, m_hat, v_hat)], m, v # ---- the real check: the bouncing-ball network, 4 → 16 → 4, with only the optimizer line changed ---------- LAB = Path(__file__).resolve().parents[4] / "classes/ml/labs/bouncing-ball" def ball_data(): data = np.load(LAB / "data/throws.npz") train = data["states"][data["split"] == "train"] test = data["states"][data["split"] == "test"] contact = data["contact"][data["split"] == "test"] pairs = lambda t: (torch.tensor(t[:, :-1].reshape(-1, 4), dtype=torch.float32), torch.tensor((t[:, 1:] - t[:, :-1]).reshape(-1, 4), dtype=torch.float32)) # Whole throws: 450 to train on, the last 50 training throws to choose the learning rate, 100 to test once. return pairs(train[:450]), pairs(train[450:]), pairs(test), contact.reshape(-1) def train_ball(name, lr, seed=0, epochs=240, batch=256): sys.path.insert(0, str(LAB)) from model import make_network torch.set_num_threads(1) torch.manual_seed(seed) (inputs, targets), validation, test, contact = ball_data() network = make_network(16) optimizer = OPTIMIZERS[name](network.parameters(), lr) # the only line that differs between the four runs loss_fn = nn.MSELoss() order = torch.Generator().manual_seed(seed) for epoch in range(epochs): for rows in torch.randperm(len(inputs), generator=order).split(batch): loss = loss_fn(network(inputs[rows]), targets[rows]) optimizer.zero_grad() loss.backward() optimizer.step() with torch.no_grad(): validation_loss = loss_fn(network(validation[0]), validation[1]).item() miss = network(test[0]) - test[1] air = torch.linalg.norm(miss[~contact, :2], dim=1).pow(2).mean().sqrt().item() * 1000 # mm, in flight bounce = torch.linalg.norm(miss[contact, 2:], dim=1).pow(2).mean().sqrt().item() # m/s, at the floor return {"validation_loss": validation_loss, "miss_air_mm": air, "miss_bounce": bounce} def tshirt_check(name, lr, seed=0, steps=3000, target=0.1): """The T-shirt network from Neural network: several inputs (scaled inputs, 80 customers): the first step whose training loss is below `target` (None if never), and how many of 200 new customers it sizes right at the end.""" from pytorch_choices import make_people, make_network, scale torch.set_num_threads(1) torch.manual_seed(seed) people, sizes = make_people(80, torch.Generator().manual_seed(seed)) network = make_network() optimizer = OPTIMIZERS[name](network.parameters(), lr) loss_fn, reached = nn.CrossEntropyLoss(), None for step in range(steps): loss = loss_fn(network(scale(people)), sizes) if reached is None and loss.item() < target: reached = step optimizer.zero_grad() loss.backward() optimizer.step() new_people, new_sizes = make_people(200, torch.Generator().manual_seed(99)) with torch.no_grad(): right = (network(scale(new_people)).argmax(dim=1) == new_sizes).sum().item() return {"reached": reached, "right": right} if __name__ == "__main__": print("The valley, from (-4, 1): first step with loss below 0.001") for name, lr in [("SGD", 0.019), ("momentum", 0.015), ("Adam", 0.3)]: print(f" {name:8s} lr {lr}: step {descend(name, lr)[1]}") g = [-4.0, 100.0] # the gradient at the start: (w1, 100 × w2) w, velocity = momentum_step(list(START), g, None, 0.015) print("momentum, step 1 by hand:", [round(x, 4) for x in w]) path, _ = descend("momentum", 0.015, steps=1) print("momentum, step 1 by torch.optim.SGD(momentum=0.9):", [round(x, 4) for x in path[1]]) w, m, v = adam_step(list(START), g, [0, 0], [0, 0], 1, 0.3) print("Adam, step 1 by hand:", [round(x, 4) for x in w]) print("Adam, step 1 by torch.optim.Adam:", [round(x, 4) for x in descend("Adam", 0.3, steps=1)[0][1]]) # 50 hand-written Adam steps against PyTorch's. w, m, v = list(START), [0.0, 0.0], [0.0, 0.0] for t in range(1, 51): w, m, v = adam_step(w, [w[0], 100 * w[1]], m, v, t, 0.3) print("Adam after 50 steps, by hand:", [f"{x:.6f}" for x in w], " torch:", [f"{x:.6f}" for x in descend("Adam", 0.3, steps=50)[0][50]]) print("\nThe ball (seed 0, 240 epochs): validation loss and test misses") for name, lr in [("SGD", 0.1), ("Adam", 0.001)]: # the rates chosen on the validation throws r = train_ball(name, lr) print(f" {name:5s} lr {lr}: validation {r['validation_loss']:.5f}, miss in the air {r['miss_air_mm']:.2f} mm, " f"at the bounce {r['miss_bounce']:.2f} m/s")