Guardado y carga de modelos PyTorch

En proyectos de deep learning, el guardado y la carga de modelos es un aspecto crucial, principalmente por las siguientes razones:

  1. Recuperación tras interrupción del entrenamiento: cuando el proceso de entrenamiento se interrumpe inesperadamente, se puede continuar entrenando desde el punto guardado
  2. Despliegue de modelos: desplegar el modelo entrenado en un entorno de producción
  3. Compartir modelos: facilita compartir los resultados del modelo entre los miembros del equipo
  4. Aprendizaje por transferencia: guardar modelos preentrenados para usarlos en otras tareas
  5. Evaluación de rendimiento: guardar modelos de diferentes etapas de entrenamiento para compararlos

Métodos básicos de guardado y carga

Guardar el modelo completo

Este es el método más sencillo, guarda la arquitectura y los parámetros del modelo:

Ejemplo

import torch
import torchvision.models as models

# Crear y entrenar un modelo
model = models.resnet18(pretrained=True)
# ... código de entrenamiento ...

# Guardar el modelo completo
torch.save(model, 'model.pth')

# Cargar el modelo completo
loaded_model = torch.load('model.pth')

Ventajas：

  • Código simple e intuitivo
  • Guarda la estructura completa del modelo

Desventajas：

  • El archivo ocupa más espacio
  • Depende de la definición de la clase del modelo

Guardar solo los parámetros del modelo (método recomendado)

El método más recomendado es guardar solo el diccionario de estado (state_dict) del modelo:

Ejemplo

# Guardar los parámetros del modelo
torch.save(model.state_dict(), 'model_weights.pth')

# Cargar los parámetros del modelo
model = models.resnet18()  # Primero se debe crear un modelo con la misma arquitectura
model.load_state_dict(torch.load('model_weights.pth'))
model.eval()  # Configurar en modo evaluación

Ventajas：

  • Archivo más pequeño
  • Más flexible, se puede cargar en diferentes arquitecturas
  • Mejor compatibilidad

Guardar y cargar el estado de entrenamiento

En proyectos reales, normalmente también necesitamos guardar el estado del optimizador, la época (epoch) y otra información:

Ejemplo

# Guardar checkpoint
checkpoint = {
    'epoch': epoch,
    'model_state_dict': model.state_dict(),
    'optimizer_state_dict': optimizer.state_dict(),
    'loss': loss,
    # Se puede agregar otra información que necesite guardarse
}

torch.save(checkpoint, 'checkpoint.pth')

# Cargar checkpoint
checkpoint = torch.load('checkpoint.pth')
model.load_state_dict(checkpoint['model_state_dict'])
optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
epoch = checkpoint['epoch']
loss = checkpoint['loss']

model.eval()  # O model.train() dependiendo de tus necesidades

Cargar modelos entre dispositivos

Manejo de compatibilidad CPU/GPU

Ejemplo

# Especificar map_location al guardar
torch.save(model.state_dict(), 'model_weights.pth')

# Cargar en CPU (cuando el modelo se entrenó en GPU)
device = torch.device('cpu')
model.load_state_dict(torch.load('model_weights.pth', map_location=device))

# Cargar en GPU
device = torch.device('cuda')
model.load_state_dict(torch.load('model_weights.pth', map_location=device))
model.to(device)

Carga de modelos entrenados con múltiples GPU

Ejemplo

# Guardar modelo de múltiples GPU
torch.save(model.module.state_dict(), 'multigpu_model.pth')

# Cargar en una sola GPU
model = ModelClass()
model.load_state_dict(torch.load('multigpu_model.pth'))

Conversión y compatibilidad de modelos

Compatibilidad de versiones de PyTorch

Ejemplo

# Especificar _use_new_zipfile_serialization=True al guardar para obtener mejor compatibilidad
torch.save(model.state_dict(), 'model.pth', _use_new_zipfile_serialization=True)

Convertir a TorchScript

Ejemplo

# Convertir el modelo al formato TorchScript
scripted_model = torch.jit.script(model)
torch.jit.save(scripted_model, 'model_scripted.pt')

# Cargar un modelo TorchScript
loaded_script = torch.jit.load('model_scripted.pt')

Mejores prácticas y problemas comunes

Mejores prácticas

  1. Convenciones de nomenclatura: usa nombres de archivo significativos, comoresnet18_epoch50.pth
  2. Guardado periódico: guarda un checkpoint cada pocas épocas
  3. Verificar la carga: prueba la funcionalidad de carga inmediatamente después de guardar
  4. Documentación: registra la arquitectura del modelo y los parámetros de entrenamiento
  5. Control de versiones: incluye los archivos del modelo en el sistema de control de versiones

Soluciones a problemas comunes

Problema 1：Missing key(s) in state_dict
Solución: asegúrate de que la arquitectura del modelo coincida exactamente, o usastrict=Falseel parámetro:

model.load_state_dict(torch.load('model.pth'), strict=False)

Problema 2：CUDA out of memory
Solución: al cargar, primero colócalo en CPU:

Ejemplo

model.load_state_dict(torch.load('model.pth', map_location='cpu'))

Problema 3: no se puede cargar un modelo de una versión anterior
Solución: intenta cargarlo en diferentes versiones de PyTorch, o convierte el formato del modelo

Ejemplo de aplicación práctica

Flujo de guardado y carga de un modelo de clasificación de imágenes

Ejemplo de código completo

Ejemplo

import torch
import torch.nn as nn
import torch.optim as optim

# Definir un modelo simple
class SimpleModel(nn.Module):
    def __init__(self):
        super(SimpleModel, self).__init__()
        self.fc = nn.Linear(10, 2)
   
    def forward(self, x):
        return self.fc(x)

# Inicialización
model = SimpleModel()
optimizer = optim.SGD(model.parameters(), lr=0.01)
criterion = nn.CrossEntropyLoss()

# Simular el proceso de entrenamiento
for epoch in range(5):
    # Simular los pasos de entrenamiento
    inputs = torch.randn(32, 10)
    labels = torch.randint(0, 2, (32,))
   
    optimizer.zero_grad()
    outputs = model(inputs)
    loss = criterion(outputs, labels)
    loss.backward()
    optimizer.step()
   
    # Guardar un checkpoint cada 2 épocas
    if epoch % 2 == 0:
        checkpoint = {
            'epoch': epoch,
            'model_state_dict': model.state_dict(),
            'optimizer_state_dict': optimizer.state_dict(),
            'loss': loss.item(),
        }
        torch.save(checkpoint, f'checkpoint_epoch{epoch}.pth')
        print(f'Checkpoint saved at epoch {epoch}')

# Guardado final
torch.save(model.state_dict(), 'final_model.pth')

# Ejemplo de carga
loaded_model = SimpleModel()
loaded_model.load_state_dict(torch.load('final_model.pth'))
loaded_model.eval()

# Probar el modelo cargado
test_input = torch.randn(1, 10)
with torch.no_grad():
    output = loaded_model(test_input)
print(f'Test output: {output}')
Otras extensiones