nn.Module: la unidad de composición de toda arquitectura

1. nn.Module: todo en PyTorch es un módulo

Cada módulo en PyTorch hereda de nn.Module. Una red neuronal completa es en sí misma un módulo que contiene otros módulos (capas) — la misma estructura anidada, recursivamente, sin límite de profundidad:

class NeuralNetwork(nn.Module):
    def __init__(self):
        super().__init__()
        self.flatten = nn.Flatten()
        self.linear_relu_stack = nn.Sequential(
            nn.Linear(28*28, 512),
            nn.ReLU(),
            nn.Linear(512, 512),
            nn.ReLU(),
            nn.Linear(512, 10),
        )

    def forward(self, x):
        x = self.flatten(x)
        return self.linear_relu_stack(x)

Dos reglas fijas: el __init__ declara las capas (con super().__init__() siempre primero, para que nn.Module pueda registrar internamente todo lo que se asigne después), y forward define cómo fluyen los datos a través de ellas. Esta separación —declarar piezas, después definir el flujo— es lo que permite que el mismo módulo se pueda inspeccionar, mover a GPU, guardar y componer dentro de módulos más grandes sin cambiar nada de su código interno.

2. Por qué nunca llamás a .forward() directamente

La documentación es enfática: nunca llamar model.forward() directamente — siempre model(x). La razón es que nn.Module sobreescribe __call__, y ese __call__ ejecuta forward además de "operaciones de fondo" necesarias: hooks registrados (para inspeccionar activaciones, como hiciste con ActivationStats en fast.ai Parte 2), el modo train/eval, y el registro para autograd. Llamar a forward a mano salta todo eso silenciosamente.

3. Las piezas: Flatten, Linear, ReLU, Sequential, Softmax

Nada de esto es nuevo si venís de fast.ai Parte 2 — son los mismos bloques, con nombre oficial:

• nn.Flatten: convierte cada imagen 2D en un vector contiguo, preservando la dimensión de batch (dim=0).
• nn.Linear(in_features, out_features): la transformación lineal x@W + b, con los pesos ya como parámetros entrenables.
• nn.ReLU: la no-linealidad entre capas lineales — sin ella, apilar Linear colapsa matemáticamente en una sola transformación lineal, como ya viste en fast.ai Lección 3.
• nn.Sequential: un contenedor que pasa los datos por sus módulos internos en orden — útil para prototipar rápido, aunque el forward explícito de una clase da más control (ramas, skip connections, múltiples inputs).
• nn.Softmax: convierte logits (valores sin acotar) en probabilidades que suman 1 sobre la dimensión indicada — el mismo softmax de atención que ya usaste en el Capítulo 1 de HF.

4. parameters(): cómo el optimizador encuentra qué ajustar

Subclasear nn.Module rastrea automáticamente cada campo asignado dentro de él — no hace falta registrar nada a mano. Esto es lo que hace posible model.parameters() y model.named_parameters():

for name, param in model.named_parameters():
    print(f"Layer: {name} | Size: {param.size()}")

# Layer: linear_relu_stack.0.weight | Size: torch.Size([512, 784])
# Layer: linear_relu_stack.0.bias   | Size: torch.Size([512])

Esta es la lista exacta que se le pasa a cualquier optimizador (torch.optim.SGD(model.parameters(), lr=...)) — sin este rastreo automático, cada capa nueva agregada a un modelo requeriría actualizar manualmente la lista de parámetros a optimizar en algún otro lugar del código. Es un detalle de diseño pequeño que evita una categoría entera de bugs silenciosos ("me olvidé de agregar esta capa a la lista de parámetros entrenables").

5. Resumen

1. Una red neuronal es un nn.Module que contiene otros nn.Module — composición recursiva, sin límite de profundidad.
2. Siempre model(x), nunca model.forward(x) — __call__ ejecuta hooks y lógica de modo train/eval que forward solo no dispara.
3. Flatten, Linear, ReLU, Sequential y Softmax son los mismos bloques que ya usaste en fast.ai, con su nombre oficial en torch.nn.
4. Subclasear nn.Module rastrea automáticamente todos los parámetros — model.parameters() es lo que cualquier optimizador necesita para saber qué ajustar.