Construcción de Modelos con nn.Module

nn.Module es la clase base para todos los modelos de PyTorch. Crea una subclase de ella para definir capas, implementar forward() y registrar parámetros que los optimizadores actualizan durante el entrenamiento.

Receta

Tarjeta de referencia rápida - lista para copiar y pegar.

import torch.nn as nn
 
class MLP(nn.Module):
    def __init__(self, in_dim: int, hidden: int, out_dim: int):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(in_dim, hidden),
            nn.ReLU(),
            nn.Linear(hidden, out_dim),
        )
 
    def forward(self, x):
        return self.net(x)

Cuándo usar esto:

• Definir cualquier arquitectura de red neuronal en PyTorch.
• Componer capas con lógica de pase hacia adelante personalizada (conexiones de salto, atención).
• Inspeccionar, congelar o entrenar selectivamente grupos de parámetros.
• Serializar modelos a través de state_dict().

Ejemplo de Trabajo

"""building_models.py - CNN personalizada para clasificación de imágenes."""
from __future__ import annotations
 
import torch
import torch.nn as nn
import torch.nn.functional as F
 
class SmallCNN(nn.Module):
    def __init__(self, num_classes: int = 10):
        super().__init__()
        self.conv1 = nn.Conv2d(1, 32, kernel_size=3, padding=1)
        self.conv2 = nn.Conv2d(32, 64, kernel_size=3, padding=1)
        self.pool = nn.MaxPool2d(2)
        self.fc1 = nn.Linear(64 * 7 * 7, 128)
        self.fc2 = nn.Linear(128, num_classes)
        self.dropout = nn.Dropout(0.25)
 
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        x = self.pool(F.relu(self.conv1(x)))   # 28->14
        x = self.pool(F.relu(self.conv2(x)))   # 14->7
        x = x.flatten(1)
        x = self.dropout(F.relu(self.fc1(x)))
        return self.fc2(x)
 
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = SmallCNN(num_classes=10).to(device)
 
x = torch.randn(8, 1, 28, 28, device=device)
logits = model(x)
print("forma de logits:", logits.shape)  # (8, 10)
 
total_params = sum(p.numel() for p in model.parameters())
trainable = sum(p.numel() for p in model.parameters() if p.requires_grad)
print(f"parámetros: {total_params:,} entrenables: {trainable:,}")

Lo que esto demuestra:

• Las capas registradas en __init__ se convierten en parte de model.parameters().
• forward define la computación; llámalo a través de model(x), no model.forward(x).
• .to(device) mueve todos los parámetros y buffers a la GPU.
• flatten(1) preserva la dimensión del lote mientras colapsa las dimensiones espaciales.

Inmersión Profunda

Cómo Funciona

• __init__ registra submódulos (nn.Linear, nn.Conv2d) y parámetros (nn.Parameter).
• forward ejecuta el grafo de computación; los hooks pueden interceptar entradas/salidas.
• nn.Module rastrea el modo training vs eval para el comportamiento de dropout y batch norm.
• state_dict() devuelve un diccionario de tensores de parámetros para guardar/cargar.
• named_modules() y named_children() recorren el árbol de módulos.

Patrones de Módulos

Notas de Python

# Congelar la base, entrenar solo la cabeza
for param in model.conv1.parameters():
    param.requires_grad = False
 
# Diferentes tasas de aprendizaje por grupo
optimizer = torch.optim.Adam([
    {"params": model.conv1.parameters(), "lr": 1e-5},
    {"params": model.fc2.parameters(), "lr": 1e-3},
])

Errores Comunes

• Llamar a forward() directamente - omite los hooks y los envoltorios de nn.Module. Solución: siempre llama a model(x).
• Olvidar super().__init__() - submódulos no registrados. Solución: llama a super().__init__() primero en __init__.
• No llamar a model.eval() para inferencia - dropout y batch norm se comportan incorrectamente. Solución: model.eval() antes de la inferencia; model.train() para entrenamiento.
• Crear capas en forward - nuevos parámetros en cada pasada, nunca entrenados. Solución: define todas las capas en __init__.
• Desajuste de dispositivo - modelo en GPU, entrada en CPU. Solución: model.to(device) y x.to(device).
• BatchNorm con batch_size=1 - las estadísticas no están definidas. Solución: usa model.eval() o nn.GroupNorm para lotes pequeños.

Alternativas

Preguntas Frecuentes

sum(p.numel() for p in model.parameters() if p.requires_grad)

def init_weights(m):
    if isinstance(m, nn.Linear):
        nn.init.xavier_uniform_(m.weight)
 
model.apply(init_weights)

print(model)
# o
from torchinfo import summary
summary(model, input_size=(8, 1, 28, 28))

def forward(self, x):
    residual = x
    out = self.block(x)
    return F.relu(out + residual)

Relacionado

• Fundamentos de PyTorch - fundamentos de tensores y nn
• Bucles de Entrenamiento - entrenando el modelo
• Aprendizaje por Transferencia y Ajuste Fino - congelación de capas
• Guardar, Cargar y Exportar - persistencia de state_dict
• PyTorch Lightning - envoltorio LightningModule

Versiones de la Pila: Esta página fue escrita para Python 3.14.0 (estable 3.14, mantenimiento 3.13), FastAPI 0.115+, Django 5.2, Flask 3.1, Pydantic 2, PyTorch 2.6+, pandas 2.2+, Polars 1.x, ruff 0.9+, y uv 0.6+.