pytorch-lightning

name	pytorch-lightning
description	PyTorch Lightning — LightningModule, Trainer, callbacks, logging, and checkpointing patterns for deep learning training loops

PyTorch Lightning

LightningModule skeleton

import lightning as L
import torch
import torch.nn as nn

class MyModel(L.LightningModule):
    def __init__(self, lr: float = 1e-3):
        super().__init__()
        self.save_hyperparameters()  # saves all __init__ args to self.hparams
        self.net = nn.Linear(128, 1)
        self.loss_fn = nn.BCEWithLogitsLoss()

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.net(x)

    def training_step(self, batch: tuple, batch_idx: int) -> torch.Tensor:
        x, y = batch
        logits = self(x)
        loss = self.loss_fn(logits, y.float())
        self.log("train/loss", loss, on_step=False, on_epoch=True, prog_bar=True)
        return loss

    def validation_step(self, batch: tuple, batch_idx: int) -> None:
        x, y = batch
        logits = self(x)
        loss = self.loss_fn(logits, y.float())
        self.log("val/loss", loss, prog_bar=True)

    def configure_optimizers(self):
        optimizer = torch.optim.Adam(self.parameters(), lr=self.hparams.lr)
        scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(optimizer, patience=5)
        return {
            "optimizer": optimizer,
            "lr_scheduler": {"scheduler": scheduler, "monitor": "val/loss"},
        }

Trainer setup

from lightning.pytorch.callbacks import ModelCheckpoint, EarlyStopping
from lightning.pytorch.loggers import CSVLogger

checkpoint_cb = ModelCheckpoint(
    dirpath="checkpoints/",
    filename="best-{epoch}-{val/loss:.4f}",
    monitor="val/loss",
    mode="min",
    save_top_k=1,
)

early_stop_cb = EarlyStopping(monitor="val/loss", patience=10, mode="min")

logger = CSVLogger("logs/", name="experiment")

trainer = L.Trainer(
    max_epochs=100,
    accelerator="auto",      # auto-detects GPU/CPU
    devices=1,
    precision="16-mixed",    # AMP; use "32" for stability
    callbacks=[checkpoint_cb, early_stop_cb],
    logger=logger,
    log_every_n_steps=10,
    deterministic=True,      # reproducibility (slower)
)

trainer.fit(model, train_dataloaders=train_loader, val_dataloaders=val_loader)

Loading checkpoints

# Resume training
trainer.fit(model, ckpt_path="checkpoints/last.ckpt")

# Load for inference
model = MyModel.load_from_checkpoint("checkpoints/best.ckpt")
model.eval()

self.log rules

# on_step=True  → logs at each batch (noisy, good for debugging)
# on_epoch=True → aggregates over epoch (recommended for val metrics)
# sync_dist=True → required for multi-GPU DDP training

self.log("train/loss", loss, on_step=False, on_epoch=True)
self.log("val/auc", auc, on_epoch=True, sync_dist=True)

# Log dict at once
self.log_dict({"val/loss": loss, "val/auc": auc}, prog_bar=True)

LightningDataModule

class FraudDataModule(L.LightningDataModule):
    def __init__(self, data_dir: str, batch_size: int = 256):
        super().__init__()
        self.save_hyperparameters()

    def setup(self, stage: str) -> None:
        # called on each GPU; load/split data here
        self.train_ds = ...
        self.val_ds = ...

    def train_dataloader(self):
        return DataLoader(self.train_ds, batch_size=self.hparams.batch_size, shuffle=True)

    def val_dataloader(self):
        return DataLoader(self.val_ds, batch_size=self.hparams.batch_size)

Pitfalls

• Never call optimizer.zero_grad(), loss.backward(), or optimizer.step() manually — Lightning handles this.
• self.save_hyperparameters() must be called in __init__ to enable checkpoint loading with load_from_checkpoint.
• For graph models (DGL/PyG), override transfer_batch_to_device if default .to(device) doesn't work on custom batch types.
• Use precision="16-mixed" (not "16") for modern AMP — the old "16" uses deprecated GradScaler path.
• deterministic=True disables some CUDA kernels; set only when reproducibility outweighs speed.