PyTorch: tensores, autograd y el bucle de entrenamiento.

Del fit de Keras al bucle a la vista

En el módulo 5.6, Keras entrena con una llamada: model.fit(x, y, epochs=…). Dentro hay un bucle que pide lotes, calcula la pérdida, deriva, actualiza los pesos y evalúa. PyTorch no esconde ese bucle: lo escribes tú. Es más código, y es la razón por la que casi toda la investigación en deep learning se publica en PyTorch: cuando algo falla, el fallo está en líneas que puedes leer.

El problema de la lección es el de la volatilidad del nivel 5: predecir el logaritmo de la varianza realizada de mañana de seis activos ficticios (vol5_diario.parquet) con las features del modelo HAR de Corsi (2009): la RV de hoy, su media de 5 sesiones y su media de 22, más la parte negativa del rendimiento de hoy. Los datos son sintéticos y la verdad es conocida: el generador es un HAR en el logaritmo de la varianza integrada, así que un HAR lineal es casi el óptimo y sirve de vara de medir. Una red que lo mejore mucho estaría midiendo un error.

import copy
import numpy as np
import pandas as pd
import torch
from torch import nn
from torch.utils.data import DataLoader, Dataset

torch.manual_seed(0)
torch.use_deterministic_algorithms(True)

v = pd.read_parquet("vol5_diario.parquet")
v["x"] = np.log(v["rv_5min"])

def har(g):
    """Features HAR de un activo y objetivo: log RV de mañana."""
    x = g["x"]
    return pd.DataFrame({
        "fecha": g["fecha"], "activo": g["activo"],
        "d": x, "w": x.rolling(5).mean(), "m": x.rolling(22).mean(),
        "neg": np.minimum(g["rendimiento"], 0),
        "y": x.shift(-1), "fecha_y": g["fecha"].shift(-1)})

h = pd.concat([har(g) for _, g in v.groupby("activo")]).dropna()
h["tramo"] = np.select(
    [h["fecha_y"] < "2016-01-01",
     (h["fecha"] >= "2016-01-01") & (h["fecha_y"] < "2018-01-01"),
     h["fecha"] >= "2018-01-01"],
    ["ent", "val", "prueba"], "purga")

ent = h["tramo"] == "ent"
centro = h[ent].groupby("activo")["d"].mean()
escala = h[ent].groupby("activo")["neg"].std()
for c in ["d", "w", "m", "y"]:
    h[c] = h[c] - h["activo"].map(centro)
h["neg"] = h["neg"] / h["activo"].map(escala)
print(h["tramo"].value_counts().to_string())
print(h.loc[h["tramo"] == "purga", ["fecha", "fecha_y"]]
      .drop_duplicates().to_string(index=False))

tramo
prueba    12060
ent       11958
val        3012
purga        12
     fecha    fecha_y
2015-12-31 2016-01-04
2017-12-29 2018-01-02

Tres decisiones que se repiten en todo el módulo:

• El corte se hace por la fecha de la etiqueta. Entrenamiento: objetivos anteriores a 2016; validación: 2016 y 2017; prueba: de 2018 en adelante. Las filas cuyo objetivo cruza una frontera (la del 31 de diciembre de 2015 apunta a la RV del 4 de enero de 2016) se purgan, como en la lección 5.4.2 con etiquetas más largas.
• Las estadísticas salen del entrenamiento. El centro de cada activo (la media de su log RV) y la escala de los rendimientos negativos se calculan con 2008-2015 y se aplican igual a validación y prueba. Calcularlas con toda la muestra mete el futuro en el pasado (lección 5.2.4).
• Un modelo global. Los seis activos van juntos, centrados cada uno en su nivel: 12.000 filas de entrenamiento en lugar de 2.000. Los cambios de nivel de la muestra (el bono en 2022, la cripto en 2019) quedan fuera del centro calculado con 2008-2015: el modelo tiene que seguirlos con la persistencia de las propias features.

Tensores

Un tensor es un array de n dimensiones con tres atributos que importan: dtype, shape y device. Y, si se pide, una historia de operaciones para derivar. Casi toda la API imita a NumPy; las diferencias están en los detalles.

a = h[["d", "w"]].to_numpy()[:3].copy()
t = torch.from_numpy(a)          # comparte memoria con `a`
print(t.dtype, tuple(t.shape))
t[0, 0] = 9.0
print("a[0, 0] tras cambiar t:", a[0, 0])
c = torch.tensor(a)              # copia
c[0, 0] = -1.0
print("a[0, 0] tras cambiar c:", a[0, 0])
print("de floats de Python:", torch.tensor([0.5, 1.0]).dtype)

x = torch.arange(6, dtype=torch.float32).reshape(3, 2)
print(x - x.mean(dim=0))         # (3, 2) - (2,)
print(x.T.is_contiguous(), x.device)

torch.float64 (3, 2)
a[0, 0] tras cambiar t: 9.0
a[0, 0] tras cambiar c: 9.0
de floats de Python: torch.float32
tensor([[-2., -2.],
        [ 0.,  0.],
        [ 2.,  2.]])
False cpu

• torch.from_numpy no copia: cambiar el tensor cambia el array. torch.tensor sí copia.
• Un array de NumPy es float64; un tensor creado con floats de Python es float32, el dtype por defecto de PyTorch y el de las capas. Mezclarlos no se convierte en silencio:

capa = nn.Linear(2, 1)           # pesos float32
capa(torch.from_numpy(a))        # entrada float64

RuntimeError: mat1 and mat2 must have the same dtype, but got Double and Float

• El broadcasting sigue las reglas de NumPy: (3, 2) - (2,) resta la media de cada columna.
• x.T es una vista con otros pasos de memoria (no contigua): no copia nada. Algunas operaciones (view) la rechazan; reshape o .contiguous() copian cuando hace falta.
• x.device dice dónde vive el tensor. Dos tensores en dispositivos distintos no se pueden operar juntos.

Autograd y cómo comprobarlo

Cuando un tensor tiene requires_grad=True, PyTorch apunta cada operación que lo usa en un grafo dirigido. Al llamar a backward() sobre un escalar, recorre ese grafo hacia atrás con la regla de la cadena y deja en cada hoja su derivada en .grad. Es diferenciación automática en modo inverso: el coste es del orden de una evaluación más de la función, tenga uno o un millón de parámetros.

Para verlo con una pérdida de las que se usan en volatilidad: un modelo lineal del logaritmo de la varianza, $f = X w$, evaluado con la QLIKE (la pérdida de la tabla de previsiones del README de los datos), que salvo constantes es

flowchart LR
  W[w hoja] --> F[f = X w]
  X[X datos] --> F
  F --> E[RV e^-f + f]
  E --> L[media L]
  L -. backward .-> W

cols = ["d", "w", "m", "neg"]
Xq = torch.tensor(h.loc[ent, cols].to_numpy())
Xq = torch.cat([torch.ones(len(Xq), 1, dtype=Xq.dtype), Xq], dim=1)
rv = torch.exp(torch.tensor(h.loc[ent, "y"].to_numpy()))

def qlike(w):
    f = Xq @ w
    return (rv * torch.exp(-f) + f).mean()

w = torch.tensor([0.0, 0.4, 0.3, 0.2, -0.1], dtype=torch.float64,
                 requires_grad=True)
perdida = qlike(w)
perdida.backward()
with torch.no_grad():
    f = Xq @ w
    analitico = Xq.T @ (1 - rv * torch.exp(-f)) / len(Xq)
print("autograd :", w.grad.numpy().round(6))
print("analítico:", analitico.numpy().round(6))
print("máx. diferencia:", f"{(w.grad - analitico).abs().max():.1e}")

autograd : [-0.127462  0.087152  0.045058  0.02778  -0.022228]
analítico: [-0.127462  0.087152  0.045058  0.02778  -0.022228]
máx. diferencia: 1.6e-15

Coinciden hasta el último decimal de float64. Cuando no hay derivada analítica (una red, una pérdida rara), la comprobación es con diferencias finitas centradas: $\partial L/\partial w_j \approx (L(w + h e_j) - L(w - h e_j)) / 2h$. Su error tiene dos partes que tiran en sentidos contrarios: el truncamiento, proporcional a $h^2$, y el redondeo de restar dos números casi iguales, proporcional a $\varepsilon/h$, con $\varepsilon$ la precisión de la máquina. El mejor paso está donde se cruzan, hacia $\varepsilon^{1/3}$.

def dif_finitas(w0, paso, dtype):
    X_, rv_ = Xq.to(dtype), rv.to(dtype)
    w0 = w0.detach().to(dtype)
    g = torch.zeros_like(w0)
    for j in range(len(w0)):
        e = torch.zeros_like(w0)
        e[j] = paso
        Lmas = (rv_ * torch.exp(-X_ @ (w0 + e)) + X_ @ (w0 + e)).mean()
        Lmen = (rv_ * torch.exp(-X_ @ (w0 - e)) + X_ @ (w0 - e)).mean()
        g[j] = (Lmas - Lmen) / (2 * paso)
    return g

errores = {}
for dt in (torch.float64, torch.float32):
    for k in range(1, 10):
        g = dif_finitas(w, 10.0 ** -k, dt).double()
        errores[(str(dt)[6:], f"1e-{k}")] = (g - w.grad).abs().max()
err = pd.Series(errores).unstack(0)
print(err.map(lambda e: f"{e:.1e}").to_string())
print("gradcheck:", torch.autograd.gradcheck(qlike, (w,)))

      float32  float64
1e-1  5.9e-03  5.9e-03
1e-2  5.5e-05  5.8e-05
1e-3  9.2e-05  5.8e-07
1e-4  9.5e-04  5.8e-09
1e-5  9.7e-03  6.5e-11
1e-6  8.7e-02  1.1e-10
1e-7  1.3e-01  1.2e-09
1e-8  1.3e-01  8.9e-09
1e-9  1.3e-01  1.1e-07
gradcheck: True

La figura pone la tabla en escala logarítmica. En float64 el error baja con $h^2$ (cada década de paso, dos décadas de error) hasta $h = 10^{-5}$, donde llega a 6,5·10⁻¹¹; con pasos menores, el redondeo lo vuelve a subir. En float32 el mínimo llega mucho antes, con $h = 10^{-2}$ y un error de 5,5·10⁻⁵, y con $h = 10^{-6}$ el error (0,087) ya es del tamaño del propio gradiente: la estimación es basura. De ahí la regla práctica: las comprobaciones de gradientes se hacen en float64, que es lo que hace torch.autograd.gradcheck (y por eso exige entradas de doble precisión).

Dataset y DataLoader con cortes temporales

Un Dataset responde a dos preguntas: cuántos ejemplos hay (__len__) y cuál es el i-ésimo (__getitem__). Un DataLoader los junta en lotes, los baraja si se le pide y puede cargarlos en procesos paralelos. La normalización de las features (media y desviación de cada columna) sale, otra vez, del entrenamiento.

FEATS = ["d", "w", "m", "neg"]
media = h.loc[ent, FEATS].mean()
desv = h.loc[ent, FEATS].std()

class SerieVol(Dataset):
    """Features tipificadas con el entrenamiento y log RV de mañana."""
    def __init__(self, df):
        z = (df[FEATS] - media) / desv
        self.X = torch.tensor(z.to_numpy(), dtype=torch.float32)
        self.y = torch.tensor(df["y"].to_numpy(), dtype=torch.float32)
        self.fecha_y = df["fecha_y"].to_numpy()

    def __len__(self):
        return len(self.y)

    def __getitem__(self, i):
        return self.X[i], self.y[i]

ds = {k: SerieVol(h[h["tramo"] == k]) for k in ("ent", "val", "prueba")}
print({k: len(d) for k, d in ds.items()})
print("último objetivo de ent:",
      ds["ent"].fecha_y.max().astype("datetime64[D]"))
print("primer objetivo de val:",
      ds["val"].fecha_y.min().astype("datetime64[D]"))

def cargador(semilla=0):
    return DataLoader(ds["ent"], batch_size=256, shuffle=True,
                      generator=torch.Generator().manual_seed(semilla))

xb, yb = next(iter(cargador()))
print(tuple(xb.shape), tuple(yb.shape), xb.dtype)

{'ent': 11958, 'val': 3012, 'prueba': 12060}
último objetivo de ent: 2015-12-31
primer objetivo de val: 2016-01-05
(256, 4) (256,) torch.float32

Barajar dentro del entrenamiento es correcto: las filas ya están cortadas por fecha y cada una lleva solo información de su pasado, así que el orden de los lotes no filtra nada; barajar ayuda al descenso por gradiente estocástico. Lo incorrecto es barajar antes de cortar (random_split sobre toda la muestra): es el k-fold aleatorio de la lección 5.4.1, con vecinos casi idénticos a los dos lados del corte.

El generator propio no es un detalle: hace que el orden de los lotes dependa solo de su semilla, no de cuántos números aleatorios haya consumido antes el resto del programa. La sección de Lightning muestra por qué importa.

nn.Module y un bucle de entrenamiento limpio

Un modelo es una subclase de nn.Module: en __init__ se declaran las capas (sus parámetros quedan registrados) y forward dice cómo se combinan. El bucle es siempre el mismo:

flowchart TB
  A[Lote x, y] --> B[zero_grad]
  B --> C[forward]
  C --> D[pérdida]
  D --> E[backward]
  E --> F[step]
  F --> A
  F --> G[val con no_grad]
  G --> H[mejor state_dict]

class Lineal(nn.Module):
    def __init__(self, n):
        super().__init__()
        self.capa = nn.Linear(n, 1)

    def forward(self, x):
        return self.capa(x).squeeze(-1)

class MLP(nn.Module):
    def __init__(self, n, oculta=32, p=0.1):
        super().__init__()
        self.red = nn.Sequential(
            nn.Linear(n, oculta), nn.ReLU(), nn.Dropout(p),
            nn.Linear(oculta, oculta), nn.ReLU(), nn.Dropout(p),
            nn.Linear(oculta, 1))

    def forward(self, x):
        return self.red(x).squeeze(-1)

def evaluar(modelo, d):
    modelo.eval()                     # dropout apagado
    with torch.no_grad():             # sin grafo
        return nn.functional.mse_loss(modelo(d.X), d.y).item()

def entrenar(modelo, epocas=40, lr=1e-3, wd=0.0, paciencia=8,
             semilla=0, restaurar=True):
    opt = torch.optim.Adam(modelo.parameters(), lr=lr,
                           weight_decay=wd)
    dl = cargador(semilla)
    hist, mejor, sin_mejora = [], (np.inf, None, -1), 0
    for epoca in range(epocas):
        modelo.train()                # dropout encendido
        suma = 0.0
        for xb, yb in dl:
            opt.zero_grad()
            perdida = nn.functional.mse_loss(modelo(xb), yb)
            perdida.backward()
            opt.step()
            suma += perdida.item() * len(yb)
        val = evaluar(modelo, ds["val"])
        hist.append((epoca, suma / len(ds["ent"]), val))
        if val < mejor[0]:
            mejor = (val, copy.deepcopy(modelo.state_dict()), epoca)
            sin_mejora = 0
        else:
            sin_mejora += 1
            if sin_mejora >= paciencia:
                break
    if restaurar:
        modelo.load_state_dict(mejor[1])
    return pd.DataFrame(hist, columns=["epoca", "ent", "val"]), mejor[2]

torch.manual_seed(0)
lineal = Lineal(4)
hist_lin, ep_lin = entrenar(lineal, epocas=30, lr=1e-2, paciencia=30)
torch.manual_seed(0)
mlp = MLP(4)
hist_mlp, ep_mlp = entrenar(mlp, lr=1e-3, wd=1e-4)
n_par = sum(p.numel() for p in mlp.parameters())
print(f"lineal: mejor época {ep_lin}; MLP ({n_par} parámetros): "
      f"mejor época {ep_mlp} de {len(hist_mlp)}")

lineal: mejor época 5; MLP (1249 parámetros): mejor época 8 de 17

Los detalles que separan un bucle limpio de uno que «casi funciona»:

• modelo.train() y modelo.eval() encienden y apagan el dropout (y, si hubiera, la normalización por lotes). Evaluar en modo entrenamiento da una pérdida de validación ruidosa y peor de la real.
• opt.zero_grad() antes de cada backward(), por la acumulación de gradientes.
• La parada temprana guarda una copia del mejor state_dict (copy.deepcopy: sin ella, el diccionario apunta a los mismos tensores, que se siguen actualizando) y la restaura al final.
• La época de parada se elige con la validación, así que la validación ya no es fuera de muestra para ese modelo: el resultado se mide en la prueba.

Un modelo lineal con pérdida cuadrática tiene solución cerrada (mínimos cuadrados). Si el bucle está bien, Adam tiene que llegar cerca de ella. Es una comprobación barata de todo lo anterior: datos, tipificación, bucle y optimizador.

Xe = np.c_[np.ones(len(ds["ent"])), ds["ent"].X.double().numpy()]
b_mco = np.linalg.lstsq(Xe, ds["ent"].y.double().numpy(),
                        rcond=None)[0]
b_torch = np.r_[lineal.capa.bias.item(),
                lineal.capa.weight.detach().double().numpy()[0]]
print(pd.DataFrame({"MCO": b_mco, "Adam": b_torch},
                   index=["const"] + FEATS).round(3).T.to_string())
# en las unidades del HAR: pendiente por unidad de log RV
print("HAR (MCO):", (b_mco[1:4] / desv[["d", "w", "m"]]).round(3)
      .to_dict())

      const      d      w      m    neg
MCO   0.000  0.120  0.159  0.138 -0.035
Adam -0.005  0.124  0.175  0.137 -0.028
HAR (MCO): {'d': 0.2, 'w': 0.342, 'm': 0.342}

Adam queda cerca de la solución cerrada, no encima: las pendientes difieren en la segunda cifra (0,175 frente a 0,159 en la media de 5 días), por el ruido de los lotes de 256 filas y porque la parada temprana se queda con la época 5. En el error de prueba la diferencia no se ve (sección siguiente). La última línea traduce los coeficientes de mínimos cuadrados a las unidades del HAR: 0,20 para la RV de hoy y 0,34 para cada media. El generador usa 0,36, 0,34 y 0,22 sobre el logaritmo de la varianza integrada; la RV de un día la mide con error, y ese error atenúa el coeficiente del día y pasa peso a las medias, que lo promedian. Es lo que se espera, no un fallo del bucle.

El resultado, contra la referencia, y el determinismo

La figura de la cabecera muestra las curvas de validación de los dos modelos por época. El lineal se estabiliza en cinco épocas sobre el valor de mínimos cuadrados, 0,197, y oscila alrededor por el ruido de los lotes; la MLP baja a 0,191 en su mejor época, la 8, y la parada temprana corta en la 16, tras ocho épocas sin mejorar. Falta lo que importa: la prueba, 2018-2025, que no se ha tocado para elegir nada.

def predecir(modelo, d):
    modelo.eval()
    with torch.no_grad():
        return modelo(d.X).double().numpy()

dp = ds["prueba"]
y = dp.y.double().numpy()
Xp = np.c_[np.ones(len(dp)), dp.X.double().numpy()]
pred = {"ingenua (RV de hoy)": (dp.X[:, 0].double().numpy()
                                * desv["d"] + media["d"]),
        "HAR por MCO": Xp @ b_mco,
        "lineal en PyTorch": predecir(lineal, dp),
        "MLP": predecir(mlp, dp)}
ecm = {k: np.mean((y - p) ** 2) for k, p in pred.items()}
print(pd.Series(ecm).round(3).to_string())

# diferencia de errores al cuadrado, media por semana
dif = (y - pred["HAR por MCO"]) ** 2 - (y - pred["MLP"]) ** 2
semana = pd.to_datetime(h.loc[h["tramo"] == "prueba", "fecha"]
                        ).dt.to_period("W").to_numpy()
s = pd.Series(dif).groupby(semana).mean()
print(f"MLP frente a HAR: {dif.mean() / ecm['HAR por MCO']:.1%} "
      f"menos ECM; t = {s.mean() / s.std() * np.sqrt(len(s)):.1f} "
      f"({len(s)} semanas)")

ingenua (RV de hoy)    0.319
HAR por MCO            0.213
lineal en PyTorch      0.213
MLP                    0.208
MLP frente a HAR: 2.5% menos ECM; t = 6.3 (418 semanas)

Lo que dice la tabla, en orden:

• La ingenua pierde de lejos. La RV de hoy sola tiene un error cuadrático un 50 % mayor que el HAR: la media de varios días quita buena parte del error de medida de un día.
• El lineal entrenado con Adam y el de mínimos cuadrados dan lo mismo en la prueba. El bucle hace lo que debe.
• La MLP mejora al HAR en un 2,5 %, con una diferencia semanal de t = 6,3 en 418 semanas. Es una no linealidad real de estos datos: el error de medida del logaritmo de la RV es asimétrico (un χ² con 78 grados de libertad en logaritmos) y el efecto de los rendimientos negativos no es lineal en el rendimiento. Pero es una mejora pequeña sobre un modelo de cuatro coeficientes; con datos reales habría que comprobarla en varios periodos (lección 5.4.3) y descontar cuántas arquitecturas se probaron (lección 5.4.4) antes de creérsela.

Que el resultado sea reproducible no viene gratis. Con lo fijado al principio de la lección (semilla, algoritmos deterministas, hilos fijos por el entorno) y el generador propio del DataLoader, dos entrenamientos con la misma semilla dan los mismos pesos, bit a bit; con otra semilla, otros pesos.

def pesos(m):
    return torch.cat([p.detach().flatten() for p in m.parameters()])

copias = []
for semilla in (0, 0, 1):
    torch.manual_seed(semilla)
    m = MLP(4)
    entrenar(m, epocas=5, semilla=semilla)
    copias.append(pesos(m))
print("semilla 0 dos veces:", torch.equal(copias[0], copias[1]))
print("semilla 0 frente a 1:", torch.equal(copias[0], copias[2]))

semilla 0 dos veces: True
semilla 0 frente a 1: False

GPU, MPS y torch.compile

En CPU, una red de 1.249 parámetros con 12.000 filas entrena en segundos. Las redes grandes necesitan un acelerador. PyTorch los trata como dispositivos: el modelo y cada lote se mueven con .to(dispositivo) y todo lo demás es igual.

if torch.cuda.is_available():
    dispositivo = torch.device("cuda")
elif torch.backends.mps.is_available():
    dispositivo = torch.device("mps")
else:
    dispositivo = torch.device("cpu")
print(dispositivo, torch.__version__)

cpu 2.14.1+cpu

En una máquina con GPU NVIDIA, el bucle de entrenamiento cambia en tres líneas, y se suele añadir precisión mixta (bfloat16 en las GPU recientes). El código no se ejecuta en el curso, que corre solo en CPU:

modelo = MLP(4).to(dispositivo)
opt = torch.optim.Adam(modelo.parameters(), lr=1e-3)
dl = DataLoader(ds["ent"], batch_size=4096, shuffle=True,
                pin_memory=True,
                generator=torch.Generator().manual_seed(0))
for xb, yb in dl:
    xb = xb.to(dispositivo, non_blocking=True)
    yb = yb.to(dispositivo, non_blocking=True)
    opt.zero_grad()
    with torch.autocast(device_type="cuda", dtype=torch.bfloat16):
        perdida = nn.functional.mse_loss(modelo(xb), yb)
    perdida.backward()
    opt.step()

• CUDA (NVIDIA): el acelerador de referencia, y donde el determinismo cuesta más: muchas operaciones en GPU suman en un orden que cambia entre ejecuciones. Con use_deterministic_algorithms(True), las que no tienen versión determinista lanzan un RuntimeError (la página de reproducibilidad de PyTorch pone el ejemplo de index_add_): mejor un error que un resultado distinto en silencio.
• MPS (GPU de los Mac con Apple Silicon): mismo código con "mps". No todas las operaciones ni todos los tipos están implementados: las que faltan fallan o, con PYTORCH_ENABLE_MPS_FALLBACK=1, se ejecutan en CPU. Las comprobaciones en float64 de este curso conviene hacerlas en CPU.
• Cuándo compensa: con modelos y lotes grandes. Con una red como la de esta lección, copiar los lotes a la GPU cuesta más que el cálculo.

torch.compile (PyTorch 2) traza el modelo y genera código optimizado para el dispositivo; en CPU, código C++ que compila con el compilador del sistema. La primera llamada es lenta (compila); las siguientes, más rápidas si el modelo es grande. Con Python 3.14 y PyTorch 2.14 funciona en CPU:

import os
import tempfile

# la caché de compilación, en una carpeta temporal
os.environ["TORCHINDUCTOR_CACHE_DIR"] = tempfile.mkdtemp()
compilado = torch.compile(mlp)
mlp.eval()
with torch.no_grad():
    a_eager = mlp(dp.X)
    a_comp = compilado(dp.X)
print("mismas predicciones:", torch.allclose(a_eager, a_comp,
                                             atol=1e-5))

mismas predicciones: True

Mismas predicciones (a la tolerancia de float32: el código generado puede sumar en otro orden). La compilación tardó entre 20 y 25 segundos en la máquina en la que se escribió la lección, para una red que predice los 12.000 días de prueba en milisegundos: con modelos pequeños no compensa. Tiene sentido con transformers y redes grandes, en GPU, y cuando el mismo modelo se va a ejecutar muchas veces.

La misma red con Lightning

Lightning reorganiza el código de PyTorch: el modelo, la pérdida de cada paso y el optimizador van en una LightningModule, y el bucle lo pone el Trainer, que añade lo que en un proyecto grande se acaba escribiendo igualmente: puntos de control, registro de métricas, varias GPU, precisión mixta, acumulación de gradientes. Es la base de neuralforecast, que se usa en las lecciones 5.7.2 y 5.7.3.

import logging
import warnings
import lightning as L

for n in ("pytorch_lightning", "lightning.pytorch",
          "lightning_fabric", "lightning"):
    logging.getLogger(n).setLevel(logging.ERROR)
warnings.filterwarnings("ignore", module="pytorch_lightning")
warnings.filterwarnings("ignore", module="lightning")

class ModuloVol(L.LightningModule):
    def __init__(self, red, lr=1e-3, wd=1e-4):
        super().__init__()
        self.red, self.lr, self.wd = red, lr, wd

    def training_step(self, lote, i):
        x, y = lote
        return nn.functional.mse_loss(self.red(x), y)

    def validation_step(self, lote, i):
        x, y = lote
        self.log("val", nn.functional.mse_loss(self.red(x), y))

    def configure_optimizers(self):
        return torch.optim.Adam(self.parameters(), lr=self.lr,
                                weight_decay=self.wd)

def con_lightning(epocas, gen_val):
    torch.manual_seed(0)
    red = MLP(4)
    val = DataLoader(ds["val"], batch_size=len(ds["val"]),
                     generator=gen_val)
    L.Trainer(max_epochs=epocas, accelerator="cpu",
              deterministic=True, logger=False,
              enable_progress_bar=False, enable_model_summary=False,
              enable_checkpointing=False, num_sanity_val_steps=0,
              ).fit(ModuloVol(red), cargador(0), val)
    return red

torch.manual_seed(0)
a_mano = MLP(4)
entrenar(a_mano, epocas=5, wd=1e-4, paciencia=99, restaurar=False)
for nombre, gen in [("sin generador", None),
                    ("con generador", torch.Generator())]:
    red = con_lightning(5, gen)
    print(f"Lightning, val {nombre}: pesos idénticos al bucle "
          f"a mano: {torch.equal(pesos(red), pesos(a_mano))}")

Lightning, val sin generador: pesos idénticos al bucle a mano: False
Lightning, val con generador: pesos idénticos al bucle a mano: True

Con el DataLoader de validación sin generador propio, los pesos de Lightning y los del bucle a mano coinciden en la primera época y divergen después. La causa no es Lightning: cada vez que se itera un DataLoader, su iterador saca una semilla base del generador del cargador o, si no tiene, del generador global de torch. El bucle a mano evalúa sin DataLoader; Lightning itera el de validación al final de cada época, consume un número del generador global y cambia las máscaras de dropout de la época siguiente. Con un generador propio en el de validación, los pesos son idénticos bit a bit. La lección vale para cualquier proyecto: todo lo aleatorio con su propio generador, y nada que dependa del orden en que el resto del programa consume números aleatorios.

Lo que Lightning no hace por ti: cortar los datos por fecha, purgar, normalizar con el entrenamiento ni elegir la métrica. Eso sigue en los DataLoader y en tu código, como en esta lección.

Ejercicios

1. Comprueba con diferencias finitas el gradiente de la pérdida de validación de la MLP respecto a tres pesos de su primera capa. Trabaja en float64 (copy.deepcopy(mlp).double()) y en modo eval (el dropout haría la pérdida aleatoria).

m64 = copy.deepcopy(mlp).double().eval()
Xv, yv = ds["val"].X.double(), ds["val"].y.double()
m64.zero_grad()
nn.functional.mse_loss(m64(Xv), yv).backward()
W = m64.red[0].weight
paso = 1e-5
for (i, j) in [(0, 0), (5, 2), (17, 3)]:
    with torch.no_grad():
        W[i, j] += paso
        Lmas = nn.functional.mse_loss(m64(Xv), yv).item()
        W[i, j] -= 2 * paso
        Lmen = nn.functional.mse_loss(m64(Xv), yv).item()
        W[i, j] += paso
    fd = (Lmas - Lmen) / (2 * paso)
    print(f"w[{i},{j}]: autograd {W.grad[i, j].item():+.8f}  "
          f"dif. finitas {fd:+.8f}")

w[0,0]: autograd +0.00071116  dif. finitas +0.00071116
w[5,2]: autograd -0.00029703  dif. finitas -0.00029703
w[17,3]: autograd +0.00020451  dif. finitas +0.00020451

2. Quita opt.zero_grad() del bucle y entrena el modelo lineal 10 épocas con el mismo lr. ¿Qué pasa con la pérdida de validación y por qué?

def entrenar_sin_cero(modelo, epocas, lr):
    opt = torch.optim.Adam(modelo.parameters(), lr=lr)
    vals = []
    for epoca in range(epocas):
        modelo.train()
        for xb, yb in cargador(0):
            nn.functional.mse_loss(modelo(xb), yb).backward()
            opt.step()
        vals.append(round(evaluar(modelo, ds["val"]), 3))
    return vals

torch.manual_seed(0)
print("sin zero_grad:", entrenar_sin_cero(Lineal(4), 10, 1e-2))
print("con zero_grad:", hist_lin["val"].round(3).tolist()[:10])

sin zero_grad: [0.432, 0.836, 2.191, 1.423, 2.192, 4.034, 4.052, 2.941, 9.944, 2.517]
con zero_grad: [0.221, 0.206, 0.2, 0.198, 0.198, 0.197, 0.199, 0.197, 0.198, 0.198]

3. Un compañero construye el Dataset con todas las filas, hace random_split(ds, [0.7, 0.15, 0.15]) y normaliza las features con la media de toda la muestra. Su MLP da un error de prueba algo menor que el de esta lección. ¿Qué le dirías?