PyTorch: tensors, autograd and the training loop.

From Keras's fit to a loop in plain sight

In module 5.6, Keras trains with one call: model.fit(x, y, epochs=…). Inside there is a loop that requests batches, computes the loss, differentiates, updates the weights and evaluates. PyTorch does not hide that loop: you write it. It is more code, and it is the reason almost all deep learning research is published in PyTorch: when something fails, the failure is in lines you can read.

The problem of the lesson is the level 5 volatility one: predict the logarithm of tomorrow's realized variance for six fictitious assets (vol5_diario.parquet) with the features of Corsi's (2009) HAR model: today's RV, its 5-session mean and its 22-session mean, plus the negative part of today's return. The data are synthetic and the truth is known: the generator is a HAR on the logarithm of integrated variance, so a linear HAR is almost optimal and serves as the yardstick. A network that beat it by a wide margin would be measuring a mistake.

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):
    """HAR features of one asset and target: tomorrow's log RV."""
    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

Three decisions that recur throughout the module (in the code, ent, val and prueba are the training, validation and test segments, and purga the purged rows):

• The split is made by the label date. Training: targets before 2016; validation: 2016 and 2017; test: 2018 onward. Rows whose target crosses a boundary (the one of December 31, 2015 points to the RV of January 4, 2016) are purged, as in lesson 5.4.2 with longer labels.
• The statistics come from training. The center of each asset (the mean of its log RV) and the scale of negative returns are computed with 2008-2015 and applied unchanged to validation and test. Computing them with the whole sample puts the future into the past (lesson 5.2.4).
• One global model. The six assets go together, each centered on its own level: 12,000 training rows instead of 2,000. The level shifts in the sample (the bond in 2022, crypto in 2019) are outside the center computed with 2008-2015: the model has to follow them through the persistence of the features themselves.

Tensors

A tensor is an n-dimensional array with three attributes that matter: dtype, shape and device. And, on request, a history of operations for differentiation. Almost the whole API imitates NumPy; the differences are in the details.

a = h[["d", "w"]].to_numpy()[:3].copy()
t = torch.from_numpy(a)          # shares memory with `a`
print(t.dtype, tuple(t.shape))
t[0, 0] = 9.0
print("a[0, 0] after changing t:", a[0, 0])
c = torch.tensor(a)              # copy
c[0, 0] = -1.0
print("a[0, 0] after changing c:", a[0, 0])
print("from Python floats:", 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] after changing t: 9.0
a[0, 0] after changing c: 9.0
from Python floats: torch.float32
tensor([[-2., -2.],
        [ 0.,  0.],
        [ 2.,  2.]])
False cpu

• torch.from_numpy does not copy: changing the tensor changes the array. torch.tensor does copy.
• A NumPy array is float64; a tensor created from Python floats is float32, PyTorch's default dtype and that of the layers. Mixing them is not converted silently:

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

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

• Broadcasting follows NumPy's rules: (3, 2) - (2,) subtracts the mean of each column.
• x.T is a view with different memory strides (not contiguous): it copies nothing. Some operations (view) reject it; reshape or .contiguous() copy when needed.
• x.device says where the tensor lives. Two tensors on different devices cannot be operated on together.

Autograd and how to check it

When a tensor has requires_grad=True, PyTorch records each operation that uses it in a directed graph. When backward() is called on a scalar, it walks that graph backward with the chain rule and leaves in each leaf its derivative in .grad. It is reverse-mode automatic differentiation: the cost is of the order of one more evaluation of the function, whether it has one parameter or a million.

To see it with one of the losses used in volatility: a linear model of the logarithm of variance, $f = X w$, evaluated with the QLIKE (the loss of the forecast table in the data README), which up to constants is

flowchart LR
  W[w leaf] --> F[f = X w]
  X[X data] --> F
  F --> E[RV e^-f + f]
  E --> L[mean 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("analytical:", analitico.numpy().round(6))
print("max. difference:", f"{(w.grad - analitico).abs().max():.1e}")

autograd  : [-0.127462  0.087152  0.045058  0.02778  -0.022228]
analytical: [-0.127462  0.087152  0.045058  0.02778  -0.022228]
max. difference: 1.6e-15

They agree to the last decimal of float64. When there is no analytical derivative (a network, an unusual loss), the check uses central finite differences: $\partial L/\partial w_j \approx (L(w + h e_j) - L(w - h e_j)) / 2h$. Its error has two parts pulling in opposite directions: truncation, proportional to $h^2$, and the rounding from subtracting two nearly equal numbers, proportional to $\varepsilon/h$, with $\varepsilon$ the machine precision. The best step is where they cross, around $\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

The figure puts the table on a log scale. In float64 the error falls with $h^2$ (each decade of step, two decades of error) down to $h = 10^{-5}$, where it reaches 6.5·10⁻¹¹; with smaller steps, rounding pushes it back up. In float32 the minimum comes much earlier, at $h = 10^{-2}$ with an error of 5.5·10⁻⁵, and at $h = 10^{-6}$ the error (0.087) is already the size of the gradient itself: the estimate is garbage. Hence the rule of thumb: gradient checks are done in float64, which is what torch.autograd.gradcheck does (and why it requires double-precision inputs).

Dataset and DataLoader with time splits

A Dataset answers two questions: how many examples there are (__len__) and which one is the i-th (__getitem__). A DataLoader groups them into batches, shuffles them if asked and can load them in parallel processes. The normalization of the features (mean and standard deviation of each column) comes, once again, from training.

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

class SerieVol(Dataset):
    """Features standardized with training and tomorrow's log RV."""
    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("last target in ent:",
      ds["ent"].fecha_y.max().astype("datetime64[D]"))
print("first target in 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}
last target in ent: 2015-12-31
first target in val: 2016-01-05
(256, 4) (256,) torch.float32

Shuffling within training is correct: the rows are already split by date and each one carries only information from its past, so the batch order leaks nothing; shuffling helps stochastic gradient descent. What is wrong is shuffling before splitting (random_split over the whole sample): it is the random k-fold of lesson 5.4.1, with almost identical neighbors on both sides of the split.

The dedicated generator is not a detail: it makes the batch order depend only on its seed, not on how many random numbers the rest of the program has consumed before. The Lightning section shows why that matters.

nn.Module and a clean training loop

A model is a subclass of nn.Module: the layers are declared in __init__ (their parameters are registered) and forward says how they combine. The loop is always the same:

flowchart TB
  A[Batch x, y] --> B[zero_grad]
  B --> C[forward]
  C --> D[loss]
  D --> E[backward]
  E --> F[step]
  F --> A
  F --> G[val with no_grad]
  G --> H[best 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 off
    with torch.no_grad():             # no graph
        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 on
        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"linear: best epoch {ep_lin}; MLP ({n_par} parameters): "
      f"best epoch {ep_mlp} of {len(hist_mlp)}")

linear: best epoch 5; MLP (1249 parameters): best epoch 8 of 17

The details that separate a clean loop from one that “almost works”:

• modelo.train() and modelo.eval() switch dropout on and off (and, if there were any, batch normalization). Evaluating in training mode gives a noisy validation loss that is worse than the real one.
• opt.zero_grad() before each backward(), because of gradient accumulation.
• Early stopping keeps a copy of the best state_dict (copy.deepcopy: without it, the dictionary points to the same tensors, which keep being updated) and restores it at the end.
• The stopping epoch is chosen with validation, so validation is no longer out of sample for that model: the result is measured on the test set.

A linear model with squared loss has a closed-form solution (least squares). If the loop is right, Adam has to get close to it. It is a cheap check of everything above: data, standardization, loop and optimizer.

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({"OLS": b_mco, "Adam": b_torch},
                   index=["const"] + FEATS).round(3).T.to_string())
# in HAR units: slope per unit of log RV
print("HAR (OLS):", (b_mco[1:4] / desv[["d", "w", "m"]]).round(3)
      .to_dict())

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

Adam ends up close to the closed-form solution, not on top of it: the slopes differ in the second digit (0.175 versus 0.159 on the 5-day mean), because of the noise of 256-row batches and because early stopping keeps epoch 5. In the test error the difference does not show (next section). The last line translates the least squares coefficients into HAR units: 0.20 for today's RV and 0.34 for each mean. The generator uses 0.36, 0.34 and 0.22 on the logarithm of integrated variance; one day's RV measures it with error, and that error attenuates the daily coefficient and shifts weight to the means, which average it out. It is what is expected, not a bug in the loop.

The result, against the benchmark, and determinism

The figure at the top shows the validation curves of the two models by epoch. The linear model settles in five epochs on the least squares value, 0.197, and oscillates around it because of batch noise; the MLP goes down to 0.191 at its best epoch, the 8th, and early stopping cuts at the 16th, after eight epochs without improvement. What matters is still missing: the test set, 2018-2025, which has not been touched to choose anything.

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 = {"naive (today's RV)": (dp.X[:, 0].double().numpy()
                               * desv["d"] + media["d"]),
        "HAR by OLS": Xp @ b_mco,
        "linear in 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())

# difference of squared errors, mean per week
dif = (y - pred["HAR by OLS"]) ** 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 versus HAR: {dif.mean() / ecm['HAR by OLS']:.1%} "
      f"lower MSE; t = {s.mean() / s.std() * np.sqrt(len(s)):.1f} "
      f"({len(s)} weeks)")

naive (today's RV)    0.319
HAR by OLS            0.213
linear in PyTorch     0.213
MLP                   0.208
MLP versus HAR: 2.5% lower MSE; t = 6.3 (418 weeks)

What the table says, in order:

• The naive forecast loses by far. Today's RV alone has a squared error 50% larger than the HAR: the mean over several days removes a good part of the measurement error of a single day.
• The linear model trained with Adam and the least squares one give the same on the test set. The loop does what it should.
• The MLP beats the HAR by 2.5%, with a weekly difference of t = 6.3 over 418 weeks. It is a real nonlinearity in these data: the measurement error of the log RV is asymmetric (a χ² with 78 degrees of freedom in logarithms) and the effect of negative returns is not linear in the return. But it is a small improvement over a four-coefficient model; with real data it would have to be checked over several periods (lesson 5.4.3), discounting how many architectures were tried (lesson 5.4.4), before believing it.

A reproducible result does not come for free. With what was fixed at the beginning of the lesson (seed, deterministic algorithms, threads fixed by the environment) and the DataLoader's dedicated generator, two training runs with the same seed give the same weights, bit for bit; with another seed, different weights.

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("seed 0 twice:", torch.equal(copias[0], copias[1]))
print("seed 0 versus 1:", torch.equal(copias[0], copias[2]))

seed 0 twice: True
seed 0 versus 1: False

GPU, MPS and torch.compile

On CPU, a 1,249-parameter network with 12,000 rows trains in seconds. Large networks need an accelerator. PyTorch treats them as devices: the model and each batch are moved with .to(device) and everything else is the same.

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

On a machine with an NVIDIA GPU, the training loop changes in three lines, and mixed precision is usually added (bfloat16 on recent GPUs). The code is not run in the course, which runs on CPU only:

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): the reference accelerator, and where determinism costs the most: many GPU operations sum in an order that changes between runs. With use_deterministic_algorithms(True), those without a deterministic version raise a RuntimeError (PyTorch's reproducibility page gives the example of index_add_): better an error than a silently different result.
• MPS (the GPU of Macs with Apple Silicon): same code with "mps". Not all operations or types are implemented: the missing ones fail or, with PYTORCH_ENABLE_MPS_FALLBACK=1, run on CPU. The float64 checks in this course are best done on CPU.
• When it pays off: with large models and batches. With a network like the one in this lesson, copying the batches to the GPU costs more than the computation.

torch.compile (PyTorch 2) traces the model and generates code optimized for the device; on CPU, C++ code that is built with the system compiler. The first call is slow (it compiles); the following ones, faster if the model is large. With Python 3.14 and PyTorch 2.14 it works on CPU:

import os
import tempfile

# the compilation cache, in a temporary folder
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("same predictions:", torch.allclose(a_eager, a_comp,
                                          atol=1e-5))

same predictions: True

Same predictions (to float32 tolerance: the generated code may sum in a different order). Compilation took between 20 and 25 seconds on the machine the lesson was written on, for a network that predicts the 12,000 test days in milliseconds: with small models it does not pay off. It makes sense with transformers and large networks, on GPU, and when the same model is going to be run many times.

The same network with Lightning

Lightning reorganizes PyTorch code: the model, the loss of each step and the optimizer go into a LightningModule, and the loop is supplied by the Trainer, which adds what a large project ends up writing anyway: checkpoints, metric logging, multiple GPUs, mixed precision, gradient accumulation. It is the base of neuralforecast, which is used in lessons 5.7.2 and 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 [("without generator", None),
                    ("with generator", torch.Generator())]:
    red = con_lightning(5, gen)
    print(f"Lightning, val {nombre}: weights identical to the "
          f"handwritten loop: {torch.equal(pesos(red), pesos(a_mano))}")

Lightning, val without generator: weights identical to the handwritten loop: False
Lightning, val with generator: weights identical to the handwritten loop: True

With the validation DataLoader without its own generator, the Lightning weights and those of the handwritten loop agree in the first epoch and diverge afterward. The cause is not Lightning: each time a DataLoader is iterated, its iterator draws a base seed from the loader's generator or, if it has none, from torch's global generator. The handwritten loop evaluates without a DataLoader; Lightning iterates the validation one at the end of each epoch, consumes a number from the global generator and changes the dropout masks of the next epoch. With a dedicated generator in the validation loader, the weights are identical bit for bit. The lesson applies to any project: everything random with its own generator, and nothing that depends on the order in which the rest of the program consumes random numbers.

What Lightning does not do for you: split the data by date, purge, normalize with training statistics or choose the metric. That stays in the DataLoaders and in your code, as in this lesson.

Exercises

1. Check with finite differences the gradient of the MLP's validation loss with respect to three weights of its first layer. Work in float64 (copy.deepcopy(mlp).double()) and in eval mode (dropout would make the loss random).

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"finite diff. {fd:+.8f}")

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

2. Remove opt.zero_grad() from the loop and train the linear model for 10 epochs with the same lr. What happens to the validation loss and why?

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("without zero_grad:", entrenar_sin_cero(Lineal(4), 10, 1e-2))
print("with zero_grad:   ", hist_lin["val"].round(3).tolist()[:10])

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

3. A colleague builds the Dataset with all the rows, does random_split(ds, [0.7, 0.15, 0.15]) and normalizes the features with the mean of the whole sample. Their MLP gives a somewhat lower test error than the one in this lesson. What would you tell them?