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

Patrón	Usar Cuando	Ejemplo
`nn.Sequential`	Cadena lineal de capas	Clasificadores MLP
Subclase `nn.Module`	Lógica de pase hacia adelante personalizada	Conexiones de salto ResNet
`nn.ModuleList`	Número variable de capas	Redes de profundidad dinámica
`nn.ModuleDict`	Submódulos con nombre	Arquitecturas de múltiples cabezas

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

Alternativa	Usar Cuando	No Usar Cuando
Subclase `nn.Module`	Arquitecturas personalizadas	MLP trivial de 3 capas (usa Sequential)
`nn.Sequential`	Pilas simples de avance	Necesidad de ramificación o conexiones de salto
PyTorch Lightning `LightningModule`	Bucles de entrenamiento estructurados	Aprender los fundamentos de PyTorch
`torch.nn.functional`	Operaciones sin estado en `forward`	Necesidad de parámetros aprendibles

Preguntas Frecuentes

¿Cuál es la diferencia entre nn.Module y nn.functional?

Las clases nn.Module contienen parámetros aprendibles (Linear, Conv2d).
nn.functional proporciona funciones sin estado (relu, conv2d con pesos explícitos).
Usa clases de Módulo en __init__; usa funcional para operaciones únicas.

¿Cómo cuento los parámetros del modelo?

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

numel() devuelve el número total de elementos en un tensor.

¿Qué son los buffers frente a los parámetros?

Los parámetros son actualizados por el optimizador.
Los buffers (estadísticas de ejecución de BatchNorm) se guardan en state_dict pero no se entrenan.
Registra buffers con self.register_buffer("nombre", tensor).

¿Cómo añado una pérdida personalizada en el modelo?

Mantén la pérdida fuera del modelo; calcúlala en el bucle de entrenamiento.
Excepción: algunos modelos multitarefa devuelven la pérdida desde forward por conveniencia.

¿Puedo anidar módulos?

Sí, cualquier nn.Module puede contener otros módulos como atributos.
model.children() devuelve los hijos directos; model.modules() es recursivo.

¿Qué hace model.train()?

Establece el modo de entrenamiento: dropout activo, batch norm usa estadísticas del lote.
Llama al inicio de cada época de entrenamiento.
model.eval() para validación e inferencia.

¿Cómo inicializo los pesos?

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

Aplica la inicialización después de la construcción del modelo.
Los modelos preentrenados omiten la inicialización manual.

¿Qué es un hook?

Callbacks registrados en módulos o tensores.
register_forward_hook inspecciona activaciones intermedias.
Elimina los hooks para evitar fugas de memoria.

¿Cómo imprimo la arquitectura del modelo?

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

torchinfo muestra el recuento de parámetros por capa.

¿Puedo usar sugerencias de tipo en forward?

Sí, anota def forward(self, x: torch.Tensor) -> torch.Tensor.
Mejora la legibilidad; no tiene efecto en tiempo de ejecución.

¿Cómo funcionan las conexiones de salto?

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

Residual = entrada sumada a la salida del bloque antes de la activación.

¿Cómo muevo solo algunas capas a la GPU?

Mueve todo el modelo con .to(device); los movimientos parciales causan errores de dispositivo.
Congela las capas en lugar de dejarlas en la CPU.

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+.

Construcción de Modelos con nn.Module

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

Patrón	Usar Cuando	Ejemplo
`nn.Sequential`	Cadena lineal de capas	Clasificadores MLP
Subclase `nn.Module`	Lógica de pase hacia adelante personalizada	Conexiones de salto ResNet
`nn.ModuleList`	Número variable de capas	Redes de profundidad dinámica
`nn.ModuleDict`	Submódulos con nombre	Arquitecturas de múltiples cabezas

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

Alternativa	Usar Cuando	No Usar Cuando
Subclase `nn.Module`	Arquitecturas personalizadas	MLP trivial de 3 capas (usa Sequential)
`nn.Sequential`	Pilas simples de avance	Necesidad de ramificación o conexiones de salto
PyTorch Lightning `LightningModule`	Bucles de entrenamiento estructurados	Aprender los fundamentos de PyTorch
`torch.nn.functional`	Operaciones sin estado en `forward`	Necesidad de parámetros aprendibles

Preguntas Frecuentes

¿Cuál es la diferencia entre nn.Module y nn.functional?

Las clases nn.Module contienen parámetros aprendibles (Linear, Conv2d).
nn.functional proporciona funciones sin estado (relu, conv2d con pesos explícitos).
Usa clases de Módulo en __init__; usa funcional para operaciones únicas.

¿Cómo cuento los parámetros del modelo?

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

numel() devuelve el número total de elementos en un tensor.

¿Qué son los buffers frente a los parámetros?

Los parámetros son actualizados por el optimizador.
Los buffers (estadísticas de ejecución de BatchNorm) se guardan en state_dict pero no se entrenan.
Registra buffers con self.register_buffer("nombre", tensor).

¿Cómo añado una pérdida personalizada en el modelo?

Mantén la pérdida fuera del modelo; calcúlala en el bucle de entrenamiento.
Excepción: algunos modelos multitarefa devuelven la pérdida desde forward por conveniencia.

¿Puedo anidar módulos?

Sí, cualquier nn.Module puede contener otros módulos como atributos.
model.children() devuelve los hijos directos; model.modules() es recursivo.

¿Qué hace model.train()?

Establece el modo de entrenamiento: dropout activo, batch norm usa estadísticas del lote.
Llama al inicio de cada época de entrenamiento.
model.eval() para validación e inferencia.

¿Cómo inicializo los pesos?

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

Aplica la inicialización después de la construcción del modelo.
Los modelos preentrenados omiten la inicialización manual.

¿Qué es un hook?

Callbacks registrados en módulos o tensores.
register_forward_hook inspecciona activaciones intermedias.
Elimina los hooks para evitar fugas de memoria.

¿Cómo imprimo la arquitectura del modelo?

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

torchinfo muestra el recuento de parámetros por capa.

¿Puedo usar sugerencias de tipo en forward?

Sí, anota def forward(self, x: torch.Tensor) -> torch.Tensor.
Mejora la legibilidad; no tiene efecto en tiempo de ejecución.

¿Cómo funcionan las conexiones de salto?

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

Residual = entrada sumada a la salida del bloque antes de la activación.

¿Cómo muevo solo algunas capas a la GPU?

Mueve todo el modelo con .to(device); los movimientos parciales causan errores de dispositivo.
Congela las capas en lugar de dejarlas en la CPU.

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+.

Construcción de Modelos con nn.Module

Buscar en la documentación

Construcción de Modelos con nn.Module