Entrenar bien: optimizadores, inicialización y regularización

import copy
import numpy as np
import matplotlib.pyplot as plt
import torch
from torch import nn
from torch.utils.data import DataLoader, TensorDataset
from torchvision import datasets

torch.manual_seed(0)
device = "cuda" if torch.cuda.is_available() else "cpu"

fashion = datasets.FashionMNIST(root="datos", train=True, download=True)
# Mezclamos una vez, con semilla, y tomamos una submuestra
orden = torch.randperm(len(fashion),
                       generator=torch.Generator().manual_seed(0))
X = fashion.data[orden].float() / 255     # píxeles entre 0 y 1
y = fashion.targets[orden]
X_ent, y_ent = X[:6000], y[:6000]
X_val, y_val = X[6000:8000], y[6000:8000]
print("Entrenamiento:", tuple(X_ent.shape),
      "  Validación:", tuple(X_val.shape))


def cargadores(lote=64, X_e=X_ent, y_e=y_ent):
    """Cargadores nuevos, con la misma mezcla en cada experimento."""
    mezcla = torch.Generator().manual_seed(0)
    ent = DataLoader(TensorDataset(X_e, y_e), batch_size=lote,
                     shuffle=True, generator=mezcla)
    val = DataLoader(TensorDataset(X_val, y_val), batch_size=1000)
    return ent, val


def mlp():
    """Un MLP como el del tema 18, pero más ancho: 784 → 256 → 10."""
    return nn.Sequential(nn.Flatten(), nn.Linear(784, 256), nn.ReLU(),
                         nn.Linear(256, 10))

import torch
from torch import nn


def evaluar(modelo, cargador, funcion_perdida, dispositivo="cpu"):
    """Pérdida media y exactitud de `modelo` sobre `cargador`."""
    modelo.eval()
    suma, aciertos, total = 0.0, 0, 0
    with torch.no_grad():
        for x, y in cargador:
            x, y = x.to(dispositivo), y.to(dispositivo)
            salida = modelo(x)
            suma += funcion_perdida(salida, y).item() * len(y)
            aciertos += (salida.argmax(dim=1) == y).sum().item()
            total += len(y)
    return suma / total, aciertos / total


def entrenar(modelo, cargador_ent, cargador_val, epocas, optimizador,
             funcion_perdida=None, planificador=None, dispositivo="cpu",
             mostrar=True):
    """Entrena `modelo` y devuelve el historial de cada época."""
    if funcion_perdida is None:
        funcion_perdida = nn.CrossEntropyLoss()
    modelo.to(dispositivo)
    historial = {"perdida_ent": [], "exactitud_ent": [],
                 "perdida_val": [], "exactitud_val": []}
    for epoca in range(1, epocas + 1):
        modelo.train()
        suma, aciertos, total = 0.0, 0, 0
        for x, y in cargador_ent:
            x, y = x.to(dispositivo), y.to(dispositivo)
            optimizador.zero_grad()
            salida = modelo(x)
            perdida = funcion_perdida(salida, y)
            perdida.backward()
            optimizador.step()
            suma += perdida.item() * len(y)
            aciertos += (salida.argmax(dim=1) == y).sum().item()
            total += len(y)
        if planificador is not None:
            planificador.step()
        perdida_val, exactitud_val = evaluar(
            modelo, cargador_val, funcion_perdida, dispositivo)
        historial["perdida_ent"].append(suma / total)
        historial["exactitud_ent"].append(aciertos / total)
        historial["perdida_val"].append(perdida_val)
        historial["exactitud_val"].append(exactitud_val)
        if mostrar:
            print(f"Época {epoca}: pérdida {suma / total:.3f}, "
                  f"validación {perdida_val:.3f}, "
                  f"exactitud de validación {exactitud_val:.3f}")
    return historial

def perdidas_por_paso(lote, pasos=200, tasa=0.05):
    """La pérdida de cada lote durante los primeros `pasos` pasos."""
    torch.manual_seed(0)
    modelo = mlp()
    optimizador = torch.optim.SGD(modelo.parameters(), lr=tasa)
    funcion_perdida = nn.CrossEntropyLoss()
    cargador_ent, _ = cargadores(lote)
    perdidas = []
    while len(perdidas) < pasos:
        for x, y_lote in cargador_ent:
            optimizador.zero_grad()
            perdida = funcion_perdida(modelo(x), y_lote)
            perdida.backward()
            optimizador.step()
            perdidas.append(perdida.item())
            if len(perdidas) == pasos:
                break
    return perdidas

lotes = [(1, "C2", "Lote de 1"), (64, "C1", "Lote de 64"),
         (6000, "C0", "Lote de 6000")]
for lote, color, nombre in lotes:
    perdidas = perdidas_por_paso(lote)
    plt.plot(perdidas, color=color, label=nombre,
             linewidth=0.8 if lote == 1 else 2)
    print(f"Lote de {lote:4d}: pérdida media de los últimos 20 pasos "
          f"{np.mean(perdidas[-20:]):.2f}")
plt.ylim(0, 5)
plt.xlabel("Paso (actualización de los pesos)")
plt.ylabel("Pérdida del lote")
plt.title("La misma tasa con tres tamaños de lote")
plt.legend(loc="upper right", frameon=True)
plt.show()

for lote in [1, 8, 64, 512, 6000]:
    torch.manual_seed(0)
    modelo = mlp()
    cargador_ent, cargador_val = cargadores(lote)
    optimizador = torch.optim.SGD(modelo.parameters(), lr=0.05)
    h = entrenar(modelo, cargador_ent, cargador_val, 1, optimizador,
                 dispositivo=device, mostrar=False)
    print(f"Lote {lote:4d}: {len(cargador_ent):4d} pasos en la época, "
          f"exactitud de validación {h['exactitud_val'][0]:.3f}")

optimizador = torch.optim.SGD(mlp().parameters(), lr=0.05, momentum=0.9)

def rosenbrock(p):
    """Un valle estrecho y curvado con el mínimo en (1, 1)."""
    x, y = p
    return (1 - x) ** 2 + 100 * (y - x ** 2) ** 2


def recorrido(crear_optimizador, pasos=500):
    """Baja por Rosenbrock desde (-1.2, 1) y devuelve el camino."""
    p = torch.tensor([-1.2, 1.0], requires_grad=True)
    optimizador = crear_optimizador([p])
    camino = [p.detach().clone()]
    for _ in range(pasos):
        optimizador.zero_grad()
        rosenbrock(p).backward()
        optimizador.step()
        camino.append(p.detach().clone())
    return torch.stack(camino)


caminos = {
    "SGD": recorrido(lambda q: torch.optim.SGD(q, lr=0.001)),
    "Momentum": recorrido(
        lambda q: torch.optim.SGD(q, lr=0.001, momentum=0.9)),
    "Adam": recorrido(lambda q: torch.optim.Adam(q, lr=0.1)),
}
for nombre, camino in caminos.items():
    x, y_ = camino[-1].tolist()
    altura = rosenbrock(camino[-1]).item()
    print(f"{nombre:8s} termina en ({x:5.2f}, {y_:5.2f}), "
          f"altura {altura:.4f}")

# El mapa: la altura en una cuadrícula, en curvas de nivel
xs = np.linspace(-2, 2, 300)
ys = np.linspace(-1, 3, 300)
XX, YY = np.meshgrid(xs, ys)
ZZ = (1 - XX) ** 2 + 100 * (YY - XX ** 2) ** 2
plt.contour(XX, YY, ZZ, levels=np.geomspace(0.1, 2000, 12),
            colors="C2", linewidths=0.8)

estilos = {"SGD": ("C0", ":"), "Momentum": ("C0", "--"),
           "Adam": ("C1", "-")}
for nombre, camino in caminos.items():
    color, linea = estilos[nombre]
    # SGD encima: su camino es corto y quedaría tapado
    plt.plot(camino[:, 0], camino[:, 1], color=color, linestyle=linea,
             linewidth=2.5, label=nombre, zorder=3 if nombre == "SGD" else 2)
plt.scatter([-1.2], [1], color="C0", s=50, zorder=3)
plt.scatter([1], [1], color="C0", marker="x", s=90, zorder=3)
plt.xlabel("x")
plt.ylabel("y")
plt.title("500 pasos por el valle de Rosenbrock")
plt.legend(loc="upper center", frameon=True)
plt.show()

optimizadores = {
    "SGD": lambda p: torch.optim.SGD(p, lr=0.05),
    "SGD+momentum": lambda p: torch.optim.SGD(p, lr=0.05, momentum=0.9),
    "Adam": lambda p: torch.optim.Adam(p, lr=0.001),
    "AdamW": lambda p: torch.optim.AdamW(p, lr=0.001, weight_decay=0.01),
}
for nombre, crear in optimizadores.items():
    torch.manual_seed(0)          # la misma red inicial para todos
    modelo = mlp()
    cargador_ent, cargador_val = cargadores()
    h = entrenar(modelo, cargador_ent, cargador_val, 10,
                 crear(modelo.parameters()), dispositivo=device,
                 mostrar=False)
    print(f"{nombre:13s} exactitud de validación: "
          f"época 1 = {h['exactitud_val'][0]:.3f}, "
          f"época 10 = {h['exactitud_val'][-1]:.3f}")

def buscar_tasa(modelo, crear_optimizador, cargador_ent,
                tasa_min=1e-5, tasa_max=10, pasos=100):
    """Sube la tasa en cada lote y anota la pérdida (suavizada)."""
    optimizador = crear_optimizador(modelo.parameters(), tasa_min)
    funcion_perdida = nn.CrossEntropyLoss()
    factor = (tasa_max / tasa_min) ** (1 / pasos)
    tasa, suave, mejor = tasa_min, None, float("inf")
    tasas, perdidas = [], []
    modelo.train()
    for x, y_lote in cargador_ent:
        for grupo in optimizador.param_groups:
            grupo["lr"] = tasa        # la tasa nueva para este paso
        optimizador.zero_grad()
        perdida = funcion_perdida(modelo(x), y_lote)
        perdida.backward()
        optimizador.step()
        # Promedio móvil: la pérdida de un lote sola es muy ruidosa
        p = perdida.item()
        suave = p if suave is None else 0.8 * suave + 0.2 * p
        tasas.append(tasa)
        perdidas.append(suave)
        mejor = min(mejor, suave)
        if suave > 4 * mejor or len(tasas) == pasos:
            break                     # la pérdida se disparó
        tasa *= factor
    return np.array(tasas), np.array(perdidas)


torch.manual_seed(0)
cargador_ent, cargador_val = cargadores()
tasas, perdidas = buscar_tasa(
    mlp(), lambda p, lr: torch.optim.SGD(p, lr=lr, momentum=0.9),
    cargador_ent)
tasa_del_minimo = tasas[np.argmin(perdidas)]
tasa_elegida = tasa_del_minimo / 10
print(f"La pérdida más baja fue con tasa {tasa_del_minimo:.3f}")
print(f"Tasa elegida (diez veces menos): {tasa_elegida:.4f}")

plt.plot(tasas, perdidas, color="C0")
plt.axvline(tasa_elegida, color="C1", label="Tasa elegida")
plt.xscale("log")
plt.xlabel("Tasa de aprendizaje (escala logarítmica)")
plt.ylabel("Pérdida (suavizada)")
plt.title("Buscador de tasa: SGD con momentum")
plt.legend()
plt.show()

def tasas_por_epoca(crear_programador, epocas=30, tasa=0.1):
    """Aplica el programador a un optimizador de prueba y anota."""
    peso = torch.zeros(1, requires_grad=True)
    optimizador = torch.optim.SGD([peso], lr=tasa)
    programador = crear_programador(optimizador)
    tasas = []
    for _ in range(epocas):
        tasas.append(optimizador.param_groups[0]["lr"])
        optimizador.step()        # en un entrenamiento real, una época
        programador.step()
    return tasas

lr_scheduler = torch.optim.lr_scheduler
programadores = {
    "StepLR": (lambda o: lr_scheduler.StepLR(o, step_size=8, gamma=0.5),
               "C0"),
    "Coseno": (lambda o: lr_scheduler.CosineAnnealingLR(o, T_max=30),
               "C1"),
    "Calentamiento + coseno": (lambda o: lr_scheduler.SequentialLR(
        o, [lr_scheduler.LinearLR(o, start_factor=0.1, total_iters=3),
            lr_scheduler.CosineAnnealingLR(o, T_max=27)],
        milestones=[3]), "C2"),
}
for nombre, (crear, color) in programadores.items():
    plt.plot(range(1, 31), tasas_por_epoca(crear), color=color,
             label=nombre)
plt.xlabel("Época")
plt.ylabel("Tasa de aprendizaje")
plt.title("Tres programadores de tasa")
plt.legend(loc="upper right")
plt.show()

for tasa in [tasa_del_minimo / 10, tasa_del_minimo / 3]:
    for usar_coseno in [False, True]:
        torch.manual_seed(0)
        modelo = mlp()
        cargador_ent, cargador_val = cargadores()
        optimizador = torch.optim.SGD(modelo.parameters(), lr=tasa,
                                      momentum=0.9)
        programador = None
        if usar_coseno:
            programador = lr_scheduler.CosineAnnealingLR(optimizador,
                                                         T_max=10)
        h = entrenar(modelo, cargador_ent, cargador_val, 10, optimizador,
                     planificador=programador, dispositivo=device,
                     mostrar=False)
        nombre = "con coseno" if usar_coseno else "tasa fija "
        print(f"Tasa {tasa:.3f} {nombre}: pérdida de validación "
              f"{h['perdida_val'][-1]:.3f}, "
              f"exactitud {h['exactitud_val'][-1]:.3f}")

torch.manual_seed(0)
red = mlp()
for capa in [red[1], red[3]]:     # las dos capas lineales
    nn.init.constant_(capa.weight, 0.01)
    nn.init.zeros_(capa.bias)
cargador_ent, cargador_val = cargadores()
optimizador = torch.optim.SGD(red.parameters(), lr=0.05, momentum=0.9)
h = entrenar(red, cargador_ent, cargador_val, 3, optimizador,
             dispositivo=device, mostrar=False)

pesos_ocultos = red[1].weight.detach()      # una fila por neurona
diferencia = (pesos_ocultos - pesos_ocultos[0]).abs().max().item()
print(f"Las 256 neuronas ocultas, tras 3 épocas: diferencia máxima "
      f"entre ellas = {diferencia:.1f}")
print(f"Exactitud de validación: {h['exactitud_val'][-1]:.3f}")

def activaciones(inicializar, capas=8, ancho=256):
    """Las salidas de cada capa tanh de una red sin entrenar."""
    torch.manual_seed(0)
    a = X_ent[:1000].reshape(1000, -1)
    salidas = []
    with torch.no_grad():
        for i in range(capas):
            capa = nn.Linear(a.shape[1], ancho)
            inicializar(capa.weight)
            nn.init.zeros_(capa.bias)
            a = torch.tanh(capa(a))
            salidas.append(a)
    return salidas

inicios = {
    "pesos chicos (σ = 0,01)": lambda w: nn.init.normal_(w, std=0.01),
    "pesos grandes (σ = 1)": lambda w: nn.init.normal_(w, std=1.0),
    "Xavier": nn.init.xavier_normal_,
}
salidas = {}
for nombre, inicializar in inicios.items():
    salidas[nombre] = activaciones(inicializar)
    primera, octava = salidas[nombre][0], salidas[nombre][-1]
    saturadas = (octava.abs() > 0.99).float().mean().item()
    print(f"{nombre:24s} desviación capa 1: {primera.std():.3f}, "
          f"capa 8: {octava.std():.3f}, saturadas: {saturadas:.0%}")

fig, ejes = plt.subplots(2, 2, figsize=(4.8, 4.8), sharex=True)
chicos = "pesos chicos (σ = 0,01)"
paneles = [(chicos, 0), (chicos, 7), ("Xavier", 0), ("Xavier", 7)]
for eje, (nombre, capa) in zip(ejes.flat, paneles):
    valores = salidas[nombre][capa].flatten().numpy()
    color = "C1" if (nombre == "Xavier" and capa == 7) else "C0"
    eje.hist(valores, bins=np.linspace(-1, 1, 41), color=color)
    eje.set_yticks([])
    titulo = "Pesos chicos" if nombre != "Xavier" else "Xavier"
    eje.set_title(f"{titulo}\ncapa {capa + 1}")
for eje in ejes[1]:
    eje.set_xlabel("Salida de tanh")
plt.show()

bn = nn.BatchNorm1d(2)            # dos neuronas
lote = torch.tensor([[1.0, 100.0], [2.0, 300.0], [3.0, 500.0]])

bn.train()
print("Entrenando:")
print(bn(lote).detach().numpy().round(2))
print("Promedios acumulados:", bn.running_mean.numpy().round(1))
bn.eval()
print("Evaluando:")
print(bn(lote).detach().numpy().round(1))

def red_profunda(con_bn):
    """Seis capas ocultas de 128 con sigmoide, con o sin BatchNorm."""
    capas, entrada = [nn.Flatten()], 784
    for _ in range(6):
        capas.append(nn.Linear(entrada, 128))
        if con_bn:
            capas.append(nn.BatchNorm1d(128))
        capas.append(nn.Sigmoid())
        entrada = 128
    capas.append(nn.Linear(128, 10))
    return nn.Sequential(*capas)

for con_bn in [False, True]:
    torch.manual_seed(0)
    modelo = red_profunda(con_bn)
    cargador_ent, cargador_val = cargadores()
    optimizador = torch.optim.SGD(modelo.parameters(), lr=0.1)
    h = entrenar(modelo, cargador_ent, cargador_val, 5, optimizador,
                 dispositivo=device, mostrar=False)
    nombre = "con BatchNorm" if con_bn else "sin BatchNorm"
    print(f"{nombre}: exactitud de validación por época",
          np.round(h["exactitud_val"], 3))

def red_grande(p_dropout=0.0):
    """Dos capas ocultas de 256, con dropout opcional."""
    return nn.Sequential(
        nn.Flatten(),
        nn.Linear(784, 256), nn.ReLU(), nn.Dropout(p_dropout),
        nn.Linear(256, 256), nn.ReLU(), nn.Dropout(p_dropout),
        nn.Linear(256, 10))

# Pocos datos: solo 1000 imágenes de entrenamiento
experimentos = {"Sin regularización": (0.0, 0.0),
                "Dropout 0,5": (0.5, 0.0),
                "Weight decay 1": (0.0, 1.0)}
regs = {}
for nombre, (p_dropout, decaimiento) in experimentos.items():
    torch.manual_seed(0)
    modelo = red_grande(p_dropout)
    cargador_poco, cargador_val = cargadores(64, X_ent[:1000],
                                             y_ent[:1000])
    optimizador = torch.optim.AdamW(modelo.parameters(), lr=0.001,
                                    weight_decay=decaimiento)
    regs[nombre] = entrenar(modelo, cargador_poco, cargador_val, 40,
                            optimizador, dispositivo=device,
                            mostrar=False)
    pv = regs[nombre]["perdida_val"]
    print(f"{nombre:18s} pérdida de validación: mínima {min(pv):.3f} "
          f"(época {np.argmin(pv) + 1}), final {pv[-1]:.3f}")

colores = {"Sin regularización": "C0", "Dropout 0,5": "C1",
           "Weight decay 1": "C2"}
epocas = range(1, 41)
for nombre, h in regs.items():
    plt.plot(epocas, h["perdida_val"], color=colores[nombre],
             label=nombre)
# La de entrenamiento, punteada, solo para la red sin regularizar
plt.plot(epocas, regs["Sin regularización"]["perdida_ent"], color="C0",
         linestyle=":")
plt.text(30, 0.2, "entrenamiento", color="C0", ha="center")
plt.ylim(0, 2)
plt.xlabel("Época")
plt.ylabel("Pérdida de validación")
plt.title("Entrenando con 1000 imágenes")
plt.legend(loc="upper right")
plt.show()

def entrenar_con_paciencia(modelo, cargador_ent, cargador_val,
                           optimizador, paciencia=5, max_epocas=100):
    """Entrena hasta que la validación no mejora en `paciencia` épocas.

    Devuelve el historial y la época de los mejores pesos, que deja
    cargados en `modelo`.
    """
    historial = {"perdida_ent": [], "exactitud_ent": [],
                 "perdida_val": [], "exactitud_val": []}
    mejor_perdida, mejor_epoca = float("inf"), 0
    mejores_pesos = copy.deepcopy(modelo.state_dict())
    for epoca in range(1, max_epocas + 1):
        h = entrenar(modelo, cargador_ent, cargador_val, 1, optimizador,
                     dispositivo=device, mostrar=False)
        for clave in historial:
            historial[clave] += h[clave]
        if h["perdida_val"][0] < mejor_perdida:
            mejor_perdida, mejor_epoca = h["perdida_val"][0], epoca
            mejores_pesos = copy.deepcopy(modelo.state_dict())
        elif epoca - mejor_epoca >= paciencia:
            print(f"Paro en la época {epoca}: {paciencia} épocas "
                  f"sin mejorar desde la época {mejor_epoca}")
            break
    modelo.load_state_dict(mejores_pesos)
    return historial, mejor_epoca


torch.manual_seed(0)
modelo = red_grande()
cargador_poco, cargador_val = cargadores(64, X_ent[:1000], y_ent[:1000])
optimizador = torch.optim.AdamW(modelo.parameters(), lr=0.001,
                                weight_decay=0.0)
h_parada, mejor_epoca = entrenar_con_paciencia(
    modelo, cargador_poco, cargador_val, optimizador, max_epocas=40)
perdida, exactitud = evaluar(modelo, cargador_val, nn.CrossEntropyLoss())
print(f"Pesos recuperados de la época {mejor_epoca}: pérdida "
      f"{perdida:.3f}, exactitud {exactitud:.3f}")

completa = regs["Sin regularización"]["perdida_val"]
plt.plot(range(1, 41), completa, color="C2", linestyle="--",
         label="Si hubiera seguido")
plt.plot(range(1, len(h_parada["perdida_val"]) + 1),
         h_parada["perdida_val"], color="C0", label="Con parada temprana")
plt.scatter([mejor_epoca], [completa[mejor_epoca - 1]], color="C1",
            s=120, zorder=3, label="Mejores pesos")
plt.axvline(len(h_parada["perdida_val"]), color="C2", linewidth=1)
plt.xlabel("Época")
plt.ylabel("Pérdida de validación")
plt.ylim(top=1.3)
plt.title("Parada temprana con paciencia 5")
plt.legend(loc="upper right")
plt.show()

torch.manual_seed(0)
modelo = mlp()
lote_chico = DataLoader(TensorDataset(X_ent[:32], y_ent[:32]),
                        batch_size=32)
optimizador = torch.optim.Adam(modelo.parameters(), lr=0.001)
# Entrenamos y "validamos" con el mismo lote: solo queremos memorizarlo
h = entrenar(modelo, lote_chico, lote_chico, 100, optimizador,
             dispositivo=device, mostrar=False)
# Hay un solo lote: cada época es un paso
for paso in [1, 10, 50, 100]:
    print(f"Paso {paso:3d}: pérdida {h['perdida_val'][paso - 1]:.4f}, "
          f"exactitud {h['exactitud_val'][paso - 1]:.2f}")