4.1 梯度下降變體、SGD、Momentum、Adam（2014）#

import numpy as np

def make_valley():                      # 病態峽谷：x 方向平緩，y 方向陡峭
    f    = lambda p: 0.05 * p[0]**2 + p[1]**2
    grad = lambda p: np.array([0.1 * p[0], 2 * p[1]])
    return f, grad

def sgd(p, g, lr):      return p - lr * g
def momentum(p, v, g, lr=0.1, beta=0.9):
    v = beta * v + g
    return p - lr * v, v
def adam(p, m, v, g, t, lr=0.2, b1=0.9, b2=0.999, eps=1e-8):
    m = b1*m + (1-b1)*g
    v = b2*v + (1-b2)*g*g
    m_hat = m/(1-b1**t); v_hat = v/(1-b2**t)
    return p - lr*m_hat/(np.sqrt(v_hat)+eps), m, v

f, grad = make_valley()
for name, runner in [("SGD",      lambda: (np.array([9.,9.]), None, sgd)),
                     ("Momentum", lambda: (np.array([9.,9.]), np.zeros(2), momentum)),
                     ("Adam",     lambda: (np.array([9.,9.]), None, adam))]:
    p, extra, step = runner()
    m = v = None; path = [p.copy()]
    for t in range(1, 101):
        g = grad(p)
        if name == "SGD":      p = step(p, g, lr=0.1)
        elif name == "Momentum": p, extra = step(p, extra, g)
        else:                  p, m, v = step(p, m or 0, v or 0, g, t)
        path.append(p.copy())
    print(f"{name:9s} 終點 = {np.round(p, 4)}, loss = {f(p):.5f}")
# 觀察：SGD 在 y 方向劇烈震盪；Momentum 壓平震盪；Adam 對兩軸自動平衡

import torch, torch.nn as nn, torchvision

device = "mps" if torch.backends.mps.is_available() else "cpu"
tf = torchvision.transforms.ToTensor()
tr = torchvision.datasets.MNIST("./data", train=True, download=True, transform=tf)
tl = torch.utils.data.DataLoader(tr, batch_size=128, shuffle=True)

def net():  # 一個小 MLP
    return nn.Sequential(nn.Flatten(), nn.Linear(784, 128), nn.ReLU(),
                         nn.Linear(128, 10)).to(device)

for opt_name in ["SGD", "SGD+Momentum", "Adam"]:
    model = net()
    params = dict(sgd=(0.5, {}), **{"SGD+Momentum": (0.2, {"momentum": 0.9})},
                  **{"Adam": (1e-3, {})})[opt_name]
    opt = getattr(torch.optim, "SGD" if "SGD" in opt_name else "Adam")(
        model.parameters(), lr=params[0], **params[1])
    for epoch in range(2):
        for x, y in tl:
            x, y = x.to(device), y.to(device)
            loss = nn.functional.cross_entropy(model(x), y)
            opt.zero_grad(); loss.backward(); opt.step()
    acc = sum((model(x.to(device)).argmax(1) == y.to(device)).sum().item()
              for x, y in tl) / len(tr)
    print(f"{opt_name:14s} 2-epoch 準確率 = {acc:.4f}")
# Adam 前期最快；SGD+動量後期泛化常更佳——「快」與「好」是不同的目標

import torch, torch.nn as nn

model = nn.Linear(100, 10)
opt = torch.optim.SGD(model.parameters(), lr=0.1, momentum=0.9)
sched = torch.optim.lr_scheduler.CosineAnnealingLR(opt, T_max=100, eta_min=1e-5)

for epoch in range(100):
    # ... train_one_epoch() ...
    sched.step()          # 學習率從 0.1 餘弦衰減到 1e-5
# 暖機+餘弦退火（warmup + cosine decay）是 Transformer 訓練的標準配方：
# 前期小步避免 Adam 初期的不穩定，後期小步精修