PyTorch Lightning Cheat Sheet

1Install & the big picture2.5+

pip install lightning★
Imports as import lightning as L. The standalone pip install pytorch-lightning still works (import pytorch_lightning as pl).
class LitModel(L.LightningModule): ...★
Your model + training logic live here (the "science"). No manual .backward(), .zero_grad(), .to(device), or .eval().
trainer = L.Trainer(...); trainer.fit(model, dm)★
The Trainer owns the loop, devices, precision, logging & checkpoints (the "engineering").
L.seed_everything(42, workers=True)
Seed Python/NumPy/PyTorch (and DataLoader workers) for reproducibility.

2LightningModule hooksthe 6 sections

def __init__(self): super().__init__(); self.save_hyperparameters()★
save_hyperparameters() stores ctor args in self.hparams and the checkpoint — enables clean reload.
def training_step(self, batch, batch_idx): loss = ...; self.log("train_loss", loss); return loss★
Define one step and return the loss — Lightning does backward + optimizer.step for you.
def validation_step(self, batch, idx): ... # + test_step, predict_step★
Same shape as training_step; no grad needed. predict_step powers trainer.predict().
def forward(self, x): return self.net(x)
Optional — defines what model(x) does for inference. Keep training logic in the *_step hooks.
on_train_epoch_end() · on_validation_epoch_end() ...
Epoch-level hooks for aggregating metrics; there's a hook for nearly every point in the loop.

3configure_optimizersoptimizer + schedule

def configure_optimizers(self): return torch.optim.AdamW(self.parameters(), lr=self.hparams.lr)★
Return an optimizer (or several). Lightning calls it after the model is on the right device.
return {"optimizer": opt, "lr_scheduler": {"scheduler": sch, "monitor": "val_loss"}}
Attach an LR scheduler; monitor drives ReduceLROnPlateau. interval:"step" to step per-batch.
self.automatic_optimization = False
Opt out to hand-code multiple optimizers (GANs, RL): use opt = self.optimizers(), self.manual_backward(loss), opt.step().

4Trainer & fit/validate/testthe loop

trainer = L.Trainer(max_epochs=10, accelerator="auto", devices="auto")★
"auto" detects GPUs/TPUs. Also max_steps, min_epochs.
trainer.fit(model, train_loader, val_loader)★
Or trainer.fit(model, datamodule=dm). Resume with ckpt_path="last" / a path.
trainer.validate(model, dm) · trainer.test(model, dm)
test should be run once at the very end. trainer.predict(model, dm) returns predictions.
trainer = L.Trainer(fast_dev_run=True)★
Run a single train/val/test batch to smoke-test the whole pipeline before a long run.

5Key Trainer flagsmost-used knobs

precision="bf16-mixed" # or "16-mixed", "32-true", "16-true"★
Mixed precision for big speed/memory wins on modern GPUs.
accumulate_grad_batches=4★
Simulate a 4× larger batch on limited memory.
gradient_clip_val=1.0 · overfit_batches=0.01
Clip grads; overfit a tiny subset to sanity-check the model can learn.
limit_train_batches=0.1 · log_every_n_steps=50 · num_sanity_val_steps=2
Shorten epochs for debugging; control logging cadence and the pre-train val sanity check.
deterministic=True · profiler="simple"
Reproducible kernels; a quick profiler to find bottlenecks.

6LightningDataModuleshareable data

class DM(L.LightningDataModule): ...★
Encapsulates the whole data pipeline so it's reproducible & swappable across models.
def prepare_data(self): # download — runs once, single process
Downloads/tokenization that must happen once. Don't set state here (it's not called on every GPU).
def setup(self, stage): # split/assign datasets per process★
Runs on every process; build self.train_ds/val_ds here. stage is "fit"/"test"/"predict".
def train_dataloader(self): return DataLoader(self.train_ds, ...)★
Plus val_dataloader/test_dataloader/predict_dataloader. Set num_workers & pin_memory for throughput.

7Logging (self.log)metrics

self.log("val_loss", loss, prog_bar=True)★
Logs to the configured logger + progress bar. Lightning auto-reduces across steps/devices.
self.log("acc", a, on_step=False, on_epoch=True, sync_dist=True)★
on_epoch aggregates; sync_dist=True averages correctly across GPUs. self.log_dict({...}) for many at once.
Trainer(logger=TensorBoardLogger("logs/")) # or WandbLogger, CSVLogger
Swap the backend without touching your self.log calls.
pip install torchmetrics # self.acc = torchmetrics.Accuracy(...)
Use TorchMetrics for correct distributed metric computation instead of manual averaging.

8Callbacks & checkpointsautomate the run

from lightning.pytorch.callbacks import ModelCheckpoint, EarlyStopping★
The two you'll almost always use. Pass a list via Trainer(callbacks=[...]).
ModelCheckpoint(monitor="val_loss", mode="min", save_top_k=3)★
Auto-saves the best checkpoints by a monitored metric; save_last=True for resumable training.
EarlyStopping(monitor="val_loss", patience=5)★
Stop when the metric stops improving. LearningRateMonitor logs the LR schedule.
class MyCb(L.Callback): def on_train_epoch_end(self, trainer, pl_module): ...
Write custom callbacks by overriding hook methods — keeps side logic out of the model.

PyTorch Lightning cheat sheet