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. While the library is self-contained, it is possible to use the models outside of solo-learn.

• [Oct 10 2021]: 👹 Restructured augmentation pipelines to allow more flexibility and multicrop. Also added multicrop for BYOL.
• [Sep 27 2021]: 🍕 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]: 🦘 Added online k-NN evaluation.
• [Sep 17 2021]: 🤖 Added ViT and Swin.
• [Sep 13 2021]: 📖 Improved Docs and added tutorials for pretraining and offline linear eval.
• [Aug 13 2021]: 🐳 DeepCluster V2 is now available.
• [Jul 31 2021]: 🦔 ReSSL is now available.
• [Jul 21 2021]: 🧪 Added Custom Dataset support.
• [Jul 21 2021]: 🎠 Added AutoUMAP.

• Barlow Twins
• BYOL
• DeepCluster V2
• DINO
• MoCo V2+
• NNBYOL
• NNCLR
• NNSiam
• ReSSL
• SimCLR + Supervised Contrastive Learning
• SimSiam
• Swav
• VICReg
• W-MSE

• Increased data processing speed by up to 100% using Nvidia Dali.
• Asymmetric and symmetric augmentations.

• Online linear evaluation via stop-gradient for easier debugging and prototyping (optionally available for the momentum encoder as well).
• Online Knn evaluation.
• Normal offline linear evaluation.
• All the perks of PyTorch Lightning (mixed precision, gradient accumulation, clipping, automatic logging and much more).
• Easy-to-extend modular code structure.
• Custom model logging with a simpler file organization.
• Automatic feature space visualization with UMAP.
• Common metrics and more to come...

• Multi-cropping dataloading following SwAV:
  • Note: currently, only SimCLR supports this.
• Exclude batchnorm and biases from LARS.
• No LR scheduler for the projection head in SimSiam.

• torch
• tqdm
• einops
• wandb
• pytorch-lightning
• lightning-bolts

Optional:

• nvidia-dali

NOTE: if you are using CUDA 10.X change nvidia-dali-cuda110 to nvidia-dali-cuda100 in setup.py, line 7.

To install the repository with Dali and/or UMAP support, use:

pip3 install .[dali,umap]

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

pip3 install .

NOTE: If you want to modify the library, install it in dev mode with -e.

NOTE 2: Soon to be on pip.

For pretraining the encoder, follow one of the many bash files in bash_files/pretrain/.

After that, for offline linear evaluation, follow the examples on bash_files/linear.

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.

Note: hyperparameters may not be the best, we will be re-running the methods with lower performance eventually.

🚀 methods where hyperparameters were heavily tuned.

😪 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!

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.

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

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

@misc{turrisi2021sololearn,
      title={Solo-learn: A Library of Self-supervised Methods for Visual Representation Learning}, 
      author={Victor G. Turrisi da Costa and Enrico Fini and Moin Nabi and Nicu Sebe and Elisa Ricci},
      year={2021},
      eprint={2108.01775},
      archivePrefix={arXiv},
      primaryClass={cs.CV},
      url={\url{https://github.com/vturrisi/solo-learn}},
}