The Training Loop: Gradient Descent in Practice

# Chapter 7: the training loop. Forward, zero the grads, backward, step, and repeat.
# Uses the Tensor engine from Chapter 6: save that chapter's code down to (not including) the line
# '# === 1.' as tensor.py next to this file. That's unbroadcast(), the Tensor class and count().
import numpy as np
from tensor import Tensor

km = np.array([2.0, 5.0, 8.0, 12.0])
minutes = np.array([11.0, 21.0, 28.0, 42.0])


def train(x, y, lr, steps, batch_size=None, seed=0):
    """Gradient descent on the line w*x + b. Returns the final knobs and the full-data loss after every update."""
    w, b = Tensor(0.0), Tensor(0.0)                     # start knowing nothing
    params = [w, b]
    rng = np.random.default_rng(seed)
    n = len(x)
    batch_size = batch_size or n                        # default: the whole dataset every step
    history = [float(np.mean((w.data * x + b.data - y) ** 2))]
    updates = 0
    while updates < steps:
        order = rng.permutation(n)                      # one epoch = one pass through the data, shuffled
        for start in range(0, n, batch_size):
            rows = order[start:start + batch_size]
            loss = ((w * x[rows] + b - y[rows]) ** 2).mean()        # 1. forward
            for p in params:                                        # 2. zero the grads
                p.grad = np.zeros_like(p.data)
            loss.backward()                                         # 3. backward
            for p in params:                                        # 4. step downhill
                p.data = p.data - lr * p.grad
            history.append(float(np.mean((w.data * x + b.data - y) ** 2)))
            updates += 1
            if updates == steps or not np.isfinite(history[-1]):
                return float(w.data), float(b.data), history
    return float(w.data), float(b.data), history


# === 1. The best line, in closed form, to know what we're aiming for ==============
w_best = np.mean((km - km.mean()) * (minutes - minutes.mean())) / np.var(km)
b_best = minutes.mean() - w_best * km.mean()
loss_best = np.mean((w_best * km + b_best - minutes) ** 2)
print(f"best line: w* = {w_best:.4f}, b* = {b_best:.4f}, loss = {loss_best:.4f}")

# === 2. Four learning rates, 2,000 steps each ======================================
print("\n   lr      w        b       loss after 2,000   first step within 1% of best")
for lr in [0.001, 0.01, 0.016, 0.017]:
    w, b, hist = train(km, minutes, lr, 2000)
    near = next((i for i, v in enumerate(hist) if v < 1.01 * loss_best), None)
    print(f"{lr:6}  {w:8.4g} {b:8.4g}   {hist[-1]:12.4g}        {near}")

# === 3. Why 0.017 explodes: the curvature of the bowl ===============================
X = np.column_stack([km, np.ones(4)])
H = 2 / len(km) * X.T @ X                          # second derivatives of the MSE in (w, b)
lam = np.linalg.eigvalsh(H)
print(f"\ncurvatures (eigenvalues of H): {lam.round(3)}, condition number {lam.max() / lam.min():.0f}")
print(f"largest stable learning rate: 2 / {lam.max():.2f} = {2 / lam.max():.4f}")

# === 4. Standardize the input: the bowl becomes round ================================
mu, sd = km.mean(), km.std()
km_std = (km - mu) / sd
for lr in [0.1, 0.5, 0.9]:
    w_s, b_s, hist = train(km_std, minutes, lr, 30)
    near = next((i for i, v in enumerate(hist) if v < 1.01 * loss_best), None)
    print(f"standardized, lr {lr}: within 1% after {near} steps;"
          f" back in km: w = {w_s / sd:.4f}, b = {b_s - w_s * mu / sd:.4f}")

# === 5. More data: a year of rides, and mini-batches ==================================
rng = np.random.default_rng(7)
year_km = np.round(rng.uniform(1, 15, 240), 1)
year_min = np.round(3 * year_km + 5 + rng.normal(0, 1.5, 240), 1)   # 3 min/km + 5, plus noise
year_std = (year_km - year_km.mean()) / year_km.std()
Xy = np.column_stack([year_std, np.ones(240)])
best = np.linalg.lstsq(Xy, year_min, rcond=None)[0]
print(f"\nyear of rides: best loss {np.mean((Xy @ best - year_min) ** 2):.3f}, starting loss {np.mean(year_min ** 2):.3f}")

print("batch  updates   loss after each of 10 epochs")
for B in [240, 16, 1]:
    steps = 10 * (240 // B)                       # 10 epochs, whatever the batch size
    w_s, b_s, hist = train(year_std, year_min, 0.05, steps, batch_size=B)
    per_epoch = hist[240 // B::240 // B]
    print(f"{B:5}  {steps:7}   {[round(v, 3) for v in per_epoch]}")
w_s, b_s, _ = train(year_std, year_min, 0.05, 150, batch_size=16)
print(f"B = 16, back in km: w = {w_s / year_km.std():.3f}, b = {b_s - w_s * year_km.mean() / year_km.std():.3f}")