solo-learn

solo-learn: a library of self-supervised methods for visual representation learning powered by Pytorch Lightning

1,424

186

1,424

View on GitHub

Top Related Projects

vissl

3,248

VISSL is FAIR's library of extensible, modular and scalable components for SOTA Self-Supervised Learning with images.

lightly

2,918

A python library for self-supervised learning on images.

lightning-bolts

1,678

Toolbox of models, callbacks, and datasets for AI/ML researchers.

pytorch-metric-learning

5,982

The easiest way to use deep metric learning in your application. Modular, flexible, and extensible. Written in PyTorch.

CMC

1,298

[arXiv 2019] "Contrastive Multiview Coding", also contains implementations for MoCo and InstDis

SimCLR

2,214

PyTorch implementation of SimCLR: A Simple Framework for Contrastive Learning of Visual Representations

Quick Overview

Solo-learn is a library of self-supervised methods for computer vision tasks, implemented in PyTorch Lightning. It provides a collection of state-of-the-art self-supervised learning algorithms, making it easier for researchers and practitioners to experiment with and compare different approaches in the field of computer vision.

Pros

Comprehensive collection of self-supervised learning methods
Built on PyTorch Lightning, offering easy integration and scalability
Well-documented and actively maintained
Includes pre-trained models and benchmarking tools

Cons

Focused primarily on computer vision tasks, limiting its applicability to other domains
May have a steeper learning curve for those unfamiliar with PyTorch Lightning
Some methods may require significant computational resources

Code Examples

Creating a BYOL model:

from solo.methods import BYOL

model = BYOL(
    backbone="resnet18",
    num_classes=10,
    proj_hidden_dim=2048,
    proj_output_dim=256,
    pred_hidden_dim=2048,
)

Training a SimCLR model:

from solo.methods import SimCLR
from solo.utils.pretrain_dataloader import prepare_data

model = SimCLR(backbone="resnet50", num_classes=1000, proj_hidden_dim=2048, proj_output_dim=128)
train_loader, val_loader = prepare_data("cifar10", batch_size=256)

trainer = pl.Trainer(max_epochs=100, gpus=1)
trainer.fit(model, train_loader, val_loader)

Evaluating a pre-trained model:

from solo.methods import BarlowTwins
from solo.utils.classification_dataloader import prepare_data

model = BarlowTwins.load_from_checkpoint("path/to/checkpoint.ckpt")
test_loader = prepare_data("imagenet", batch_size=64, num_workers=4)

trainer = pl.Trainer(gpus=1)
trainer.test(model, test_loader)

Getting Started

To get started with solo-learn, follow these steps:

Install the library:

pip install solo-learn

Import the desired method and create a model:

from solo.methods import SimSiam

model = SimSiam(backbone="resnet18", num_classes=10, proj_hidden_dim=2048, pred_hidden_dim=512)

Prepare your data and train the model:

from solo.utils.pretrain_dataloader import prepare_data
import pytorch_lightning as pl

train_loader, val_loader = prepare_data("cifar10", batch_size=256)
trainer = pl.Trainer(max_epochs=100, gpus=1)
trainer.fit(model, train_loader, val_loader)

Competitor Comparisons

vissl

3,248

VISSL is FAIR's library of extensible, modular and scalable components for SOTA Self-Supervised Learning with images.

Pros of VISSL

More comprehensive and feature-rich, offering a wider range of self-supervised learning algorithms
Backed by Facebook AI Research, potentially providing more robust support and regular updates
Includes advanced features like distributed training and mixed precision support

Cons of VISSL

Steeper learning curve due to its complexity and extensive feature set
Heavier codebase, which may be overkill for simpler projects or beginners
Less focused on specific self-supervised learning methods compared to solo-learn

Code Comparison

VISSL example:

from vissl.engines.train import train_main
from vissl.utils.hydra_config import compose_hydra_configuration, convert_to_attrdict

cfg = compose_hydra_configuration(["config=pretrain/simclr/simclr_8node_resnet"])
args = convert_to_attrdict(cfg)
train_main(args)

solo-learn example:

from solo.methods import SimCLR
from solo.utils.classification_dataloader import prepare_data

model = SimCLR(backbone="resnet18")
train_loader, val_loader = prepare_data("cifar10", batch_size=256)
model.fit(train_loader, val_loader)

lightly

2,918

A python library for self-supervised learning on images.

Pros of Lightly

More comprehensive documentation and tutorials
Broader range of self-supervised learning methods implemented
Active development and community support

Cons of Lightly

Larger codebase, potentially more complex to understand and modify
Focuses on a wider range of tasks, which may be overwhelming for specific use cases

Code Comparison

Solo-learn:

class SimCLR(BaseMethod):
    def __init__(self, proj_output_dim: int, proj_hidden_dim: int, temperature: float):
        super().__init__()
        self.temperature = temperature
        self.projector = nn.Sequential(
            nn.Linear(self.features_dim, proj_hidden_dim),
            nn.ReLU(),
            nn.Linear(proj_hidden_dim, proj_output_dim),
        )

Lightly:

class SimCLR(nn.Module):
    def __init__(self, backbone, num_ftrs, out_dim):
        super().__init__()
        self.backbone = backbone
        self.projection_head = ProjectionHead(num_ftrs, out_dim)

    def forward(self, x):
        f = self.backbone(x)
        z = self.projection_head(f)
        return z

Both repositories implement SimCLR, but Lightly's implementation is more modular, separating the projection head into its own class. Solo-learn's implementation includes the temperature parameter directly in the SimCLR class, while Lightly handles it separately in the loss function.

lightning-bolts

1,678

Toolbox of models, callbacks, and datasets for AI/ML researchers.

Pros of lightning-bolts

More comprehensive collection of models, including GANs, VAEs, and self-supervised learning techniques
Integrated with PyTorch Lightning, providing a structured and scalable framework
Extensive documentation and examples for various use cases

Cons of lightning-bolts

Steeper learning curve due to PyTorch Lightning integration
Less focused on self-supervised learning compared to solo-learn
May have more dependencies and overhead for simpler projects

Code Comparison

solo-learn example:

from solo.methods import SimCLR

model = SimCLR(backbone="resnet18", num_classes=10, proj_hidden_dim=2048)

lightning-bolts example:

from pl_bolts.models.self_supervised import SimCLR

model = SimCLR(encoder="resnet18", num_samples=10, batch_size=32, dataset="cifar10")

Both repositories provide implementations of self-supervised learning methods, but solo-learn focuses specifically on this domain, while lightning-bolts offers a broader range of models and techniques. solo-learn may be more suitable for researchers focusing solely on self-supervised learning, while lightning-bolts provides a more comprehensive toolkit for various deep learning tasks within the PyTorch Lightning ecosystem.

pytorch-metric-learning

5,982

The easiest way to use deep metric learning in your application. Modular, flexible, and extensible. Written in PyTorch.

Pros of pytorch-metric-learning

More comprehensive set of metric learning techniques and loss functions
Extensive documentation and tutorials for easier implementation
Flexible integration with various deep learning frameworks

Cons of pytorch-metric-learning

Steeper learning curve due to its broader scope
May have more overhead for simple self-supervised learning tasks
Less focused on specific self-supervised learning algorithms

Code Comparison

solo-learn example:

from solo.methods import SimCLR

model = SimCLR(backbone, num_classes, proj_hidden_dim, proj_output_dim)

pytorch-metric-learning example:

from pytorch_metric_learning import losses, miners, distances

loss_func = losses.NTXentLoss(temperature=0.07)
mining_func = miners.MultiSimilarityMiner()

The solo-learn code focuses on implementing specific self-supervised learning methods, while pytorch-metric-learning provides more general-purpose metric learning tools that can be applied to various tasks, including self-supervised learning.

CMC

1,298

[arXiv 2019] "Contrastive Multiview Coding", also contains implementations for MoCo and InstDis

Pros of CMC

Focuses specifically on contrastive multiview coding, offering a more specialized approach
Provides implementations for multiple datasets, including ImageNet and CIFAR
Includes detailed documentation and explanations of the methodology

Cons of CMC

Less actively maintained, with fewer recent updates
More limited in terms of supported self-supervised learning methods
Smaller community and fewer contributors

Code Comparison

CMC:

def NT_xent(x, t=0.5):
    x = F.normalize(x, dim=1)
    x_dist = torch.pdist(x, p=2).pow(2).neg().div(2 * t).exp()
    return -torch.log(x_dist.sum() / (x.size(0) - 1))

solo-learn:

def nt_xent_loss(out_1, out_2, temperature=0.5):
    out = torch.cat([out_1, out_2], dim=0)
    n_samples = len(out)
    cov = torch.mm(out, out.t().contiguous())
    sim = torch.exp(cov / temperature)
    mask = ~torch.eye(n_samples, device=sim.device).bool()
    neg = sim.masked_select(mask).view(n_samples, -1).sum(dim=-1)
    pos = torch.exp(torch.sum(out_1 * out_2, dim=-1) / temperature)
    loss = -torch.log(pos / neg).mean()
    return loss

Both repositories implement contrastive loss functions, but solo-learn's implementation is more detailed and flexible, allowing for separate input tensors and explicit positive pair handling.

SimCLR

2,214

PyTorch implementation of SimCLR: A Simple Framework for Contrastive Learning of Visual Representations

Pros of SimCLR

Focused implementation of the SimCLR framework, making it easier to understand and modify for specific use cases
Includes detailed documentation and explanations of the SimCLR algorithm within the codebase
Lightweight implementation with minimal dependencies, suitable for educational purposes

Cons of SimCLR

Limited to SimCLR method, while solo-learn offers multiple self-supervised learning algorithms
Less actively maintained compared to solo-learn, with fewer recent updates and contributions
Lacks some advanced features and optimizations present in solo-learn, such as multi-GPU training support

Code Comparison

SimCLR:

class SimCLR(nn.Module):
    def __init__(self, base_encoder, projection_dim=128):
        super(SimCLR, self).__init__()
        self.encoder = base_encoder(num_classes=projection_dim)

solo-learn:

class SimCLR(BaseMethod):
    def __init__(
        self,
        proj_output_dim: int,
        proj_hidden_dim: int,
        temperature: float,
        **kwargs
    ):
        super().__init__(**kwargs)
        self.temperature = temperature
        self.projector = Projector(
            self.features_dim, proj_hidden_dim, proj_output_dim
        )

Convert designs to code with AI

Introducing Visual Copilot: A new AI model to turn Figma designs to high quality code using your components.

Try Visual Copilot

README

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

All pretrained models avaiable can be downloaded directly via the tables below or programmatically by running one of the following scripts zoo/cifar10.sh, zoo/cifar100.sh, zoo/imagenet100.sh and zoo/imagenet.sh.

Results

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

CIFAR-10

Method	Backbone	Epochs	Dali	Acc@1	Acc@5	Checkpoint
All4One	ResNet18	1000	:x:	93.24	99.88	:link:
Barlow Twins	ResNet18	1000	:x:	92.10	99.73	:link:
BYOL	ResNet18	1000	:x:	92.58	99.79	:link:
DeepCluster V2	ResNet18	1000	:x:	88.85	99.58	:link:
DINO	ResNet18	1000	:x:	89.52	99.71	:link:
MoCo V2+	ResNet18	1000	:x:	92.94	99.79	:link:
MoCo V3	ResNet18	1000	:x:	93.10	99.80	:link:
NNCLR	ResNet18	1000	:x:	91.88	99.78	:link:
ReSSL	ResNet18	1000	:x:	90.63	99.62	:link:
SimCLR	ResNet18	1000	:x:	90.74	99.75	:link:
Simsiam	ResNet18	1000	:x:	90.51	99.72	:link:
SupCon	ResNet18	1000	:x:	93.82	99.65	:link:
SwAV	ResNet18	1000	:x:	89.17	99.68	:link:
VIbCReg	ResNet18	1000	:x:	91.18	99.74	:link:
VICReg	ResNet18	1000	:x:	92.07	99.74	:link:
W-MSE	ResNet18	1000	:x:	88.67	99.68	:link:

CIFAR-100

Method	Backbone	Epochs	Dali	Acc@1	Acc@5	Checkpoint
All4One	ResNet18	1000	:x:	72.17	93.35	:link:
Barlow Twins	ResNet18	1000	:x:	70.90	91.91	:link:
BYOL	ResNet18	1000	:x:	70.46	91.96	:link:
DeepCluster V2	ResNet18	1000	:x:	63.61	88.09	:link:
DINO	ResNet18	1000	:x:	66.76	90.34	:link:
MoCo V2+	ResNet18	1000	:x:	69.89	91.65	:link:
MoCo V3	ResNet18	1000	:x:	68.83	90.57	:link:
NNCLR	ResNet18	1000	:x:	69.62	91.52	:link:
ReSSL	ResNet18	1000	:x:	65.92	89.73	:link:
SimCLR	ResNet18	1000	:x:	65.78	89.04	:link:
Simsiam	ResNet18	1000	:x:	66.04	89.62	:link:
SupCon	ResNet18	1000	:x:	70.38	89.57	:link:
SwAV	ResNet18	1000	:x:	64.88	88.78	:link:
VIbCReg	ResNet18	1000	:x:	67.37	90.07	:link:
VICReg	ResNet18	1000	:x:	68.54	90.83	:link:
W-MSE	ResNet18	1000	:x:	61.33	87.26	:link:

ImageNet-100

Method	Backbone	Epochs	Dali	Acc@1 (online)	Acc@1 (offline)	Acc@5 (online)	Acc@5 (offline)	Checkpoint
All4One	ResNet18	400	:heavy_check_mark:	81.93	-	96.23	-	:link:
Barlow Twins :rocket:	ResNet18	400	:heavy_check_mark:	80.38	80.16	95.28	95.14	:link:
BYOL :rocket:	ResNet18	400	:heavy_check_mark:	80.16	80.32	95.02	94.94	:link:
DeepCluster V2	ResNet18	400	:x:	75.36	75.4	93.22	93.10	:link:
DINO	ResNet18	400	:heavy_check_mark:	74.84	74.92	92.92	92.78	:link:
DINO :sleepy:	ViT Tiny	400	:x:	63.04	TODO	87.72	TODO	:link:
MoCo V2+ :rocket:	ResNet18	400	:heavy_check_mark:	78.20	79.28	95.50	95.18	:link:
MoCo V3 :rocket:	ResNet18	400	:heavy_check_mark:	80.36	80.36	95.18	94.96	:link:
MoCo V3 :rocket:	ResNet50	400	:heavy_check_mark:	85.48	84.58	96.82	96.70	:link:
NNCLR :rocket:	ResNet18	400	:heavy_check_mark:	79.80	80.16	95.28	95.30	:link:
ReSSL	ResNet18	400	:heavy_check_mark:	76.92	78.48	94.20	94.24	:link:
SimCLR :rocket:	ResNet18	400	:heavy_check_mark:	77.64	TODO	94.06	TODO	:link:
Simsiam	ResNet18	400	:heavy_check_mark:	74.54	78.72	93.16	94.78	:link:
SupCon	ResNet18	400	:heavy_check_mark:	84.40	TODO	95.72	TODO	:link:
SwAV	ResNet18	400	:heavy_check_mark:	74.04	74.28	92.70	92.84	:link:
VIbCReg	ResNet18	400	:heavy_check_mark:	79.86	79.38	94.98	94.60	:link:
VICReg :rocket:	ResNet18	400	:heavy_check_mark:	79.22	79.40	95.06	95.02	:link:
W-MSE	ResNet18	400	:heavy_check_mark:	67.60	69.06	90.94	91.22	:link:

: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

Method	Backbone	Epochs	Dali	Acc@1 (online)	Acc@1 (offline)	Acc@5 (online)	Acc@5 (offline)	Checkpoint	Finetuned Checkpoint
Barlow Twins	ResNet50	100	:heavy_check_mark:	67.18	67.23	87.69	87.98	:link:
BYOL	ResNet50	100	:heavy_check_mark:	68.63	68.37	88.80	88.66	:link:
MoCo V2+	ResNet50	100	:heavy_check_mark:	62.61	66.84	85.40	87.60	:link:
MAE	ViT-B/16	100	:x:	~	81.60 (finetuned)	~	95.50 (finetuned)	:link:	:link:

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.

Method	Dali	Total time for 20 epochs	Time for 1 epoch	GPU memory (per GPU)
Barlow Twins	:x:	1h 38m 27s	4m 55s	5097 MB
	:heavy_check_mark:	43m 2s	2m 10s (56% faster)	9292 MB
BYOL	:x:	1h 38m 46s	4m 56s	5409 MB
	:heavy_check_mark:	50m 33s	2m 31s (49% faster)	9521 MB
NNCLR	:x:	1h 38m 30s	4m 55s	5060 MB
	:heavy_check_mark:	42m 3s	2m 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}
}

Top Related Projects

Convert designs to code with AI

Introducing Visual Copilot: A new AI model to turn Figma designs to high quality code using your components.

Try Visual Copilot