Introducción a PyTorch Lightning

Modelos de IA escalables con PyTorch Lightning

Sergiy Tkachuk

Director, GenAI Productivity

PyTorch y PyTorch Lightning

PyTorch estándar:

Mucho trabajo manual
Bucles de entrenamiento explícitos
Gestión de GPU/TPU, logging y checkpoints

PyTorch y PyTorch Lightning

PyTorch Lightning:

Construido sobre PyTorch
Automatiza:
- Entrenamiento
- Checkpointing
- Logging
Reduce boilerplate
Mejora escalabilidad y reproducibilidad

Logotipo de PyTorch Lightning

Visión general de PyTorch Lightning

Ejemplo: e-commerce global optimizando flujos
- Desarrollo de búsqueda visual
- Bucles de entrenamiento automáticos
- Iteración rápida con poco boilerplate

Componentes base: LightningModule y Trainer

from lightning.pytorch import LightningModule
from lightning.pytorch import Trainer

Estructura de Lightning

Componentes clave:

LightningModule: lógica central del modelo

Una bombilla que representa la lógica

Estructura de Lightning

Componentes clave:

LightningModule: lógica central del modelo
Lightning Trainer: orquesta el entrenamiento

Una bombilla que representa la lógica y un engranaje que representa las operaciones

Estructura de Lightning

Componentes clave:

LightningModule: lógica central del modelo
Lightning Trainer: orquesta el entrenamiento
DataModule: organiza los flujos de datos
Callbacks: automatiza eventos
Logger: registra experimentos

Una bombilla que representa la lógica y un engranaje que representa las operaciones

LightningModule en acción

Puntos clave:

__init__: Define la arquitectura
forward(): Pasa datos por el modelo
training_step(): Define el entrenamiento
Hooks personalizados disponibles

import lightning.pytorch as pl

class LightClassifier(pl.LightningModule):
    def __init__(self, model, criterion, optimizer):

        super().__init__()
        self.model = model
        self.criterion = criterion
        self.optimizer = optimizer

    def forward(self, x):
        return self.model(x)

    def training_step(self, batch, batch_idx):
        x, y = batch
        logits = self(x)
        loss = self.criterion(logits, y)
        return loss

Lightning Trainer en acción

Puntos clave:

Gestiona el bucle de entrenamiento
Soporta entrenamiento distribuido
Maneja callbacks y logging
Optimiza recursos

model = LightClassifier()



trainer = Trainer(max_epochs=10, accelerator="gpu", devices=1)
trainer.fit(model, train_dataloader, val_dataloader)

Presentación del dataset Afro-MNIST

Conjunto de datasets sintéticos tipo MNIST para cuatro ortografías de lenguas afroasiáticas y Níger-Congo: Ge'ez (etíope), Vai, Osmanya y N'Ko.

¹ Wu, Daniel J., Andrew C. Yang, and Vinay U. Prabhu. "Afro-MNIST: Synthetic generation of MNIST-style datasets for low-resource languages." arXiv preprint arXiv:2009.13509 (2020).

Resumen de PyTorch Lightning

Diagrama de PyTorch Lightning

¡Vamos a practicar!

Modelos de IA escalables con PyTorch Lightning