stable-pretraining
stable-pretraining is a PyTorch Lightning library that lets researchers swap and compare self-supervised learning objectives, such as SimCLR, DINO, MAE, BYOL, VICReg and LeJEPA, while keeping the same training loop, evaluation and checkpointing.
stable-pretraining helps researchers compare and customize self-supervised learning methods without rebuilding the surrounding training infrastructure each time. A method is defined as a forward function returning a dictionary of outputs, online probes and representation diagnostics can be attached to monitor training, and everything runs through a Manager class that handles checkpointing, resuming and a local experiment registry.
The library builds on established third-party components, PyTorch, Lightning, Hugging Face and TorchMetrics, so it can focus on assembling a flexible, modular framework rather than reimplementing architectures, losses, metrics or augmentations from scratch. Its built-in method catalog includes SimCLR, DINO, MAE, BYOL, VICReg and LeJEPA, among others, with an experimental JAX backend available as an opt-in. A quickstart command runs one CPU epoch of SimCLR with an online classifier on synthetic images as an installation check, needing no dataset download or GPU, and a companion viewer displays training metrics and probe accuracy from the resulting run directory.