Definir modelos con LightningModule

Modelos de IA escalables con PyTorch Lightning

Sergiy Tkachuk

Director, GenAI Productivity

LightningModule en foco

Encapsula la arquitectura del modelo
Organiza el entrenamiento en una unidad manejable
Plano que aporta orden y claridad a proyectos de deep learning

Diagrama de PyTorch LightningModule

Definir el método init

Tareas clave:

Inicialización del modelo
super():
- Manejo automático de loops de entrenamiento
- Logging
- Checkpointing
Capas del modelo tras la inicialización
Modular y fácil de mantener

import lightning.pytorch as pl
import torch.nn as nn

class ClassificationModel(pl.LightningModule):
    def __init__(self, input_dim,
                 hidden_dim, num_class):
          # Initialize parent class
        super().__init__()

        # First layer
        self.layer1 = nn.Linear(input_dim,
                                hidden_dim)
        # Activation function
        self.relu = nn.ReLU()
        # Output layer
        self.layer2 = nn.Linear(hidden_dim,
                                num_class)

Implementar el método forward

Pasos clave:

Define el flujo de datos por la red
Procesa la entrada por capas secuenciales
- Transformación lineal
- Activación
- Última capa y salida

import lightning.pytorch as pl
import torch.nn as nn

class ClassificationModel(pl.LightningModule):
    def __init__(self, input_dim,
                 hidden_dim, num_class):
          ...

    def forward(self, x):

        x = self.layer1(x) # Pass input

        x = nn.ReLU(x) # Apply activation

        x = self.layer2(x) # Compute output

        return x # Return result

Ejemplo: clasificar dígitos escritos a mano

import lightning.pytorch as pl
from torch.utils.data import DataLoader
from torchvision.datasets import MNIST
from torchvision import transforms

transform = transforms.ToTensor()
train_ds = MNIST(root='.', train=True,  download=True, transform=transform)
test_ds  = MNIST(root='.', train=False, download=True, transform=transform)

train_loader = DataLoader(train_ds, batch_size=64, shuffle=True)
test_loader  = DataLoader(test_ds,  batch_size=64)


model = ClassificationModel(input_dim=28*28, hidden_dim=128, num_class=10)


trainer = pl.Trainer(max_epochs=3, accelerator='auto')
trainer.fit(model, train_loader, test_loader)

Integrar el modelo con tareas de clasificación

Enfoque en el caso de clasificación
Flujo completo dentro de LightningModule
Salida en bruto para activación softmax
Integración con Lightning Trainer

class ClassificationModel(pl.LightningModule):
  def __init__(self, input_dim, 
               hidden_dim, output_dim):
    super().__init__()

    self.hid = nn.Linear(input_dim, hidden_dim)
    self.out = nn.Linear(hidden_dim, output_dim)


  def forward(self, x):

    x = self.hidden(x)
    x = nn.ReLU(x)
    x = self.output(x)

    return x

¡Vamos a practicar!

Modelos de IA escalables con PyTorch Lightning