PySSL

A PyTorch implementation of Self-Supervised Learning (SSL) methods

Updates

• 19/06/2023 - Post on medium

1. Prerequisites

Before proceeding, create a conda environment:

conda create -n pyssl

Activate the newly created environment:

conda activate pyssl

Once the environment is activated, install the required packages from the "requirements.txt" file using the following command:

pip install -r requirements.txt

2. Methods

The SSL method implementation was based on lucidrains implementation of BYOL.

For the following section, assume a backbone, e.g., a ResNet-50, and an input image of size 224 in a batch of 4, i.e.,

import torchvision

backbone = torchvision.models.resnet50(pretrained=False)
feature_size = backbone.fc.in_features
backbone.fc = torch.nn.Identity()

x = torch.rand(4, 3, 224, 224)

2.1 Barlow Twins

Barlow Twins is a SSL method that aims to learn representations by making the outputs of a neural network to be as similar as possible for two augmented views of the same image, while reducing the redundancy between the output components.

model = BarlowTwins(backbone, feature_size, projection_dim=8192, hidden_dim=8192, lamda=0.005)

2.2 BYOL

BYOL (Bootstrap Your Own Latent) is an SSL method that learns representations by predicting one view of an input from another view of the same input, without the necessity for negative pairs which is the case for most contrastive learning methods.

model = BYOL(backbone, feature_size, projection_dim=256, hidden_dim=4096, tau=0.996)

2.3 DINO

DINO (Distillation of Self-supervised Networks) is a method that trains a student network using the outputs of a teacher network, where the teacher network is an exponentially moving average of the student network. The key idea is to use distillation loss to capture information from different viewpoints of the same image.

model = DINO(backbone, feature_size, projection_dim=256, hidden_dim=2048, bottleneck_dim=256, temp_s=0.1, temp_t=0.5, m=0.5, lamda=0.996, num_crops=6)

2.4 MoCo, MoCov2, MoCov3

MoCo (Momentum Contrast) series are methods that utilize a dynamic dictionary implemented with a queue and a moving-averaged encoder. The methods aim to maximize similarity between a query and its positive key and minimize similarity between the query and negative keys (v2 and v3 are updated versions with improvements over the original).

model = MoCo(backbone, feature_size, projection_dim=128, K=65536, m=0.999, temperature=0.07)
model = MoCoV2(backbone, feature_size, projection_dim=128, K=65536, m=0.999, temperature=0.07)
model = MoCoV3(backbone, feature_size, projection_dim=256, hidden_dim=2048, temperature=0.5, m=0.999)

2.5 SimCLR, SimCLRv2

SimCLR (Simple Framework for Contrastive Learning of Visual Representations) is a method that utilizes a contrastive loss function to distinguish between similar and dissimilar images. The v2 version introduces a simple method for augmenting the data sample with a learnable nonlinear transformation.

model = SimCLR(backbone, feature_size, projection_dim=128, temperature=0.5)
model = SimCLRv2(backbone, feature_size, projection_dim=128, temperature=0.5)

2.6 SimSiam

SimSiam (Simple Siamese) aims to learn representations by predicting one view of an input from another view of the same input, similar to BYOL. However, unlike BYOL, SimSiam does not use batch normalization in the prediction MLP, and does not require the use of a momentum encoder or moving average.

model = SimSiam(backbone, feature_size, projection_dim=2048, hidden_dim_proj=2048, hidden_dim_pred=512)

2.7 SupCon

SupCon (Supervised Contrastive Learning) is a method that uses a contrastive loss function but also makes use of labels, if available, to further refine the learning of representations by encouraging dissimilarities between representations of different classes.

model = SupCon(backbone, feature_size, projection_dim=128, temperature=0.07)

2.8 SwAV

SwAV (Swapping Assignments between multiple Views of the same image) is a clustering-based method which learns representations by swapping the cluster assignments between different views of the same image, and then minimizing the difference between the swapped assignments and the original assignments.

model = SwAV(backbone, feature_size, projection_dim=128, hidden_dim=2048, temperature=0.1, epsilon=0.05, sinkhorn_iterations=3, num_prototypes=3000, queue_length=64, use_the_queue=True, num_crops=6)

3. Training

The models can directly output the loss, i.e., loss = model(x) or loss = model.forward(x) so as to integrate smoothly with the training loop (see main.py).

import torch
import torchvision

# get device
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

# initialize backbone (resnet50)
backbone = torchvision.models.resnet50(pretrained=False)
feature_size = backbone.fc.in_features
backbone.fc = torch.nn.Identity()

# initialize ssl method
model = builders.SimCLR(backbone, feature_size, image_size=32)
model = model.to(device)
    
# load fake CIFAR-like dataset
dataset = torchvision.datasets.FakeData(2000, (3, 32, 32), 10, torchvision.transforms.ToTensor())
loader = torch.utils.data.DataLoader(dataset, batch_size=128, shuffle=True, num_workers=2)

# set optimizer
optimizer = torch.optim.Adam(model.parameters(), lr=0.001)

# switch to train mode
model.train()

# epoch training
for epoch in range(10):
    for i, (x, _) in enumerate(loader):
        x = x.to(device)

        # zero the parameter gradients
        model.zero_grad()

        # compute loss
        loss = model(x)
        
        # compute gradient and do SGD step
        loss.backward()
        optimizer.step()

Note, that each method uses a default augmentation for training, according to original implementation, and the model requires the image size, mean, and standard deviation of the dataset.

For example:

kwargs = {
    'image_size': 224,
    'mean': [0.485, 0.456, 0.406],
    'std': [0.229, 0.224, 0.225]
}

model = BYOL(backbone, feature_size, projection_dim=8192, hidden_dim=8192, lamda=0.005, **kwargs)

4. Inference

For inference, use either model.encoder(x) to get the projection vector (backbone and projector's representations), or use model.backbone(x) to get the feature vector (backbone representations). The model automatically assigns the encoder to the corresponding encoder of each method (e.g., student network for BYOL or queue encoder for MoCo series).

5. Citation

In Bibtex format:

@misc{pyssl2023giakoumoglou,
   author = {Nikolaos Giakoumoglou and Paschalis Giakoumoglou},
   title = {PySSL: A PyTorch implementation of Self-Supervised Learning (SSL) methods},
   year = {2023},
   publisher = {GitHub},
   journal = {GitHub repository},
   howpublished = {\url{https://github.com/giakou4/pyssl}},
}

5. Support

Reach out to me:

• giakou4's email