pytorch-lightning

model = MyModel()
optimizer = torch.optim.Adam(model.parameters())
model.to('cuda')
 
for epoch in range(max_epochs):
    for batch in train_loader:
        batch = batch.to('cuda')
        optimizer.zero_grad()
        loss = model(batch)
        loss.backward()
        optimizer.step()

class LitModel(L.LightningModule):
    def __init__(self):
        super().__init__()
        self.model = MyModel()
 
    def training_step(self, batch, batch_idx):
        loss = self.model(batch)  # No .to('cuda') needed!
        return loss
 
    def configure_optimizers(self):
        return torch.optim.Adam(self.parameters())
 
# Train
trainer = L.Trainer(max_epochs=10, accelerator='gpu')
trainer.fit(LitModel(), train_loader)

class LitModel(L.LightningModule):
    def __init__(self):
        super().__init__()
        self.model = MyModel()
 
    def training_step(self, batch, batch_idx):
        x, y = batch
        y_hat = self.model(x)
        loss = nn.functional.cross_entropy(y_hat, y)
        self.log('train_loss', loss)
        return loss
 
    def validation_step(self, batch, batch_idx):
        x, y = batch
        y_hat = self.model(x)
        val_loss = nn.functional.cross_entropy(y_hat, y)
        acc = (y_hat.argmax(dim=1) == y).float().mean()
        self.log('val_loss', val_loss)
        self.log('val_acc', acc)
 
    def test_step(self, batch, batch_idx):
        x, y = batch
        y_hat = self.model(x)
        test_loss = nn.functional.cross_entropy(y_hat, y)
        self.log('test_loss', test_loss)
 
    def configure_optimizers(self):
        return torch.optim.Adam(self.parameters(), lr=1e-3)
 
# Train with validation
trainer = L.Trainer(max_epochs=10)
trainer.fit(model, train_loader, val_loader)
 
# Test
trainer.test(model, test_loader)

# Same code as single GPU!
model = LitModel()
 
# 8 GPUs with DDP (automatic!)
trainer = L.Trainer(
    accelerator='gpu',
    devices=8,
    strategy='ddp'  # Or 'fsdp', 'deepspeed'
)
 
trainer.fit(model, train_loader)

# Single command, Lightning handles the rest
python train.py

from lightning.pytorch.callbacks import ModelCheckpoint, EarlyStopping, LearningRateMonitor
 
# Create callbacks
checkpoint = ModelCheckpoint(
    monitor='val_loss',
    mode='min',
    save_top_k=3,
    filename='model-{epoch:02d}-{val_loss:.2f}'
)
 
early_stop = EarlyStopping(
    monitor='val_loss',
    patience=5,
    mode='min'
)
 
lr_monitor = LearningRateMonitor(logging_interval='epoch')
 
# Add to Trainer
trainer = L.Trainer(
    max_epochs=100,
    callbacks=[checkpoint, early_stop, lr_monitor]
)
 
trainer.fit(model, train_loader, val_loader)

class LitModel(L.LightningModule):
    # ... (training_step, etc.)
 
    def configure_optimizers(self):
        optimizer = torch.optim.Adam(self.parameters(), lr=1e-3)
 
        # Cosine annealing
        scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(
            optimizer,
            T_max=100,
            eta_min=1e-5
        )
 
        return {
            'optimizer': optimizer,
            'lr_scheduler': {
                'scheduler': scheduler,
                'interval': 'epoch',  # Update per epoch
                'frequency': 1
            }
        }
 
# Learning rate auto-logged!
trainer = L.Trainer(max_epochs=100)
trainer.fit(model, train_loader)

# Add to training_step
def training_step(self, batch, batch_idx):
    if batch_idx == 0:
        print(f"Batch shape: {batch[0].shape}")
        print(f"Labels: {batch[1]}")
    loss = ...
    return loss

trainer = L.Trainer(
    accumulate_grad_batches=4,  # Effective batch = batch_size × 4
    precision='bf16'  # Or 'fp16', reduces memory 50%
)

# WRONG
trainer.fit(model, train_loader)
 
# CORRECT
trainer.fit(model, train_loader, val_loader)

# Test on CPU first
trainer = L.Trainer(accelerator='cpu', devices=1)
 
# Then GPU
trainer = L.Trainer(accelerator='gpu', devices=1)