Función PyTorch torch.nn.Module
Manual de referencia de PyTorch torch.nn
torch.nn.ModuleEs la clase base de todos los módulos de redes neuronales en PyTorch.
Todos los modelos de red personalizados deben heredar de esta clase, que proporciona funciones como gestión de parámetros, migración de dispositivos y guardado de modelos.
Definición de clase
torch.nn.Module
Atributos principales
parameters(): Devuelve todos los parámetros aprendibles del modelo.named_parameters(): Devuelve un iterador de nombres y valores de parámetros.children(): Devuelve los submódulos del modelo.named_children(): Devuelve un iterador de nombres y módulos de los submódulos.modules(): Devuelve todos los módulos.state_dict(): Devuelve un diccionario que contiene todos los parámetros.
Ejemplos de uso
Ejemplo 1: Crear un módulo personalizado
Hereda de nn.Module para crear una red personalizada:
Ejemplo
import torch
import torch.nn as nn
class SimpleNet(nn.Module):
def __init__(self, input_dim, hidden_dim, output_dim):
super(SimpleNet, self).__init__()
# Definir capas de red
self.fc1 = nn.Linear(input_dim, hidden_dim)
self.relu = nn.ReLU()
self.fc2 = nn.Linear(hidden_dim, output_dim)
def forward(self, x):
# Definir la propagación hacia adelante
x = self.fc1(x)
x = self.relu(x)
x = self.fc2(x)
return x
# Instanciar el modelo
model = SimpleNet(input_dim=784, hidden_dim=256, output_dim=10)
# Probar
x = torch.randn(32, 784)
output = model(x)
print("Estructura del modelo:")
print(model)
print("nForma de entrada:", x.shape)
print("Forma de salida:", output.shape)
import torch.nn as nn
class SimpleNet(nn.Module):
def __init__(self, input_dim, hidden_dim, output_dim):
super(SimpleNet, self).__init__()
# Definir capas de red
self.fc1 = nn.Linear(input_dim, hidden_dim)
self.relu = nn.ReLU()
self.fc2 = nn.Linear(hidden_dim, output_dim)
def forward(self, x):
# Definir la propagación hacia adelante
x = self.fc1(x)
x = self.relu(x)
x = self.fc2(x)
return x
# Instanciar el modelo
model = SimpleNet(input_dim=784, hidden_dim=256, output_dim=10)
# Probar
x = torch.randn(32, 784)
output = model(x)
print("Estructura del modelo:")
print(model)
print("nForma de entrada:", x.shape)
print("Forma de salida:", output.shape)
Ejemplo 2: Gestionar parámetros
Acceder y gestionar los parámetros del modelo:
Ejemplo
import torch
import torch.nn as nn
class Net(nn.Module):
def __init__(self):
super(Net, self).__init__()
self.conv1 = nn.Conv2d(3, 16, 3)
self.bn1 = nn.BatchNorm2d(16)
self.fc = nn.Linear(16 * 14 * 14, 10)
def forward(self, x):
x = self.conv1(x)
x = self.bn1(x)
x = x.view(x.size(0), -1)
x = self.fc(x)
return x
model = Net()
# Contar el número de parámetros
total_params = sum(p.numel() for p in model.parameters())
trainable_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
print("Número total de parámetros:", total_params)
print("Número de parámetros entrenables:", trainable_params)
# Acceder a parámetros específicos
print("nForma de los pesos de conv1:", model.conv1.weight.shape)
print("Forma del sesgo de fc:", model.fc.bias.shape)
import torch.nn as nn
class Net(nn.Module):
def __init__(self):
super(Net, self).__init__()
self.conv1 = nn.Conv2d(3, 16, 3)
self.bn1 = nn.BatchNorm2d(16)
self.fc = nn.Linear(16 * 14 * 14, 10)
def forward(self, x):
x = self.conv1(x)
x = self.bn1(x)
x = x.view(x.size(0), -1)
x = self.fc(x)
return x
model = Net()
# Contar el número de parámetros
total_params = sum(p.numel() for p in model.parameters())
trainable_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
print("Número total de parámetros:", total_params)
print("Número de parámetros entrenables:", trainable_params)
# Acceder a parámetros específicos
print("nForma de los pesos de conv1:", model.conv1.weight.shape)
print("Forma del sesgo de fc:", model.fc.bias.shape)
Ejemplo 3: Guardar y cargar modelos
Guardar y cargar el modelo:
Ejemplo
import torch
import torch.nn as nn
import tempfile
import os
class Net(nn.Module):
def __init__(self):
super(Net, self).__init__()
self.fc = nn.Linear(10, 5)
def forward(self, x):
return self.fc(x)
model = Net()
# Guardar el modelo completo
with tempfile.NamedTemporaryFile(delete=False, suffix='.pth') as f:
torch.save(model, f.name)
path_full = f.name
# Guardar state_dict (método recomendado)
with tempfile.NamedTemporaryFile(delete=False, suffix='.pth') as f:
torch.save(model.state_dict(), f.name)
path_state = f.name
# Cargar el modelo
loaded_model = Net()
loaded_model.load_state_dict(torch.load(path_state))
loaded_model.eval()
# Probar el modelo cargado
x = torch.randn(2, 10)
output1 = model(x)
output2 = loaded_model(x)
print("Salida original:", output1[0].tolist())
print("Salida después de cargar:", output2[0].tolist())
# Limpiar
os.remove(path_full)
os.remove(path_state)
import torch.nn as nn
import tempfile
import os
class Net(nn.Module):
def __init__(self):
super(Net, self).__init__()
self.fc = nn.Linear(10, 5)
def forward(self, x):
return self.fc(x)
model = Net()
# Guardar el modelo completo
with tempfile.NamedTemporaryFile(delete=False, suffix='.pth') as f:
torch.save(model, f.name)
path_full = f.name
# Guardar state_dict (método recomendado)
with tempfile.NamedTemporaryFile(delete=False, suffix='.pth') as f:
torch.save(model.state_dict(), f.name)
path_state = f.name
# Cargar el modelo
loaded_model = Net()
loaded_model.load_state_dict(torch.load(path_state))
loaded_model.eval()
# Probar el modelo cargado
x = torch.randn(2, 10)
output1 = model(x)
output2 = loaded_model(x)
print("Salida original:", output1[0].tolist())
print("Salida después de cargar:", output2[0].tolist())
# Limpiar
os.remove(path_full)
os.remove(path_state)
Ejemplo 4: Migración de dispositivos
Migrar el modelo entre diferentes dispositivos:
Ejemplo
import torch
import torch.nn as nn
class Net(nn.Module):
def __init__(self):
super(Net, self).__init__()
self.fc = nn.Linear(10, 5)
def forward(self, x):
return self.fc(x)
model = Net()
# Ver el dispositivo actual
print("Dispositivo de los parámetros:", model.fc.weight.device)
# Mover a GPU (si está disponible)
if torch.cuda.is_available():
model = model.cuda()
print("Después de mover a GPU:", model.fc.weight.device)
# Mover de vuelta a CPU
model = model.cpu()
print("Después de mover de vuelta a CPU:", model.fc.weight.device)
import torch.nn as nn
class Net(nn.Module):
def __init__(self):
super(Net, self).__init__()
self.fc = nn.Linear(10, 5)
def forward(self, x):
return self.fc(x)
model = Net()
# Ver el dispositivo actual
print("Dispositivo de los parámetros:", model.fc.weight.device)
# Mover a GPU (si está disponible)
if torch.cuda.is_available():
model = model.cuda()
print("Después de mover a GPU:", model.fc.weight.device)
# Mover de vuelta a CPU
model = model.cpu()
print("Después de mover de vuelta a CPU:", model.fc.weight.device)
Ejemplo 5: Inicialización con apply
Inicialización recursiva con apply:
Ejemplo
import torch
import torch.nn as nn
def init_weights(module):
if isinstance(module, nn.Linear):
nn.init.xavier_uniform_(module.weight)
if module.bias is not None:
nn.init.zeros_(module.bias)
elif isinstance(module, nn.Conv2d):
nn.init.kaiming_normal_(module.weight)
class Net(nn.Module):
def __init__(self):
super(Net, self).__init__()
self.fc1 = nn.Linear(10, 20)
self.fc2 = nn.Linear(20, 5)
def forward(self, x):
return self.fc2(self.fc1(x))
model = Net()
model.apply(init_weights)
print("Pesos de fc1:", model.fc1.weight[0, :3].tolist())
print("Pesos de fc2:", model.fc2.weight[0, :3].tolist())
import torch.nn as nn
def init_weights(module):
if isinstance(module, nn.Linear):
nn.init.xavier_uniform_(module.weight)
if module.bias is not None:
nn.init.zeros_(module.bias)
elif isinstance(module, nn.Conv2d):
nn.init.kaiming_normal_(module.weight)
class Net(nn.Module):
def __init__(self):
super(Net, self).__init__()
self.fc1 = nn.Linear(10, 20)
self.fc2 = nn.Linear(20, 5)
def forward(self, x):
return self.fc2(self.fc1(x))
model = Net()
model.apply(init_weights)
print("Pesos de fc1:", model.fc1.weight[0, :3].tolist())
print("Pesos de fc2:", model.fc2.weight[0, :3].tolist())
Ejemplo 6: Estructura de red compleja
Construir una red con ramas:
Ejemplo
import torch
import torch.nn as nn
class BranchNet(nn.Module):
def __init__(self):
super(BranchNet, self).__init__()
# Backbone
self.shared = nn.Linear(10, 20)
# Rama
self.branch_a = nn.Linear(20, 5)
self.branch_b = nn.Linear(20, 3)
def forward(self, x):
feat = self.shared(x)
out_a = self.branch_a(feat)
out_b = self.branch_b(feat)
return out_a, out_b
model = BranchNet()
x = torch.randn(4, 10)
out_a, out_b = model(x)
print("Salida de la rama A:", out_a.shape)
print("Salida de la rama B:", out_b.shape)
import torch.nn as nn
class BranchNet(nn.Module):
def __init__(self):
super(BranchNet, self).__init__()
# Backbone
self.shared = nn.Linear(10, 20)
# Rama
self.branch_a = nn.Linear(20, 5)
self.branch_b = nn.Linear(20, 3)
def forward(self, x):
feat = self.shared(x)
out_a = self.branch_a(feat)
out_b = self.branch_b(feat)
return out_a, out_b
model = BranchNet()
x = torch.randn(4, 10)
out_a, out_b = model(x)
print("Salida de la rama A:", out_a.shape)
print("Salida de la rama B:", out_b.shape)
Preguntas frecuentes
P1: ¿Por qué se debe llamar a super().__init__()?
Llama al método de inicialización de la clase padre, asegurando que el módulo registre correctamente los parámetros.
<h3 P2: ¿Cómo ver la estructura del modelo?Simplemente usa print(model) o la función summary de torchvision.
<h3 P3: ¿Cuál es la diferencia entre parameters() y modules()?parameters() devuelve los parámetros aprendibles, modules() devuelve todos los módulos.
Casos de uso
nn.ModuleEs la base para construir todas las redes neuronales personalizadas. Los principales casos de uso incluyen:
- Definir la estructura de la red: Cualquier red neuronal personalizada
- Gestión de modelos: Guardar, cargar, migrar
- Gestión de parámetros: Actualización del optimizador, cálculo de gradientes
Consejo: Todos los módulos personalizados deben implementar el método forward.
Otras extensiones