PyTorch Lightning

General

PyTorch Lightning is a lightweight wrapper around PyTorch that separates research code from engineering code. You write the model logic in a LightningModule. The Trainer handles the loop, devices, precision, checkpointing and logging. No more .to(device), loss.backward() or optimizer.step() by hand.

Older code uses import pytorch_lightning as pl. Same API, legacy package name.

LightningModule hooks

Prepare data

A LightningDataModule bundles download, split, transforms and the DataLoaders so the same data pipeline is reusable across experiments. Plain DataLoaders passed to trainer.fit() work too.

class MNISTData(L.LightningDataModule):
    def __init__(self, data_dir="./data", batch_size=64):
        super().__init__()
        self.data_dir, self.bs = data_dir, batch_size

    def prepare_data(self):     # download once (rank 0 only)
        MNIST(self.data_dir, download=True)

    def setup(self, stage):     # split/transform on every GPU
        full = MNIST(self.data_dir, transform=ToTensor())
        self.train, self.val = random_split(full, [55000, 5000])

    def train_dataloader(self):
        return DataLoader(self.train, batch_size=self.bs,
                          shuffle=True, num_workers=4)

    def val_dataloader(self):
        return DataLoader(self.val, batch_size=self.bs)

Callbacks (lightning.pytorch.callbacks)

Define LightningModule

A LightningModule is an nn.Module with training hooks added. Any PyTorch model, e.g. an nn.Sequential, goes inside it.

import torch, torch.nn as nn, torch.nn.functional as F
from torch.optim.lr_scheduler import CosineAnnealingLR
import lightning as L

class LitClassifier(L.LightningModule):
    def __init__(self, lr=1e-3):
        super().__init__()
        self.save_hyperparameters()   # -> self.hparams.lr
        self.net = nn.Sequential(
            nn.Flatten(), nn.Linear(28*28, 128),
            nn.ReLU(), nn.Dropout(0.2), nn.Linear(128, 10))

    def forward(self, x):
        return self.net(x)

    def training_step(self, batch, batch_idx):
        x, y = batch                  # already on device
        loss = F.cross_entropy(self(x), y)
        self.log("train_loss", loss, prog_bar=True)
        return loss          # Lightning does backward + step

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

    def configure_optimizers(self):
        opt = torch.optim.Adam(self.parameters(), lr=self.hparams.lr)
        sched = CosineAnnealingLR(opt, T_max=self.trainer.max_epochs)
        return {"optimizer": opt, "lr_scheduler": sched}

Logging

Loggers (lightning.pytorch.loggers)

Pass with L.Trainer(logger=WandbLogger(...)). For robust metrics use torchmetrics, e.g. Accuracy(task="multiclass", num_classes=10).

Configure Trainer & fit

from lightning.pytorch.callbacks import (
    ModelCheckpoint, EarlyStopping)

ckpt = ModelCheckpoint(monitor="val_loss", mode="min")
stop = EarlyStopping(monitor="val_loss", patience=3)
trainer = L.Trainer(
    max_epochs=10,
    accelerator="auto",    # "gpu", "cpu", "mps"
    devices="auto",        # 1, 4, [0, 2], -1 = all
    precision="16-mixed",  # "bf16-mixed", "32-true"
    callbacks=[ckpt, stop])
model, dm = LitClassifier(lr=1e-3), MNISTData()
trainer.fit(model, datamodule=dm)   # train + val
trainer.test(model, datamodule=dm)  # test_step
preds = trainer.predict(model, loader)

Useful Trainer flags

Debugging flags

Save/Load checkpoint

A .ckpt holds weights, optimizer, scheduler, epoch & hparams. ModelCheckpoint writes them to the log folder.

model = LitClassifier.load_from_checkpoint("best.ckpt")
model.eval()                      # for inference
trainer.fit(model, datamodule=dm, ckpt_path="last")
trainer.save_checkpoint("manual.ckpt")

Beyond the Trainer