Pytorch Lightning Guide

# Install PyTorch Lightning
pip install lightning

# Or install with specific extras
pip install lightning[extra]

# For development/research with all features
pip install lightning[all]

# Check GPU availability
python -c "import torch; print(torch.cuda.is_available())"

import lightning as L
print(L.__version__)

import lightning as L
import torch
import torch.nn.functional as F
from torch import nn

class ResearchModel(L.LightningModule):
    def __init__(self, input_dim, hidden_dim, output_dim, lr=1e-3):
        super().__init__()
        self.save_hyperparameters()
        self.encoder = nn.Sequential(
            nn.Linear(input_dim, hidden_dim),
            nn.ReLU(),
            nn.Dropout(0.2),
            nn.Linear(hidden_dim, hidden_dim),
            nn.ReLU(),
        )
        self.classifier = nn.Linear(hidden_dim, output_dim)
        self.lr = lr

    def forward(self, x):
        features = self.encoder(x)
        return self.classifier(features)

    def training_step(self, batch, batch_idx):
        x, y = batch
        logits = self(x)
        loss = F.cross_entropy(logits, y)
        acc = (logits.argmax(dim=-1) == y).float().mean()
        self.log("train_loss", loss, prog_bar=True)
        self.log("train_acc", acc, prog_bar=True)
        return loss

    def validation_step(self, batch, batch_idx):
        x, y = batch
        logits = self(x)
        loss = F.cross_entropy(logits, y)
        acc = (logits.argmax(dim=-1) == y).float().mean()
        self.log("val_loss", loss, prog_bar=True)
        self.log("val_acc", acc, prog_bar=True)

    def configure_optimizers(self):
        optimizer = torch.optim.AdamW(self.parameters(), lr=self.lr)
        scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(
            optimizer, T_max=self.trainer.max_epochs
        )
        return [optimizer], [scheduler]

class ResearchDataModule(L.LightningDataModule):
    def __init__(self, data_dir, batch_size=32, num_workers=4):
        super().__init__()
        self.data_dir = data_dir
        self.batch_size = batch_size
        self.num_workers = num_workers

    def setup(self, stage=None):
        # Load and split data
        dataset = load_research_dataset(self.data_dir)
        self.train_data, self.val_data, self.test_data = random_split(
            dataset, [0.8, 0.1, 0.1]
        )

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

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

trainer = L.Trainer(
    max_epochs=100,
    accelerator="gpu",
    devices=4,
    strategy="ddp",
    precision="16-mixed",
    gradient_clip_val=1.0,
    accumulate_grad_batches=4,
    callbacks=[
        L.callbacks.EarlyStopping(monitor="val_loss", patience=10),
        L.callbacks.ModelCheckpoint(monitor="val_loss", save_top_k=3),
        L.callbacks.LearningRateMonitor(),
    ],
    logger=L.loggers.WandbLogger(project="my-research"),
)

# Train the model
trainer.fit(model, datamodule=data_module)

# Test with best checkpoint
trainer.test(model, datamodule=data_module, ckpt_path="best")

# FSDP for large model training
trainer = L.Trainer(
    strategy="fsdp",
    devices=8,
    precision="bf16-mixed",
)

class GANModule(L.LightningModule):
    def training_step(self, batch, batch_idx):
        optimizer_g, optimizer_d = self.optimizers()

        # Train discriminator
        real_loss = self.discriminator_loss(batch, real=True)
        fake_loss = self.discriminator_loss(batch, real=False)
        d_loss = (real_loss + fake_loss) / 2
        optimizer_d.zero_grad()
        self.manual_backward(d_loss)
        optimizer_d.step()

        # Train generator
        g_loss = self.generator_loss(batch)
        optimizer_g.zero_grad()
        self.manual_backward(g_loss)
        optimizer_g.step()

    @property
    def automatic_optimization(self):
        return False

trainer = L.Trainer(
    profiler="advanced",  # or "simple", "pytorch"
    detect_anomaly=True,
    overfit_batches=10,   # Quick sanity check
)

# Seed everything for reproducibility
L.seed_everything(42, workers=True)

# Hyperparameters are automatically saved
model = ResearchModel(input_dim=768, hidden_dim=256, output_dim=10)
# model.hparams is automatically populated and logged

# Checkpoints include full training state
# Resume training from a checkpoint
trainer.fit(model, ckpt_path="path/to/checkpoint.ckpt")