04 · DataModules & Callbacks 🧰
Lightning can absorb two more chunks of boilerplate you hand-coded earlier: the data pipeline from Module 04 and the checkpointing/early-stopping from Module 06.
1. The 'Why'
The LightningModule and Trainer cleaned up your model and loop, but two other pieces of your projects are still ad-hoc: the data pipeline and the training-management logic. In Module 04 you wrote dataset downloads, splits, transforms, and DataLoader construction, usually scattered across a script. In Module 06 you hand-coded best-model saving and early stopping with manual bookkeeping. Both work, but both are boilerplate that's easy to get subtly wrong and tedious to reproduce across projects. Lightning offers clean, reusable homes for each: the LightningDataModule and callbacks.
A LightningDataModule bundles all your data logic — downloading, splitting, transforms, and the three DataLoaders — into one self-contained, shareable class. Hand it to the Trainer and your data setup becomes a portable component you can drop into any project or share with a teammate, with train/val/test cleanly separated (guarding against the leakage you learned to fear). Callbacks, meanwhile, are hooks that run at defined points in the training loop, and Lightning ships battle-tested ones for exactly the tasks you did by hand: ModelCheckpoint saves your best model automatically by monitoring a metric, and EarlyStopping halts training when it stops improving. Together these complete the refactor: everything you built across the course — data, model, loop, checkpointing — now has a clean, standard place to live.
2. Core Concepts
LightningDataModule. A class with standardized methods: prepare_data (download once), setup (build datasets/splits/transforms), and train_dataloader/val_dataloader/test_dataloader (return the loaders). The Trainer calls them at the right time.
Why a DataModule. It makes your data pipeline reproducible and shareable, keeps splits cleanly separated, and decouples data from model so you can swap either independently.
Callbacks. Self-contained objects that hook into loop events (epoch end, etc.) without cluttering your LightningModule. You pass a list of them to the Trainer.
ModelCheckpoint. Monitors a logged metric (e.g. val_acc) and automatically saves the best checkpoint(s) — the Module 06 best-model pattern, automated and configurable.
EarlyStopping. Monitors a metric and stops training after a patience window without improvement — the Module 06 early-stopping logic, as a one-line callback.
3. Code in Action
A LightningDataModule (packaging Module 04's pipeline)
import lightning as L
from torch.utils.data import DataLoader, random_split
from torchvision import datasets, transforms
class FashionMNISTData(L.LightningDataModule):
def __init__(self, data_dir="./data", batch_size=64):
super().__init__()
self.data_dir = data_dir
self.batch_size = batch_size
self.tfms = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize((0.2860,), (0.3530,)),
])
def prepare_data(self): # download ONCE (no state set here)
datasets.FashionMNIST(self.data_dir, train=True, download=True)
datasets.FashionMNIST(self.data_dir, train=False, download=True)
def setup(self, stage=None): # build splits (runs on each device)
full = datasets.FashionMNIST(self.data_dir, train=True, transform=self.tfms)
self.train_ds, self.val_ds = random_split(full, [55000, 5000])
self.test_ds = datasets.FashionMNIST(self.data_dir, train=False, transform=self.tfms)
def train_dataloader(self):
return DataLoader(self.train_ds, batch_size=self.batch_size, shuffle=True)
def val_dataloader(self):
return DataLoader(self.val_ds, batch_size=self.batch_size)
def test_dataloader(self):
return DataLoader(self.test_ds, batch_size=self.batch_size)Callbacks that automate Module 06
from lightning.pytorch.callbacks import ModelCheckpoint, EarlyStopping
checkpoint = ModelCheckpoint(
monitor="val_acc", # watch the metric you logged in validation_step
mode="max", # higher accuracy is better
save_top_k=1, # keep only the single best model
filename="best-{epoch}-{val_acc:.3f}",
)
early_stop = EarlyStopping(
monitor="val_loss", # watch validation loss
mode="min", # lower is better
patience=5, # stop after 5 epochs without improvement
)Putting it all together
import lightning as L
model = LitClassifier() # sub-module 02
data = FashionMNISTData() # the DataModule above
trainer = L.Trainer(
max_epochs=20,
accelerator="auto",
callbacks=[checkpoint, early_stop], # <-- automated checkpointing + early stopping
)
trainer.fit(model, datamodule=data) # pass the DataModule directly
trainer.test(model, datamodule=data) # test split comes from the same component
print("Best model saved at:", checkpoint.best_model_path)4. Common Pitfalls
Pitfall 1 — Setting state in prepare_data. prepare_data runs once on a single process and must only download — assigning self.train_ds there breaks in multi-GPU settings. Build datasets in setup, which runs on every process. Mixing these up is the classic DataModule bug.
Pitfall 2 — Monitoring a metric you never logged. ModelCheckpoint(monitor="val_acc") and EarlyStopping(monitor="val_loss") only work if your validation_step actually calls self.log("val_acc", ...) / self.log("val_loss", ...) with those exact names. A typo or missing log silently disables the callback (or errors). Keep logged names and monitored names in sync.
Pitfall 3 — Reinventing what callbacks already do. After learning Lightning, it's tempting to hand-code best-model saving inside validation_step (Module 06 habit). That fights the framework and duplicates logic. Use ModelCheckpoint/EarlyStopping — they're tested, configurable, and multi-device-safe.
5. Further Reading & Watch List
- 📖 Deep Learning with PyTorch + Lightning docs for production workflows
- 🌐 Lightning — LightningDataModule
- 🛠️
ModelCheckpoint·EarlyStopping· Callbacks overview
⬅️ Prev: 03 · The Trainer · 🏁 Module complete! Back to the module index or on to Module 10 — Deployment & Optimization.