pytorch-lightning-guide

PyTorch Lightning Guide

Overview

PyTorch Lightning is a deep learning framework with over 31,000 GitHub stars that provides a high-level interface for PyTorch, enabling researchers to focus on model design rather than engineering boilerplate. Developed by Lightning AI, it decouples the science (model architecture, loss functions, data processing) from the engineering (distributed training, mixed precision, gradient accumulation, checkpointing) through a structured LightningModule abstraction.

For academic researchers, Lightning eliminates the need to write repetitive training loops, device management code, and distributed training logic. You define your model, training step, and data loaders, and Lightning handles everything else -- from single GPU to multi-node distributed training, from FP32 to mixed precision, from local development to cloud deployment. This means faster iteration on research ideas with production-quality training infrastructure.

Lightning is used extensively in AI research labs and has become a standard tool for reproducible deep learning experiments. It integrates seamlessly with experiment tracking tools like Weights & Biases, MLflow, and TensorBoard, and supports all PyTorch-compatible model architectures.

Installation and Setup

# 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]

Lightning requires Python 3.9+ and PyTorch 2.1+. For GPU training, ensure your PyTorch installation includes CUDA support:

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

Verify your installation:

import lightning as L
print(L.__version__)

Core Architecture

The LightningModule

The LightningModule is the central abstraction. It organizes your PyTorch code into clearly defined methods:

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]