PyTorch Lightning

1. LightningModule - Model Definition

Organize PyTorch models into six logical sections:

1. Initialization - __init__() and 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()

Quick template reference: See scripts/template_lightning_module.py for a complete boilerplate.

Detailed documentation: Read references/lightning_module.md for comprehensive method documentation, hooks, properties, and best practices.

2. Trainer - Training Automation

The Trainer automates the training loop, device management, gradient operations, and callbacks. Key features:

• Multi-GPU/TPU support with strategy selection (DDP, FSDP, DeepSpeed)
• Automatic mixed precision training
• Gradient accumulation and clipping
• Checkpointing and early stopping
• Progress bars and logging

Quick setup reference: See scripts/quick_trainer_setup.py for common Trainer configurations.

Detailed documentation: Read references/trainer.md for all parameters, methods, and configuration options.

3. LightningDataModule - Data Pipeline Organization

Encapsulate all data processing steps in a reusable class:

1. prepare_data() - Download and process data (single-process)
2. setup() - Create datasets and apply transforms (per-GPU)
3. train_dataloader() - Return training DataLoader
4. val_dataloader() - Return validation DataLoader
5. test_dataloader() - Return test DataLoader

Quick template reference: See scripts/template_datamodule.py for a complete boilerplate.

Detailed documentation: Read references/data_module.md for method details and usage patterns.