Construindo Modelos com nn.Module

nn.Module é a classe base para todos os modelos PyTorch. Subclasseie-a para definir camadas, implementar forward() e registrar parâmetros que os otimizadores atualizam durante o treinamento.

Receita

Cartão de receita de referência rápida - pronto para copiar e colar.

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)

Quando usar isso:

• Definir qualquer arquitetura de rede neural em PyTorch.
• Compor camadas com lógica forward customizada (conexões de atalho, atenção).
• Inspecionar, congelar ou treinar seletivamente grupos de parâmetros.
• Serializar modelos via state_dict().

Exemplo de Trabalho

"""building_models.py - CNN customizada para classificação de imagens."""
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("logits shape:", 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"params: {total_params:,} trainable: {trainable:,}")

O que isso demonstra:

• Camadas registradas em __init__ tornam-se parte de model.parameters().
• forward define a computação; chame via model(x), não model.forward(x).
• .to(device) move todos os parâmetros e buffers para a GPU.
• flatten(1) preserva a dimensão do lote enquanto colapsa as dimensões espaciais.

Mergulho Profundo

Como Funciona

• __init__ registra submódulos (nn.Linear, nn.Conv2d) e parâmetros (nn.Parameter).
• forward executa o grafo de computação; hooks podem interceptar entradas/saídas.
• nn.Module rastreia o modo training vs eval para comportamento de dropout e batch norm.
• state_dict() retorna um dicionário de tensores de parâmetros para salvar/carregar.
• named_modules() e named_children() percorrem a árvore de módulos.

Padrões de Módulo

Notas de Python

# Congela o backbone, treina apenas a cabeça
for param in model.conv1.parameters():
    param.requires_grad = False
 
# Taxas de aprendizado diferentes por grupo
optimizer = torch.optim.Adam([
    {"params": model.conv1.parameters(), "lr": 1e-5},
    {"params": model.fc2.parameters(), "lr": 1e-3},
])

Armadilhas

• Chamar forward() diretamente - pula hooks e wrappers nn.Module. Correção: sempre chame model(x).
• Esquecer super().__init__() - submódulos não registrados. Correção: chame super().__init__() primeiro em __init__.
• Não chamar model.eval() para inferência - dropout e batch norm se comportam incorretamente. Correção: model.eval() antes da inferência; model.train() para treinamento.
• Criar camadas em forward - novos parâmetros a cada passagem, nunca treinados. Correção: defina todas as camadas em __init__.
• Incompatibilidade de dispositivo - modelo na GPU, entrada na CPU. Correção: model.to(device) e x.to(device).
• BatchNorm com batch_size=1 - estatísticas indefinidas. Correção: use model.eval() ou nn.GroupNorm para lotes pequenos.

Alternativas

FAQs

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)
# ou
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)

Relacionados

• Fundamentos do PyTorch - fundamentos de tensor e nn
• Loops de Treinamento - treinando o modelo
• Transfer Learning & Fine-Tuning - congelando camadas
• Salvando, Carregando & Exportando - persistência de state_dict
• PyTorch Lightning - wrapper LightningModule

Versões do Stack: Esta página foi escrita para Python 3.14.0 (estável 3.14, manutenção 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+, e uv 0.6+.