Pytorch Lightning

PyTorch Lightning

Visão Geral

PyTorch Lightning é um framework de deep learning que organiza código PyTorch para eliminar boilerplate mantendo total flexibilidade. Automatize workflows de treinamento, orquestração multi-device e implemente boas práticas para treinamento e scaling de redes neurais em múltiplas GPUs/TPUs.

Quando Usar Esta Skill

Esta skill deve ser usada quando:

• Construir, treinar ou fazer deploy de redes neurais usando PyTorch Lightning
• Organizar código PyTorch em LightningModules
• Configurar Trainers para treinamento multi-GPU/TPU
• Implementar data pipelines com LightningDataModules
• Trabalhar com callbacks, logging e estratégias de distributed training (DDP, FSDP, DeepSpeed)
• Estruturar projetos de deep learning profissionalmente

Capacidades Principais

1. LightningModule - Definição de Modelo

Organize modelos PyTorch em seis seções lógicas:

1. Inicialização - __init__() e setup()
2. Training Loop - training_step(batch, batch_idx)
3. Validation Loop - validation_step(batch, batch_idx)
4. Test Loop - test_step(batch, batch_idx)
5. Prediction - predict_step(batch, batch_idx)
6. Optimizer Configuration - configure_optimizers()

Referência de template rápido: Veja scripts/template_lightning_module.py para boilerplate completo.

Documentação detalhada: Leia references/lightning_module.md para documentação abrangente de métodos, hooks, propriedades e boas práticas.

2. Trainer - Automação de Treinamento

O Trainer automatiza o loop de treinamento, gerenciamento de devices, operações de gradiente e callbacks. Recursos principais:

• Suporte multi-GPU/TPU com seleção de strategy (DDP, FSDP, DeepSpeed)
• Treinamento com precisão mista automática
• Gradient accumulation e clipping
• Checkpointing e early stopping
• Barras de progresso e logging

Referência de setup rápido: Veja scripts/quick_trainer_setup.py para configurações comuns de Trainer.

Documentação detalhada: Leia references/trainer.md para todos os parâmetros, métodos e opções de configuração.

3. LightningDataModule - Organização de Data Pipeline

Encapsule todos os passos de processamento de dados em uma classe reutilizável:

1. prepare_data() - Download e processe dados (single-process)
2. setup() - Crie datasets e aplique transforms (per-GPU)
3. train_dataloader() - Retorne training DataLoader
4. val_dataloader() - Retorne validation DataLoader
5. test_dataloader() - Retorne test DataLoader

Referência de template rápido: Veja scripts/template_datamodule.py para boilerplate completo.

Documentação detalhada: Leia references/data_module.md para detalhes de métodos e padrões de uso.

4. Callbacks - Lógica de Treinamento Extensível

Adicione funcionalidade customizada em hooks específicos de treinamento sem modificar seu LightningModule. Callbacks built-in incluem:

• ModelCheckpoint - Salve modelos melhores/mais recentes
• EarlyStopping - Pare quando métricas platô
• LearningRateMonitor - Rastreie mudanças do LR scheduler
• BatchSizeFinder - Determine automaticamente batch size ótimo

Documentação detalhada: Leia references/callbacks.md para callbacks built-in e criação de callbacks customizados.

5. Logging - Rastreamento de Experimentos

Integre com múltiplas plataformas de logging:

• TensorBoard (padrão)
• Weights & Biases (WandbLogger)
• MLflow (MLFlowLogger)
• Neptune (NeptuneLogger)
• Comet (CometLogger)
• CSV (CSVLogger)

Log de métricas usando self.log("metric_name", value) em qualquer método de LightningModule.

Documentação detalhada: Leia references/logging.md para setup de loggers e configuração.

6. Distributed Training - Escale para Múltiplos Devices

Escolha a strategy correta baseada no tamanho do modelo:

• DDP - Para modelos <500M parâmetros (ResNet, transformers menores)
• FSDP - Para modelos 500M+ parâmetros (transformers grandes, recomendado para usuários Lightning)
• DeepSpeed - Para recursos de ponta e controle fine-grained

Configure com: Trainer(strategy="ddp", accelerator="gpu", devices=4)

Documentação detalhada: Leia references/distributed_training.md para comparação de strategies e configuração.

7. Boas Práticas

• Código device agnostic - Use self.device ao invés de .cuda()
• Salvamento de hiperparâmetros - Use self.save_hyperparameters() em __init__()
• Logging de métricas - Use self.log() para agregação automática entre devices
• Reproducibilidade - Use seed_everything() e Trainer(deterministic=True)
• Debugging - Use Trainer(fast_dev_run=True) para testar com 1 batch

Documentação detalhada: Leia references/best_practices.md para padrões comuns e armadilhas.

Workflow Rápido

1. Defina o modelo:

   class MyModel(L.LightningModule):
       def __init__(self):
           super().__init__()
           self.save_hyperparameters()
           self.model = YourNetwork()
   
       def training_step(self, batch, batch_idx):
           x, y = batch
           loss = F.cross_entropy(self.model(x), y)
           self.log("train_loss", loss)
           return loss
   
       def configure_optimizers(self):
           return torch.optim.Adam(self.parameters())
   

2. Prepare dados:

   # Option 1: DataLoaders diretos
   train_loader = DataLoader(train_dataset, batch_size=32)
   
   # Option 2: LightningDataModule (recomendado para reusabilidade)
   dm = MyDataModule(batch_size=32)
   

3. Treine:

   trainer = L.Trainer(max_epochs=10, accelerator="gpu", devices=2)
   trainer.fit(model, train_loader)  # ou trainer.fit(model, datamodule=dm)
   

Recursos

scripts/

Templates Python executáveis para padrões comuns de PyTorch Lightning:

• template_lightning_module.py - Boilerplate completo de LightningModule
• template_datamodule.py - Boilerplate completo de LightningDataModule
• quick_trainer_setup.py - Exemplos de configuração comum de Trainer

references/

Documentação detalhada para cada componente de PyTorch Lightning:

• lightning_module.md - Guia abrangente de LightningModule (métodos, hooks, propriedades)
• trainer.md - Configuração de Trainer e parâmetros
• data_module.md - Padrões e métodos de LightningDataModule
• callbacks.md - Callbacks built-in e customizados
• logging.md - Integrações de loggers e uso
• distributed_training.md - Comparação de DDP, FSDP, DeepSpeed e setup
• best_practices.md - Padrões comuns, dicas e armadilhas