# Source: content/notes/math/optimization.md
# Independent CPU example; use the curriculum environment.
# See /notes/ml/#example-environment or /notes/deep-learning/#example-environment.

import math
import torch
from torch import nn

torch.manual_seed(17)
torch.set_num_threads(1)
X = torch.randn(23, 3)
y = X @ torch.tensor([1., -2., .5]) + .1
model = nn.Linear(3, 1)
optimizer = torch.optim.AdamW([
    {"params": [model.weight], "weight_decay": .01},
    {"params": [model.bias], "weight_decay": 0.},
], lr=.01)
batches = [(X[i:i+4], y[i:i+4]) for i in range(0, len(X), 4)]
accum_steps, epochs = 4, 10
total_steps = epochs * math.ceil(len(batches) / accum_steps)
scheduler = torch.optim.lr_scheduler.OneCycleLR(
    optimizer, max_lr=.05, total_steps=total_steps,
    pct_start=.3, cycle_momentum=False,
)
before = nn.functional.mse_loss(model(X).squeeze(1), y).item()
updates = 0
for _ in range(epochs):
    for start in range(0, len(batches), accum_steps):
        window = batches[start:start+accum_steps]
        count = sum(len(xb) for xb, _ in window)
        optimizer.zero_grad(set_to_none=True)
        for xb, yb in window:
            loss = nn.functional.mse_loss(
                model(xb).squeeze(1), yb, reduction="sum"
            ) / count
            loss.backward()
        nn.utils.clip_grad_norm_(model.parameters(), 1.)
        optimizer.step()
        scheduler.step()
        updates += 1
after = nn.functional.mse_loss(model(X).squeeze(1), y).item()
assert updates == total_steps
assert after < before
print("optimizer updates:", updates, "MSE:", before, "->", after)