Gestión de datos con LightningDataModule

Modelos de IA escalables con PyTorch Lightning

Sergiy Tkachuk

Director, GenAI Productivity

Preparar datos para entrenar el modelo

Una mala preparación de datos provoca problemas de entrenamiento
- Entrenamiento lento
- Interrupciones frecuentes
- No converge

Preparación de datos para el entrenamiento.png

¿Por qué usar LightningDataModule?

📂 Centraliza la gestión del dataset

📊 Estandariza los flujos de preparación de datos

🚀 Simplifica el entrenamiento y la evaluación

Gestión de datos con LightningDataModule

Métodos clave:

prepare_data: descarga y configura los datos
setup: divide en train, validación y test

class ImageDataModule(pl.LightningDataModule):
    def __init__(self, data_dir="./data", batch_size=32):
        super().__init__()
        ...

    def prepare_data(self):
        datasets.MNIST(self.data_dir, train=True, download=True)

    def setup(self, stage=None):
        dataset = datasets.MNIST(self.data_dir, train=True, transform=self.transform)
        self.train_data, self.val_data = random_split(dataset, [55000, 5000])
        self.test_data = datasets.MNIST(self.data_dir, train=False, transform=self.transform)

Crear el DataLoader de entrenamiento

Suministra lotes de datos de entrenamiento
Ayuda a optimizar el uso de la GPU
Permite iterar eficientemente en grandes datasets

def train_dataloader(self):
    return DataLoader(self.train_data, batch_size=self.batch_size, shuffle=True)

Crear el DataLoader de validación

Proporciona datos para validar el modelo
Ayuda a vigilar la generalización
Asegura consistencia en evaluaciones con barajado

def val_dataloader(self):
    return DataLoader(self.val_data, batch_size=self.batch_size)

Crear el DataLoader de test

Aporta datos para la evaluación final tras el entrenamiento
Simula una evaluación en condiciones reales
Garantiza una medición imparcial

def test_dataloader(self):
    return DataLoader(self.test_data, batch_size=self.batch_size)

Conectar DataModule con LightningModule

El diseño modular separa la lógica de datos y del modelo

Diagrama de PyTorch Lightning

Conectar DataModule con LightningModule

El diseño modular separa la lógica de datos y del modelo
LightningDataModule se empareja con LightningModule
Flujo estandarizado mejora la reproducibilidad

Diagrama de PyTorch Lightning con DataModule y LightningModule

¡Vamos a practicar!

Modelos de IA escalables con PyTorch Lightning