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.

El mecanismo detrás de "registrar" es concreto: nn.Module sobreescribe __setattr__. Cada vez que hacés self.algo = valor, mira qué es valor: si es un nn.Module lo guarda en self._modules, si es un nn.Parameter en self._parameters, y si no, es un atributo de Python común. De ahí salen parameters(), state_dict(), .to(), train() — todos recorren esos diccionarios recursivamente. Y de ahí sale el error más clásico al olvidar super().__init__(): AttributeError: cannot assign module before Module.__init__() call, porque los diccionarios donde iría la capa todavía no existen.

Imprimir un módulo muestra el árbol completo, con los nombres que vos elegiste y los índices de los contenedores — los mismos nombres que después aparecen en named_parameters() y en el state_dict:

print(NeuralNetwork())

Cuando en el Capítulo 1 de HF hiciste print(model) sobre un transformer y viste 32 bloques idénticos con self_attn, mlp, input_layernorm, era exactamente esto: un nn.Module con un nn.ModuleList de bloques, cada bloque un módulo con submódulos.

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.

Vale precisar qué salta y qué no. Autograd registra las operaciones en el momento en que se ejecutan sobre tensores con requires_grad, así que forward() a secas sí produce gradientes; y training es un atributo que forward lee, así que Dropout también se comporta según el modo. Lo que sí se pierde son los hooks — los forward pre-hooks, forward hooks y backward hooks que vos, una librería (PEFT para inyectar LoRA, Accelerate para offloading, herramientas de profiling y de interpretabilidad) o torch.compile hayan registrado sobre ese módulo. Como no sabés qué hooks tiene un modelo ajeno, la regla es absoluta: siempre model(x). La sección 6 muestra cómo se registran.

3. Las piezas: Flatten, Linear, ReLU, Sequential — y dónde no va 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.T + b, con los pesos ya como parámetros entrenables. El weight se guarda como (out, in) — por eso en el ejemplo de §4 la primera capa muestra [512, 784] y no [784, 512].
• 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. Es stateless (no tiene parámetros), así que da lo mismo usar una instancia o F.relu; los transformers modernos usan variantes suaves (GELU, SiLU/SwiGLU) por la misma razón.
• 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). Acepta un OrderedDict si querés nombres en vez de índices.
• nn.Softmax(dim): 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. El tutorial oficial lo aplica afuera del modelo, sobre los logits, y solo para leer la predicción.

logits = model(X)                          # (B, 10), sin normalizar
loss = nn.CrossEntropyLoss()(logits, y)    # log_softmax + NLL adentro
probs = logits.softmax(dim=1)              # solo para mostrar / calibrar
pred = logits.argmax(dim=1)                # la clase, sin necesidad de softmax

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").

Un nn.Parameter es un tensor con requires_grad=True por defecto y la marca de "esto se registra". Contar cuántos hay, y cuántos se entrenan, es la línea que más veces vas a escribir:

total = sum(p.numel() for p in model.parameters())
entrenables = sum(p.numel() for p in model.parameters() if p.requires_grad)
print(f"{total/1e6:.2f}M parámetros, {entrenables/1e6:.2f}M entrenables")

Verificalo a mano con el modelo de §1: 784×512 + 512 + 512×512 + 512 + 512×10 + 10 = 669 706. Con un adapter LoRA sobre un modelo base congelado, la segunda cifra baja a una fracción de un por ciento de la primera — el print_trainable_parameters() de PEFT es exactamente esta línea. Congelar es simplemente p.requires_grad_(False); y el optimizador conviene construirlo con [p for p in model.parameters() if p.requires_grad] para no arrastrar estado de Adam (dos tensores fp32 por parámetro) sobre pesos que nunca van a cambiar.

5. Lo que nn.Module registra y lo que no: buffers, ModuleList, ModuleDict

El registro automático de §1 tiene una regla simple y dos consecuencias que hay que conocer: se registra lo que es un nn.Module o un nn.Parameter asignado directamente como atributo. Una lista de Python con capas adentro no es ninguna de las dos cosas.

Listas y diccionarios de módulos

class MalHecho(nn.Module):
    def __init__(self):
        super().__init__()
        self.layers = [nn.Linear(8, 8) for _ in range(3)]      # lista de Python: invisible

class BienHecho(nn.Module):
    def __init__(self):
        super().__init__()
        self.layers = nn.ModuleList([nn.Linear(8, 8) for _ in range(3)])
    def forward(self, x):
        for layer in self.layers:      # ModuleList se itera e indexa como una lista
            x = torch.relu(layer(x))
        return x

len(list(MalHecho().parameters())), len(list(BienHecho().parameters()))

El modelo "mal hecho" corre sin ningún error: el forward funciona, la loss baja o no, pero parameters() está vacío, el optimizador no actualiza nada, .to("cuda") deja esas capas en CPU y el state_dict las omite. nn.ModuleList se comporta como una lista pero registra; nn.ModuleDict hace lo mismo con claves (útil para ramas por nombre: self.heads = nn.ModuleDict({"cls": ..., "reg": ...})); y nn.ParameterList / nn.ParameterDict son los equivalentes para tensores sueltos. A diferencia de nn.Sequential, ModuleList no define forward: vos decidís el orden y podés meter lógica entre capas — así están escritos los bloques de todos los transformers de HF.

Buffers: estado que no se entrena pero viaja con el modelo

Hay tensores que forman parte del modelo pero no son parámetros: la media y varianza acumuladas de BatchNorm, las tablas de senos y cosenos de RoPE, una máscara causal precalculada, los factores de escala de un modelo cuantizado. Si los guardás como atributo común (self.mask = torch.tril(...)) se quedan en CPU cuando movés el modelo a GPU y desaparecen del checkpoint. register_buffer resuelve las dos cosas:

class CausalAttentionLite(nn.Module):
    def __init__(self, dim, max_len=1024):
        super().__init__()
        self.qkv = nn.Linear(dim, 3 * dim)
        mask = torch.tril(torch.ones(max_len, max_len, dtype=torch.bool))
        self.register_buffer("mask", mask, persistent=False)   # viaja con .to(), no va al state_dict
        self.register_buffer("calls", torch.zeros(()))            # persistent=True (default): sí se guarda

    def forward(self, x):
        T = x.shape[1]
        q, k, v = self.qkv(x).chunk(3, dim=-1)
        att = (q @ k.transpose(-2, -1)) / q.shape[-1] ** 0.5
        att = att.masked_fill(~self.mask[:T, :T], float("-inf")).softmax(-1)
        self.calls += 1
        return att @ v

m = CausalAttentionLite(16)
[n for n, _ in m.named_buffers()], list(m.state_dict().keys())

Los dos buffers se mueven con m.to("cuda") y con m.bfloat16() (ojo: un buffer booleano no cambia con el cast de dtype, uno float sí); solo calls aparece en el state_dict porque mask se marcó como no persistente — sentido común para cualquier cosa que se pueda recomputar desde la config y ocupe espacio en el checkpoint. Cuando en un state_dict de HF ves claves como rotary_emb.inv_freq o bn.running_mean, son buffers.

6. Herramientas del día a día: to(), apply(), hooks

model.to(device) — in-place, a diferencia de los tensores

Para un tensor, .to() devuelve una copia. Para un módulo, .to() modifica el módulo en el lugar (recorre parámetros y buffers y los reemplaza) y devuelve self solo para poder encadenar. Así que model.to("cuda") sin asignar funciona, y x.to("cuda") sin asignar no. El mismo método acepta dtype (model.to(torch.bfloat16), o los atajos model.half() / model.bfloat16()), y es lo que hace from_pretrained(..., dtype=torch.bfloat16, device_map="cuda") por vos.

device = "cuda" if torch.cuda.is_available() else "cpu"
model.to(device)                         # in-place; el return es el mismo objeto
X, y = X.to(device), y.to(device)        # tensores: hay que reasignar
next(model.parameters()).device          # truco para saber dónde vive un modelo

model.apply(fn) — inicialización recursiva

apply llama fn(module) sobre cada submódulo, hijos primero, y termina en el propio modelo. Es la forma canónica de inicializar pesos según el tipo de capa, y de ahí vienen los _init_weights que ves en todos los modelos de Transformers:

def init_weights(m):
    if isinstance(m, nn.Linear):
        nn.init.trunc_normal_(m.weight, std=0.02)     # estilo GPT-2
        if m.bias is not None:
            nn.init.zeros_(m.bias)
    elif isinstance(m, nn.Embedding):
        nn.init.normal_(m.weight, std=0.02)

model.apply(init_weights)

Los nn.init.*_ son in-place sobre parámetros — funcionan porque corren bajo no_grad por dentro (recordá la página anterior). El default de nn.Linear es una uniforme tipo Kaiming; para redes profundas con ReLU la inicialización de He/Kaiming (kaiming_normal_) es la que viste en fast.ai Parte 2 cuando mediste que las activaciones no se fueran a cero capa a capa.

register_forward_hook — ver adentro sin tocar el código

Un forward hook es una función fn(module, input, output) que __call__ ejecuta después del forward de ese módulo. Sirve para imprimir shapes, guardar activaciones, medir estadísticas o incluso reemplazar la salida (si devuelve algo). El ejemplo clásico, que aplica a cualquier modelo aunque no sepas cómo está escrito:

handles = []
def print_shape(name):
    def hook(module, inputs, output):
        print(f"{name:28s} {tuple(inputs[0].shape)} -> {tuple(output.shape)}")
    return hook

for name, module in model.named_modules():
    if isinstance(module, (nn.Linear, nn.Flatten)):
        handles.append(module.register_forward_hook(print_shape(name)))

with torch.no_grad():
    model(torch.rand(2, 1, 28, 28))

for h in handles: h.remove()       # siempre: los hooks quedan pegados al módulo

named_modules() recorre el árbol completo (incluida la raíz), children() solo el primer nivel. Los mismos hooks, con register_forward_pre_hook (antes del forward) y register_full_backward_hook (durante backward), son la base de ActivationStats de fast.ai, de las herramientas de interpretabilidad que leen residual streams, y de los "hooks de parcheo" con los que se inyecta un adapter sin reescribir el modelo. Y son la razón por la que model.forward(x) directo es un error: se saltea todo esto.

7. Errores comunes

La mayoría de estos bugs no lanzan excepción: el modelo entrena "algo" y el problema aparece como una loss que no baja o un checkpoint incompleto.

8. Resumen

1. Una red neuronal es un nn.Module que contiene otros nn.Module — composición recursiva, sin límite de profundidad. El registro automático funciona porque __setattr__ intercepta cada asignación de módulo o parámetro.
2. Siempre model(x), nunca model.forward(x) — __call__ ejecuta los hooks (tuyos, de PEFT, de profilers) que forward solo no dispara.
3. Flatten, Linear, ReLU y Sequential son los mismos bloques que ya usaste en fast.ai, con su nombre oficial en torch.nn. Softmax no va al final del modelo si la loss es CrossEntropyLoss: el modelo devuelve logits.
4. Subclasear nn.Module rastrea automáticamente todos los parámetros — model.parameters() es lo que cualquier optimizador necesita para saber qué ajustar; sum(p.numel() for p in model.parameters() if p.requires_grad) es la verificación de una línea.
5. Solo se registra lo asignado directamente: listas de Python no. nn.ModuleList/nn.ModuleDict para capas, register_buffer para estado no entrenable que debe moverse con .to() y (opcionalmente) guardarse.
6. model.to() es in-place (a diferencia de los tensores), model.apply(fn) es la forma canónica de inicializar, y register_forward_hook permite ver shapes y activaciones de cualquier modelo sin tocar su código.

9. Ponete a prueba

Respondé antes de abrir cada respuesta. Si fallás una, volvé a la sección indicada.

10. Para practicar

import torch, torch.nn as nn
torch.manual_seed(0)
X = torch.randn(2000, 20); y = (X[:, :5].sum(1) > 0).long()   # 2 clases, fácil

def mlp(con_softmax):
    capas = [nn.Linear(20, 64), nn.ReLU(), nn.Linear(64, 2)]
    if con_softmax: capas.append(nn.Softmax(dim=1))
    return nn.Sequential(*capas)

for con_softmax in (False, True):
    torch.manual_seed(1)
    m = mlp(con_softmax); opt = torch.optim.Adam(m.parameters(), lr=1e-2); loss_fn = nn.CrossEntropyLoss()
    for step in range(300):
        loss = loss_fn(m(X), y); opt.zero_grad(); loss.backward(); opt.step()
    with torch.no_grad():
        out = m(X); probs = out if con_softmax else out.softmax(1)
        acc = (out.argmax(1) == y).float().mean()
    print(f"softmax al final={con_softmax}: loss={loss.item():.3f} acc={acc:.3f} confianza media={probs.max(1).values.mean():.3f}")

import torch, torch.nn as nn
# pegá acá la clase CausalAttentionLite de la sección 5

class MiniLM(nn.Module):
    def __init__(self, dim=32, n=2):
        super().__init__()
        self.blocks = nn.ModuleList([CausalAttentionLite(dim) for _ in range(n)])
    def forward(self, x):
        for b in self.blocks: x = x + b(x)
        return x

m = MiniLM()
print([k for k in m.state_dict()])
print(sum(p.numel() for p in m.parameters()))
m.to(torch.bfloat16)
print(m.blocks[0].calls.dtype, m.blocks[0].mask.dtype)
hs = [b.register_forward_hook(lambda mod, i, o: print("bloque ->", tuple(o.shape))) for b in m.blocks]
m(torch.randn(4, 10, 32, dtype=torch.bfloat16))
for h in hs: h.remove()

11. Fuentes y para seguir

• Build the Neural Network — Learn the Basics — el tutorial original; el recorrido capa por capa con shapes impresos que esta página resume.
• Modules (PyTorch notes) — la nota oficial que explica parámetros, buffers, hooks, apply y el ciclo de vida de un módulo; es la mejor lectura después de esta página.
• torch.nn.Module — referencia — la lista completa de métodos (named_modules, register_buffer, register_forward_hook, to, apply) con sus firmas exactas.
• torch.nn.CrossEntropyLoss — la documentación dice explícitamente que la entrada son logits sin normalizar; leela una vez para no volver a poner Softmax al final.
• Delving Deep into Rectifiers (He et al., 2015, arXiv:1502.01852) — el paper de la inicialización de Kaiming/He que nn.init.kaiming_normal_ implementa; la sección 2.2 explica por qué la varianza depende del fan-in y de ReLU.