solo-learn

April 22, 2026 ยท View on GitHub

tests Documentation Status codecov

solo-learn

A library of self-supervised methods for unsupervised visual representation learning powered by PyTorch Lightning. We aim at providing SOTA self-supervised methods in a comparable environment while, at the same time, implementing training tricks. The library is self-contained, but it is possible to use the models outside of solo-learn. More details in our paper.

News

  • [Jan 14 2024]: :clap: Bunch of stability improvements during 2023 :) Also added All4One.
  • [Jan 07 2023]: :diving_mask: Added results, checkpoints and configs for MAE on ImageNet. Thanks to HuangChiEn.
  • [Dec 31 2022]: :stars: Shiny new logo! Huge thanks to Luiz!
  • [Sep 27 2022]: :pencil: Brand new config system using OmegaConf/Hydra. Adds more clarity and flexibility. New tutorials will follow soon!
  • [Aug 04 2022]: :paintbrush: Added MAE and supports finetuning the backbone with main_linear.py, mixup, cutmix and random augment.
  • [Jul 13 2022]: :sparkling_heart: Added support for H5 data, improved scripts and data handling.
  • [Jun 26 2022]: :fire: Added MoCo V3.
  • [Jun 10 2022]: :bomb: Improved LARS.
  • [Jun 09 2022]: :lollipop: Added support for WideResnet, multicrop for SwAV and equalization data augmentation.
  • [May 02 2022]: :diamond_shape_with_a_dot_inside: Wrapped Dali with a DataModule, added auto resume for linear eval and Wandb run resume.
  • [Apr 12 2022]: :rainbow: Improved design of models and added support to train with a fraction of data.
  • [Apr 01 2022]: :mag: Added the option to use channel last conversion which considerably decreases training times.
  • [Feb 04 2022]: :partying_face: Paper got accepted to JMLR.
  • [Jan 31 2022]: :eye: Added ConvNeXt support with timm.
  • [Dec 20 2021]: :thermometer: Added ImageNet results, scripts and checkpoints for MoCo V2+.
  • [Dec 05 2021]: :notes: Separated SupCon from SimCLR and added runs.
  • [Dec 01 2021]: :fountain: Added PoolFormer.
  • [Nov 29 2021]: :bangbang: Breaking changes! Update your versions!!!
  • [Nov 29 2021]: :book: New tutorials!
  • [Nov 29 2021]: :houses: Added offline K-NN and offline UMAP.
  • [Nov 29 2021]: :rotating_light: Updated PyTorch and PyTorch Lightning versions. 10% faster.
  • [Nov 29 2021]: :beers: Added code of conduct, contribution instructions, issue templates and UMAP tutorial.
  • [Nov 23 2021]: :space_invader: Added VIbCReg.
  • [Oct 21 2021]: :triumph: Added support for object recognition via Detectron v2 and auto resume functionally that automatically tries to resume an experiment that crashed/reached a timeout.
  • [Oct 10 2021]: :japanese_ogre: Restructured augmentation pipelines to allow more flexibility and multicrop. Also added multicrop for BYOL.
  • [Sep 27 2021]: :pizza: Added NNSiam, NNBYOL, new tutorials for implementing new methods 1 and 2, more testing and fixed issues with custom data and linear evaluation.
  • [Sep 19 2021]: :kangaroo: Added online k-NN evaluation.
  • [Sep 17 2021]: :robot: Added ViT and Swin.
  • [Sep 13 2021]: :book: Improved Docs and added tutorials for pretraining and offline linear eval.
  • [Aug 13 2021]: :whale: DeepCluster V2 is now available.

Roadmap and help needed

  • Redoing the documentation to improve clarity.
  • Better and up-to-date tutorials.
  • Add performance-related testing to ensure that methods perform the same across updates.
  • Adding new methods (continuous effort).

Methods available

Extra flavor

Backbones

Data

  • Increased data processing speed by up to 100% using Nvidia Dali.
  • Flexible augmentations.

Evaluation

  • Online linear evaluation via stop-gradient for easier debugging and prototyping (optionally available for the momentum backbone as well).
  • Standard offline linear evaluation.
  • Online and offline K-NN evaluation.
  • Automatic feature space visualization with UMAP.

Training tricks

  • All the perks of PyTorch Lightning (mixed precision, gradient accumulation, clipping, and much more).
  • Channel last conversion
  • Multi-cropping dataloading following SwAV:
    • Note: currently, only SimCLR, BYOL and SwAV support this.
  • Exclude batchnorm and biases from weight decay and LARS.
  • No LR scheduler for the projection head (as in SimSiam).

Logging

  • Metric logging on the cloud with WandB
  • Custom model checkpointing with a simple file organization.

Requirements

  • torch
  • torchvision
  • tqdm
  • einops
  • wandb
  • pytorch-lightning
  • lightning-bolts
  • torchmetrics
  • scipy
  • timm

Optional:

  • nvidia-dali
  • matplotlib
  • seaborn
  • pandas
  • umap-learn

Installation

First clone the repo.

Then, to install solo-learn with Dali and/or UMAP support, use:

pip3 install .[dali,umap,h5] --extra-index-url https://developer.download.nvidia.com/compute/redist

If no Dali/UMAP/H5 support is needed, the repository can be installed as:

pip3 install .

For local development:

pip3 install -e .[umap,h5]
# Make sure you have pre-commit hooks installed
pre-commit install

NOTE: if you are having trouble with dali, install it following their guide.

NOTE 2: consider installing Pillow-SIMD for better loading times when not using Dali.

NOTE 3: Soon to be on pip.

Training

For pretraining the backbone, follow one of the many bash files in scripts/pretrain/. We are now using Hydra to handle the config files, so the common syntax is something like:

python3 main_pretrain.py \
    # path to training script folder
    --config-path scripts/pretrain/imagenet-100/ \
    # training config name
    --config-name barlow.yaml
    # add new arguments (e.g. those not defined in the yaml files)
    # by doing ++new_argument=VALUE
    # pytorch lightning's arguments can be added here as well.

After that, for offline linear evaluation, follow the examples in scripts/linear or scripts/finetune for finetuning the whole backbone.

For k-NN evaluation and UMAP visualization check the scripts in scripts/{knn,umap}.

NOTE: Files try to be up-to-date and follow as closely as possible the recommended parameters of each paper, but check them before running.

Tutorials

Please, check out our documentation and tutorials:

If you want to contribute to solo-learn, make sure you take a look at how to contribute and follow the code of conduct

Model Zoo

Checkpoints are no longer available as of 2026. Re-running the scripts will yield the same results indicated below.

Results

CIFAR-10

MethodBackboneEpochsDaliAcc@1Acc@5
All4OneResNet181000:x:93.2499.88
Barlow TwinsResNet181000:x:92.1099.73
BYOLResNet181000:x:92.5899.79
DeepCluster V2ResNet181000:x:88.8599.58
DINOResNet181000:x:89.5299.71
MoCo V2+ResNet181000:x:92.9499.79
MoCo V3ResNet181000:x:93.1099.80
NNCLRResNet181000:x:91.8899.78
ReSSLResNet181000:x:90.6399.62
SimCLRResNet181000:x:90.7499.75
SimsiamResNet181000:x:90.5199.72
SupConResNet181000:x:93.8299.65
SwAVResNet181000:x:89.1799.68
VIbCRegResNet181000:x:91.1899.74
VICRegResNet181000:x:92.0799.74
W-MSEResNet181000:x:88.6799.68

CIFAR-100

MethodBackboneEpochsDaliAcc@1Acc@5
All4OneResNet181000:x:72.1793.35
Barlow TwinsResNet181000:x:70.9091.91
BYOLResNet181000:x:70.4691.96
DeepCluster V2ResNet181000:x:63.6188.09
DINOResNet181000:x:66.7690.34
MoCo V2+ResNet181000:x:69.8991.65
MoCo V3ResNet181000:x:68.8390.57
NNCLRResNet181000:x:69.6291.52
ReSSLResNet181000:x:65.9289.73
SimCLRResNet181000:x:65.7889.04
SimsiamResNet181000:x:66.0489.62
SupConResNet181000:x:70.3889.57
SwAVResNet181000:x:64.8888.78
VIbCRegResNet181000:x:67.3790.07
VICRegResNet181000:x:68.5490.83
W-MSEResNet181000:x:61.3387.26

ImageNet-100

MethodBackboneEpochsDaliAcc@1 (online)Acc@1 (offline)Acc@5 (online)Acc@5 (offline)
All4OneResNet18400:heavy_check_mark:81.93-96.23-
Barlow Twins :rocket:ResNet18400:heavy_check_mark:80.3880.1695.2895.14
BYOL :rocket:ResNet18400:heavy_check_mark:80.1680.3295.0294.94
DeepCluster V2ResNet18400:x:75.3675.4093.2293.10
DINOResNet18400:heavy_check_mark:74.8474.9292.9292.78
DINO :sleepy:ViT Tiny400:x:63.04-87.72-
MoCo V2+ :rocket:ResNet18400:heavy_check_mark:78.2079.2895.5095.18
MoCo V3 :rocket:ResNet18400:heavy_check_mark:80.3680.3695.1894.96
MoCo V3 :rocket:ResNet50400:heavy_check_mark:85.4884.5896.8296.70
NNCLR :rocket:ResNet18400:heavy_check_mark:79.8080.1695.2895.30
ReSSLResNet18400:heavy_check_mark:76.9278.4894.2094.24
SimCLR :rocket:ResNet18400:heavy_check_mark:77.64-94.06-
SimsiamResNet18400:heavy_check_mark:74.5478.7293.1694.78
SupConResNet18400:heavy_check_mark:84.40-95.72-
SwAVResNet18400:heavy_check_mark:74.0474.2892.7092.84
VIbCRegResNet18400:heavy_check_mark:79.8679.3894.9894.60
VICReg :rocket:ResNet18400:heavy_check_mark:79.2279.4095.0695.02
W-MSEResNet18400:heavy_check_mark:67.6069.0690.9491.22

:rocket: methods where hyperparameters were heavily tuned.

:sleepy: ViT is very compute intensive and unstable, so we are slowly running larger architectures and with a larger batch size. Atm, total batch size is 128 and we needed to use float32 precision. If you want to contribute by running it, let us know!

ImageNet

MethodBackboneEpochsDaliAcc@1 (online)Acc@1 (offline)Acc@5 (online)Acc@5 (offline)
Barlow TwinsResNet50100:heavy_check_mark:67.1867.2387.6987.98
BYOLResNet50100:heavy_check_mark:68.6368.3788.8088.66
MoCo V2+ResNet50100:heavy_check_mark:62.6166.8485.4087.60
MAEViT-B/16100:x:~81.60 (finetuned)~95.50 (finetuned)

Training efficiency for DALI

We report the training efficiency of some methods using a ResNet18 with and without DALI (4 workers per GPU) in a server with an Intel i9-9820X and two RTX2080ti.

MethodDaliTotal time for 20 epochsTime for 1 epochGPU memory (per GPU)
Barlow Twins:x:1h 38m 27s4m 55s5097 MB
:heavy_check_mark:43m 2s2m 10s (56% faster)9292 MB
BYOL:x:1h 38m 46s4m 56s5409 MB
:heavy_check_mark:50m 33s2m 31s (49% faster)9521 MB
NNCLR:x:1h 38m 30s4m 55s5060 MB
:heavy_check_mark:42m 3s2m 6s (64% faster)9244 MB

Note: GPU memory increase doesn't scale with the model, rather it scales with the number of workers.

Citation

If you use solo-learn, please cite our paper:

@article{JMLR:v23:21-1155,
  author  = {Victor Guilherme Turrisi da Costa and Enrico Fini and Moin Nabi and Nicu Sebe and Elisa Ricci},
  title   = {solo-learn: A Library of Self-supervised Methods for Visual Representation Learning},
  journal = {Journal of Machine Learning Research},
  year    = {2022},
  volume  = {23},
  number  = {56},
  pages   = {1-6},
  url     = {http://jmlr.org/papers/v23/21-1155.html}
}