pytorch-lightning

1# PyTorch Lightning - High-Level Training Framework
2 
3## Quick start
4 
5PyTorch Lightning organizes PyTorch code to eliminate boilerplate while maintaining flexibility.
6 
7**Installation**:
8```bash
9pip install lightning
10```
11 
12**Convert PyTorch to Lightning** (3 steps):
13 
14```python
15import lightning as L
16import torch
17from torch import nn
18from torch.utils.data import DataLoader, Dataset
19 
20# Step 1: Define LightningModule (organize your PyTorch code)
21class LitModel(L.LightningModule):
22    def __init__(self, hidden_size=128):
23        super().__init__()
24        self.model = nn.Sequential(
25            nn.Linear(28 * 28, hidden_size),
26            nn.ReLU(),
27            nn.Linear(hidden_size, 10)
28        )
29 
30    def training_step(self, batch, batch_idx):
31        x, y = batch
32        y_hat = self.model(x)
33        loss = nn.functional.cross_entropy(y_hat, y)
34        self.log('train_loss', loss)  # Auto-logged to TensorBoard
35        return loss
36 
37    def configure_optimizers(self):
38        return torch.optim.Adam(self.parameters(), lr=1e-3)
39 
40# Step 2: Create data
41train_loader = DataLoader(train_dataset, batch_size=32)
42 
43# Step 3: Train with Trainer (handles everything else!)
44trainer = L.Trainer(max_epochs=10, accelerator='gpu', devices=2)
45model = LitModel()
46trainer.fit(model, train_loader)
47```
48 
49**That's it!** Trainer handles:
50- GPU/TPU/CPU switching
51- Distributed training (DDP, FSDP, DeepSpeed)
52- Mixed precision (FP16, BF16)
53- Gradient accumulation
54- Checkpointing
55- Logging
56- Progress bars
57…