Gradient Descent

"""Gradient Descent thuần bằng NumPy.

Ví dụ với hàm 1D: L(w) = (w - 2)^2 + 1
Cực tiểu tại w = 2, L = 1.
"""
import numpy as np


def loss_fn(w: float) -> float:
    return (w - 2.0) ** 2 + 1.0


def grad_fn(w: float) -> float:
    return 2.0 * (w - 2.0)


def gradient_descent(
    w_init: float,
    lr: float = 0.1,
    steps: int = 200,
    tol: float = 1e-6,
) -> tuple[float, list[float]]:
    """Trả về (w_cuối, lịch sử loss)."""
    w = w_init
    history: list[float] = [loss_fn(w)]
    for step in range(steps):
        g = grad_fn(w)
        w = w - lr * g
        history.append(loss_fn(w))
        if abs(g) < tol:
            print(f"Hội tụ tại bước {step}, |grad| = {abs(g):.2e}")
            break
    return w, history


if __name__ == "__main__":
    w_star, hist = gradient_descent(w_init=6.0, lr=0.1, steps=100)
    print(f"w* = {w_star:.6f}, loss = {loss_fn(w_star):.6f}")

"""Huấn luyện một MLP nhỏ với 4 optimizer, so sánh đường hội tụ.

Yêu cầu: torch >= 2.0
"""
import torch
import torch.nn as nn
from torch.utils.data import DataLoader, TensorDataset


# ---------------------------------------------------------
# 1. Dữ liệu giả lập — hồi quy y = 3x + nhiễu
# ---------------------------------------------------------
torch.manual_seed(42)
X = torch.randn(1000, 4)
y = X @ torch.tensor([3.0, -1.5, 2.0, 0.5]) + 0.1 * torch.randn(1000)
loader = DataLoader(TensorDataset(X, y), batch_size=32, shuffle=True)


# ---------------------------------------------------------
# 2. Kiến trúc MLP đơn giản
# ---------------------------------------------------------
def build_model() -> nn.Module:
    return nn.Sequential(
        nn.Linear(4, 16),
        nn.ReLU(),
        nn.Linear(16, 1),
    )


# ---------------------------------------------------------
# 3. Bốn optimizer, cùng learning rate
# ---------------------------------------------------------
def make_optimizers(model: nn.Module, lr: float = 1e-2):
    return {
        "SGD":      torch.optim.SGD(model.parameters(), lr=lr),
        "Momentum": torch.optim.SGD(model.parameters(), lr=lr, momentum=0.9),
        "RMSprop":  torch.optim.RMSprop(model.parameters(), lr=lr, alpha=0.99),
        "Adam":     torch.optim.Adam(model.parameters(), lr=lr,
                                     betas=(0.9, 0.999), eps=1e-8),
    }


# ---------------------------------------------------------
# 4. Vòng huấn luyện + LR scheduler
# ---------------------------------------------------------
def train_one(name: str, model: nn.Module, opt: torch.optim.Optimizer,
              epochs: int = 20):
    loss_fn = nn.MSELoss()
    scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(opt, T_max=epochs)
    losses: list[float] = []

    for epoch in range(epochs):
        epoch_loss = 0.0
        for xb, yb in loader:
            opt.zero_grad(set_to_none=True)
            pred = model(xb).squeeze(-1)
            loss = loss_fn(pred, yb)
            loss.backward()                 # tính gradient = backpropagation
            opt.step()                      # cập nhật trọng số
            epoch_loss += loss.item() * len(xb)
        scheduler.step()                    # cập nhật learning rate
        epoch_loss /= len(loader.dataset)
        losses.append(epoch_loss)
        if epoch % 5 == 0:
            lr_now = scheduler.get_last_lr()[0]
            print(f"[{name}] epoch {epoch:3d}  loss={epoch_loss:.4f}  lr={lr_now:.5f}")
    return losses


if __name__ == "__main__":
    for name in ["SGD", "Momentum", "RMSprop", "Adam"]:
        model = build_model()
        opt = make_optimizers(model)[name]
        train_one(name, model, opt)