This post was originally published on zablo.net

Transfer learning and pre-training schemas for both NLP and Computer Vision have gained a lot of attention in the last months. Research showed that carefully designed unsupervised/self-supervised training could produce high-quality base models and embeddings that greatly decrease the amount of data needed to obtain good classification models downstream. This approach becomes more and more important as the companies collect a lot of data from which only a fraction can be labeled by humans—either due to the large cost of the labeling process or some time constraints.

Here I explore the SimCLR pre-training framework proposed by Google in this arxiv paper. I will explain the SimCLR and its contrastive loss function step by step, starting from a naive implementation followed by a faster, vectorized one. Then I will show how to use SimCLR’s pre-training routine to build image embeddings using EfficientNet network architecture, and finally, I will explain how to build a classifier on top of it.

TL;DR

This post covers:

  • understanding the SimCLR framework
  • from-scratch explanation & implementation of SimCLR’s loss function (NT-Xent) in PyTorch
  • pre-training image embeddings using EfficientNet architecture
  • training a classifier using transfer learning from the pre-trained embeddings

Understanding the SimCLR framework

In general, SimCLR is a simple framework for contrastive learning of visual representations. It’s not any new framework for deep learning, it’s a set of fixed steps that one should follow in order to train good-quality image embeddings.

I drew a schema that explains the flow and the whole representation learning process.

The flow is as follows (from left to right):

  1. Take an input image
  2. Prepare 2 random augmentations on the image, including rotations, hue/saturation/brightness changes, zooming, cropping, etc. The range of augmentations as well as analysis of which ones work best are discussed in detail in the paper.
  3. Run a deep neural network (preferably a convolutional one, like ResNet50) to obtain image representations (embeddings) of those augmented images.
  4. Run a small, fully connected linear neural network to project embeddings into another vector space.
  5. Calculate the contrastive loss and run backpropagation through both networks. Contrastive loss decreases when projections coming from the same image are similar. The similarity between projections can be arbitrary, here I will use cosine similarity, same as in the paper.

Contrastive loss function

The theory behind contrastive loss function

One can reason about contrastive loss function from two angles:

  1. Contrastive loss decreases when projections of augmented images coming from the same input image are similar.
  2. For two augmented images: (i), (j) (coming from the same input image—I will call them a “positive” pair later on), the contrastive loss for (i) tries to identify (j) among other images (“negative” examples) that are in the same batch.

The formal definition of the loss for a pair of positive examples (i) and (j) is as follows:

The final loss is an arithmetic mean of the losses for all positive pairs in the batch:

(keep in mind that the indexing in l(2k-1, 2k) + l(2k, 2k-1) is purely dependent on how you implement the loss—I find it easier to understand when I reason about them as l(i,j) + l(j,i)).

Contrastive loss function – implementation in PyTorch, ELI5 version

It’s much easier to implement the loss function without vectorization first, and then follow up with the vectorization phase.

import torch
from torch import nn
import torch.nn.functional as F
class ContrastiveLossELI5(nn.Module):
    def __init__(self, batch_size, temperature=0.5, verbose=True):
        super().__init__()
        self.batch_size = batch_size
        self.register_buffer("temperature", torch.tensor(temperature))
        self.verbose = verbose
            
    def forward(self, emb_i, emb_j):
        """
        emb_i and emb_j are batches of embeddings, where corresponding indices are pairs
        z_i, z_j as per SimCLR paper
        """
        z_i = F.normalize(emb_i, dim=1)
        z_j = F.normalize(emb_j, dim=1)

        representations = torch.cat([z_i, z_j], dim=0)
        similarity_matrix = F.cosine_similarity(representations.unsqueeze(1), representations.unsqueeze(0), dim=2)
        if self.verbose: print("Similarity matrix\n", similarity_matrix, "\n")
            
        def l_ij(i, j):
            z_i_, z_j_ = representations[i], representations[j]
            sim_i_j = similarity_matrix[i, j]
            if self.verbose: print(f"sim({i}, {j})={sim_i_j}")
                
            numerator = torch.exp(sim_i_j / self.temperature)
            one_for_not_i = torch.ones((2 * self.batch_size, )).to(emb_i.device).scatter_(0, torch.tensor([i]), 0.0)
            if self.verbose: print(f"1{{k!={i}}}",one_for_not_i)
            
            denominator = torch.sum(
                one_for_not_i * torch.exp(similarity_matrix[i, :] / self.temperature)
            )    
            if self.verbose: print("Denominator", denominator)
                
            loss_ij = -torch.log(numerator / denominator)
            if self.verbose: print(f"loss({i},{j})={loss_ij}\n")
                
            return loss_ij.squeeze(0)

        N = self.batch_size
        loss = 0.0
        for k in range(0, N):
            loss += l_ij(k, k + N) + l_ij(k + N, k)
        return 1.0 / (2*N) * loss

Explanation

Contrastive loss needs to know the batch size and the temperature (scaling) parameter. You can find details about setting the optimal temperature parameter in the paper.

My implementation of the forward of the contrastive loss takes two parameters. The first one will be a batch projection of images after the first augmentation, the second will be a batch projection of images after the second augmentation.

Projections need to be normalized first, hence:

z_i = F.normalize(emb_i, dim=1)
z_j = F.normalize(emb_j, dim=1)

All representations are concatenated together in order to efficiently calculate cosine similarities between each image pair.

representations = torch.cat([z_i, z_j], dim=0)

similarity_matrix = F.cosine_similarity(representations.unsqueeze(1), representations.unsqueeze(0), dim=2)

Next is the naive implementation of l(i,j) for clarity and easiness of understanding. The code below implements the equation almost directly:

def l_ij(i, j):
   z_i_, z_j_ = representations[i], representations[j]
   sim_i_j = similarity_matrix[i, j]
   numerator = torch.exp(sim_i_j / self.temperature)
   one_for_not_i = torch.ones((2 * self.batch_size, )).to(emb_i.device).scatter_(0, torch.tensor([i]), 0.0)
   denominator = torch.sum(
       one_for_not_i * torch.exp(similarity_matrix[i, :] / self.temperature)
   )

   loss_ij = -torch.log(numerator / denominator)
   return loss_ij.squeeze(0)

Then, the final loss for the batch is computed as an arithmetic mean of all combinations of positive examples:

N = self.batch_size
loss = 0.0
for k in range(0, N):
   loss += l_ij(k, k + N) + l_ij(k + N, k)
return 1.0 / (2*N) * loss

Now, let’s run it with verbose mode to see what’s inside.

I = torch.tensor([[1.0, 2.0], [3.0, -2.0], [1.0, 5.0]])
J = torch.tensor([[1.0, 0.75], [2.8, -1.75], [1.0, 4.7]])
loss_eli5 = ContrastiveLossELI5(batch_size=3, temperature=1.0, verbose=True)
loss_eli5(I, J)
Similarity matrix
   tensor([[ 1.0000, -0.1240,  0.9648,  0.8944, -0.0948,  0.9679],
       [-0.1240,  1.0000, -0.3807,  0.3328,  0.9996, -0.3694],
       [ 0.9648, -0.3807,  1.0000,  0.7452, -0.3534,  0.9999],
       [ 0.8944,  0.3328,  0.7452,  1.0000,  0.3604,  0.7533],
       [-0.0948,  0.9996, -0.3534,  0.3604,  1.0000, -0.3419],
       [ 0.9679, -0.3694,  0.9999,  0.7533, -0.3419,  1.0000]])

sim(0, 3)=0.8944272398948669
1{k!=0} tensor([0., 1., 1., 1., 1., 1.])
Denominator tensor(9.4954)
loss(0,3)=1.3563847541809082

sim(3, 0)=0.8944272398948669
1{k!=3} tensor([1., 1., 1., 0., 1., 1.])
Denominator tensor(9.5058)
loss(3,0)=1.357473373413086

sim(1, 4)=0.9995677471160889
1{k!=1} tensor([1., 0., 1., 1., 1., 1.])
Denominator tensor(6.3699)
loss(1,4)=0.8520082831382751

sim(4, 1)=0.9995677471160889
1{k!=4} tensor([1., 1., 1., 1., 0., 1.])
Denominator tensor(6.4733)
loss(4,1)=0.8681114912033081

sim(2, 5)=0.9999250769615173
1{k!=2} tensor([1., 1., 0., 1., 1., 1.])
Denominator tensor(8.8348)
loss(2,5)=1.1787779331207275

sim(5, 2)=0.9999250769615173
1{k!=5} tensor([1., 1., 1., 1., 1., 0.])
Denominator tensor(8.8762)
loss(5,2)=1.1834462881088257

tensor(1.1327)

A few things happened there, but by going back and forward between the verbose logs and the equation, everything should become clear.

The indexing jumps by batch size (first l(0,3), l(3,0)then l(1,4), l(4,1)because of the way the similarity matrix was constructed. First row of the similarity_matrixis:

[ 1.0000, -0.1240,  0.9648,  0.8944, -0.0948,  0.9679]

Remember the input:

I = torch.tensor([[1.0, 2.0], [3.0, -2.0], [1.0, 5.0]])
J = torch.tensor([[1.0, 0.75], [2.8, -1.75], [1.0, 4.7]])

Now:

1.0000is the cosine similarity between I[0] and I[0]([1.0, 2.0] and [1.0, 2.0])

-0.1240 is the cosine similarity between I[0] and I[1]([1.0, 2.0] and [3.0, -2.0])

-0.0948is the cosine similarity between I[0] and J[2]([1.0, 2.0] and [2.8, -1.75])

… and so on.

Let’s see if the loss decreases if the similarity between first image projection increases:

I = torch.tensor([[1.0, 2.0], [3.0, -2.0], [1.0, 5.0]])
J = torch.tensor([[1.0, 0.75], [2.8, -1.75], [1.0, 4.7]])
J = torch.tensor([[1.0, 1.75], [2.8, -1.75], [1.0, 4.7]]) # note the change
ContrastiveLossELI5(3, 1.0, verbose=False)(I, J)
tensor(1.0996)

Indeed the loss decreases! Now I will follow up with the vectorized implementation.

Contrastive loss function – implementation in PyTorch, vectorized version

The performance of naive implementation is really poor (mostly due to the manual loop), see the results:

contrastive_loss_eli5 = ContrastiveLossELI5(3, 1.0, verbose=False)
I = torch.tensor([[1.0, 2.0], [3.0, -2.0], [1.0, 5.0]], requires_grad=True)
J = torch.tensor([[1.0, 0.75], [2.8, -1.75], [1.0, 4.7]], requires_grad=True)
%%timeit
contrastive_loss_eli5(I, J)
838 µs ± 23.8 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)

Once I understood the internals of the loss, it’s easy to vectorize it and remove the manual loop:

class ContrastiveLoss(nn.Module):
   def __init__(self, batch_size, temperature=0.5):
       super().__init__()
       self.batch_size = batch_size
       self.register_buffer("temperature", torch.tensor(temperature))
       self.register_buffer("negatives_mask", (~torch.eye(batch_size * 2, batch_size * 2, dtype=bool)).float())
          

   def forward(self, emb_i, emb_j):
       """
       emb_i and emb_j are batches of embeddings, where corresponding indices are pairs
       z_i, z_j as per SimCLR paper
       """
       z_i = F.normalize(emb_i, dim=1)
       z_j = F.normalize(emb_j, dim=1)

       representations = torch.cat([z_i, z_j], dim=0)
       similarity_matrix = F.cosine_similarity(representations.unsqueeze(1), representations.unsqueeze(0), dim=2)
      
       sim_ij = torch.diag(similarity_matrix, self.batch_size)
       sim_ji = torch.diag(similarity_matrix, -self.batch_size)
       positives = torch.cat([sim_ij, sim_ji], dim=0)
      
       nominator = torch.exp(positives / self.temperature)
       denominator = self.negatives_mask * torch.exp(similarity_matrix / self.temperature)
  
       loss_partial = -torch.log(nominator / torch.sum(denominator, dim=1))
       loss = torch.sum(loss_partial) / (2 * self.batch_size)
       return loss
contrastive_loss = ContrastiveLoss(3, 1.0)
contrastive_loss(I, J).item() - contrastive_loss_eli5(I, J).item()
0.0

The difference should be zero or close to zero (< 1e-6 due to fp arithmetics). Performance comparison:

I = torch.tensor([[1.0, 2.0], [3.0, -2.0], [1.0, 5.0]], requires_grad=True)
J = torch.tensor([[1.0, 0.75], [2.8, -1.75], [1.0, 4.7]], requires_grad=True)
%%timeit
contrastive_loss_eli5(I, J)
918 µs ± 60.2 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)
%%timeit
contrastive_loss(I, J)
272 µs ± 9.18 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)

Almost 4x improvement, it works.

Pre-training image embeddings using SimCLR with EfficientNet

Once the loss function is established and understood, it’s time to make good use of it. I will pre-train image embeddings using EfficientNet architecture, following the SimCLR framework. For convenience, I’ve implemented a few utility functions and classes that I will explain briefly below. The training code is structured using PyTorch-Lightning.

I will use a great EfficientNet [https://arxiv.org/abs/1905.11946] implementation by Luke Melas-Kyriazi from GitHub, already pre-trained on ImageNet (transfer learning inception!). The dataset I choose is STL10 (from torchvision) as it contains both training and unlabeled splits for unsupervised / self-supervised learning tasks.

My goal here is to demonstrate the whole SimCLR flow from start to finish. I had no intent to reach new SOTA with the presented configuration.

Utility functions for image augmentations

Training with SimCLR produces good image embeddings that are not affected by image transformations. It happens because, during training, various data augmentations are done to force the network to understand the contents of the images regardless of i.e. the color of the image or the position where the object in the image is placed.

SimCLR’s authors say that composition of data augmentations plays a critical role in defining effective predictive tasks and also *contrastive learning needs stronger data augmentation than supervised learning*. To sum this up: when pre-training image embeddings, it’s good to make this task difficult for the network to learn by strongly augmenting the images in order to generalize better afterwards.

I strongly advise reading the SimCLR’s paper and its appendix as they did ablation studies on which data augmentations bring the best effects on the embeddings.

To keep this blog post simple to go through, I will mostly use built-in Torchvision’s data augmentations, with a single additional one – random resized rotation.

def random_rotate(image):
   if random.random() &gt; 0.5:
       return tvf.rotate(image, angle=random.choice((0, 90, 180, 270)))
   return image
class ResizedRotation():
   def __init__(self, angle, output_size=(96, 96)):
       self.angle = angle
       self.output_size = output_size
      
   def angle_to_rad(self, ang): return np.pi * ang / 180.0
      
   def __call__(self, image):
       w, h = image.size
       new_h = int(np.abs(w * np.sin(self.angle_to_rad(90 - self.angle))) + np.abs(h * np.sin(self.angle_to_rad(self.angle))))
       new_w = int(np.abs(h * np.sin(self.angle_to_rad(90 - self.angle))) + np.abs(w * np.sin(self.angle_to_rad(self.angle))))
       img = tvf.resize(image, (new_w, new_h))
       img = tvf.rotate(img, self.angle)
       img = tvf.center_crop(img, self.output_size)
       return img
class WrapWithRandomParams():
   def __init__(self, constructor, ranges):
       self.constructor = constructor
       self.ranges = ranges
  
   def __call__(self, image):
       randoms = [float(np.random.uniform(low, high)) for _, (low, high) in zip(range(len(self.ranges)), self.ranges)]
       return self.constructor(*randoms)(image)

from torchvision.datasets import STL10
import torchvision.transforms.functional as tvf
from torchvision import transforms
import numpy as np

A brief look on the transformation results:

stl10_unlabeled = STL10(".", split="unlabeled", download=True)
idx = 123
random_resized_rotation = WrapWithRandomParams(lambda angle: ResizedRotation(angle), [(0.0, 360.0)])
random_resized_rotation(tvf.resize(stl10_unlabeled[idx][0], (96, 96)))

Automatic data augmentation wrapper

Here I’ve also implemented a utility dataset wrapper that automatically applies random data augmentations whenever an image is retrieved. It can be easily used with any image dataset as long as it follows the simple interface of returning tuple with (PIL Image, anything). This wrapper can be set to return a deterministic transformation with the debug flag set to True. Note that there is a preprocess step that applies ImageNet-originated data standardization as I’m using already pre-trained EfficientNet.

 

from torch.utils.data import Dataset, DataLoader, SubsetRandomSampler, SequentialSampler
import random
class PretrainingDatasetWrapper(Dataset):
   def __init__(self, ds: Dataset, target_size=(96, 96), debug=False):
       super().__init__()
       self.ds = ds
       self.debug = debug
       self.target_size = target_size
       if debug:
           print("DATASET IN DEBUG MODE")
      

       # I will be using a network pre-trained on ImageNet first, which uses this normalization.
       # Remove this, if you're training from scratch or apply different transformations accordingly
       self.preprocess = transforms.Compose([
               transforms.ToTensor(),
               transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
       ])

      
       random_resized_rotation = WrapWithRandomParams(lambda angle: ResizedRotation(angle, target_size), [(0.0, 360.0)])
       self.randomize = transforms.Compose([
           transforms.RandomResizedCrop(target_size, scale=(1/3, 1.0), ratio=(0.3, 2.0)),
           transforms.RandomChoice([
               transforms.RandomHorizontalFlip(p=0.5),
               transforms.Lambda(random_rotate)
           ]),
           transforms.RandomApply([
               random_resized_rotation
           ], p=0.33),
           transforms.RandomApply([
               transforms.ColorJitter(brightness=0.5, contrast=0.5, saturation=0.5, hue=0.2)
           ], p=0.8),
           transforms.RandomGrayscale(p=0.2)
       ])

   def __len__(self): return len(self.ds)

   def __getitem_internal__(self, idx, preprocess=True):
       this_image_raw, _ = self.ds[idx]

       if self.debug:
           random.seed(idx)
           t1 = self.randomize(this_image_raw)
           random.seed(idx + 1)
           t2 = self.randomize(this_image_raw)
       else:
           t1 = self.randomize(this_image_raw)
           t2 = self.randomize(this_image_raw)
      
       if preprocess:
           t1 = self.preprocess(t1)
           t2 = self.preprocess(t2)
       else:
           t1 = transforms.ToTensor()(t1)
           t2 = transforms.ToTensor()(t2)

       return (t1, t2), torch.tensor(0)


   def __getitem__(self, idx):
       return self.__getitem_internal__(idx, True)

   def raw(self, idx):
       return self.__getitem_internal__(idx, False)
ds = PretrainingDatasetWrapper(stl10_unlabeled, debug=False)
tvf.to_pil_image(ds[idx][0][0])

tvf.to_pil_image(ds[idx][0][1])

tvf.to_pil_image(ds.raw(idx)[0][1])

tvf.to_pil_image(ds.raw(idx)[0][0])

SimCLR neural network for embeddings

Here I define the ImageEmbedding neural network which is based on EfficientNet-b0 architecture. I swap out the last layer of pre-trained EfficientNet with identity function and add projection for image embeddings on top of it (following the SimCLR paper) with Linear-ReLU-Linear layers. It was shown in the paper that the non-linear projection head (i.e Linear-ReLU-Linear) improves the quality of the embeddings.

from efficientnet_pytorch import EfficientNet
class ImageEmbedding(nn.Module):       
    class Identity(nn.Module):
        def __init__(self): super().__init__()

        def forward(self, x):
            return x


    def __init__(self, embedding_size=1024):
        super().__init__()
        
        base_model = EfficientNet.from_pretrained("efficientnet-b0")
        internal_embedding_size = base_model._fc.in_features
        base_model._fc = ImageEmbedding.Identity()
        
        self.embedding = base_model
        
        self.projection = nn.Sequential(
            nn.Linear(in_features=internal_embedding_size, out_features=embedding_size),
            nn.ReLU(),
            nn.Linear(in_features=embedding_size, out_features=embedding_size)
        )

    def calculate_embedding(self, image):
        return self.embedding(image)

    def forward(self, X):
        image = X
        embedding = self.calculate_embedding(image)
        projection = self.projection(embedding)
        return embedding, projection

Next is the implementation of a PyTorch-Lightning-based training module that orchestrates everything together:

  • hyper-parameters handling
  • SimCLR ImageEmbedding network
  • STL10 dataset
  • optimizer
  • forward step

As the PretrainingDatasetWrapper I’ve implemented returns a tuple of: (Image1, Image2), dummy class, the forward step for this module is straightforward—it needs to produce two batches of embeddings and calculate the contrastive loss function:

(X, Y), y = batch
embX, projectionX = self.forward(X)
embY, projectionY = self.forward(Y)
loss = self.loss(projectionX, projectionY)
from torch.multiprocessing import cpu_count
from torch.optim import RMSprop
import pytorch_lightning as pl
class ImageEmbeddingModule(pl.LightningModule):
    def __init__(self, hparams):
        hparams = Namespace(**hparams) if isinstance(hparams, dict) else hparams
        super().__init__()
        self.hparams = hparams
        self.model = ImageEmbedding()
        self.loss = ContrastiveLoss(hparams.batch_size)
    
    def total_steps(self):
        return len(self.train_dataloader()) // self.hparams.epochs
    
    def train_dataloader(self):
        return DataLoader(PretrainingDatasetWrapper(stl10_unlabeled, 
                                             debug=getattr(self.hparams, "debug", False)),
                          batch_size=self.hparams.batch_size, 
                          num_workers=cpu_count(),
                          sampler=SubsetRandomSampler(list(range(hparams.train_size))),
                         drop_last=True)
    
    def val_dataloader(self):
        return DataLoader(PretrainingDatasetWrapper(stl10_unlabeled,
                                            debug=getattr(self.hparams, "debug", False)),
                          batch_size=self.hparams.batch_size, 
                          shuffle=False,
                          num_workers=cpu_count(),
                          sampler=SequentialSampler(list(range(hparams.train_size + 1, hparams.train_size + hparams.validation_size))),
                         drop_last=True)
    
    def forward(self, X):
        return self.model(X)
    
    def step(self, batch, step_name = "train"):
        (X, Y), y = batch
        embX, projectionX = self.forward(X)
        embY, projectionY = self.forward(Y)
        loss = self.loss(projectionX, projectionY)
        loss_key = f"{step_name}_loss"
        tensorboard_logs = {loss_key: loss}

        return { ("loss" if step_name == "train" else loss_key): loss, 'log': tensorboard_logs,
                        "progress_bar": {loss_key: loss}}
    
    def training_step(self, batch, batch_idx):
        return self.step(batch, "train")
    
    def validation_step(self, batch, batch_idx):
        return self.step(batch, "val")
    
    def validation_end(self, outputs):
        if len(outputs) == 0:
            return {"val_loss": torch.tensor(0)}
        else:
            loss = torch.stack([x["val_loss"] for x in outputs]).mean()
            return {"val_loss": loss, "log": {"val_loss": loss}}

    def configure_optimizers(self):
        optimizer = RMSprop(self.model.parameters(), lr=self.hparams.lr)
        return [optimizer], []

Initial hyper-parameters. Batch size of 128 works fine with EfficientNet-B0 on GTX1070. Note that I’ve limited the training dataset to first 10k images from STL10 for convenience of running this blogpost in the form of a Jupyter Notebook / Google Colab.

 

Important! SimCLR greatly benefits from large batch sizes—it should be set to as high a value as possible given the GPU / cluster limits.

from argparse import Namespace
hparams = Namespace(
    lr=1e-3,
    epochs=50,
    batch_size=160,
    train_size=10000,
    validation_size=1000
)

Finding a good initial-learning rate using the LRFinder algorithm

I use pytorch-lightning’s built-in LRFinder algorithm to find an initial learning rate.

module = ImageEmbeddingModule(hparams)
t = pl.Trainer(gpus=1)
lr_finder = t.lr_find(module)
GPU available: True, used: True
TPU available: False, using: 0 TPU cores
CUDA_VISIBLE_DEVICES: [0]

   | Name  | Type            | Params
------------------------------------------
0 | model | ImageEmbedding  | 6 M  
1 | loss  | ContrastiveLoss | 0  
lr_finder.plot(show=False, suggest=True)
lr_finder.suggestion()
0.000630957344480193

I also use W&B Logging to keep track of my experiments:

from pytorch_lightning.loggers import WandbLogger
hparams = Namespace(
    lr=0.000630957344480193,
    epochs=10,
    batch_size=160,
    train_size=20000,
    validation_size=1000
)
module = ImageEmbeddingModule(hparams)
logger = WandbLogger(project="simclr-blogpost")
logger.watch(module, log="all", log_freq=50)
trainer = pl.Trainer(gpus=1, logger=logger)
trainer.fit(module)
| Name  | Type            | Params
------------------------------------------
0 | model | ImageEmbedding  | 6 M   
1 | loss  | ContrastiveLoss | 0 

 

After the training finishes, embeddings are ready for usage for the downstream tasks.

Building an image classifier on top of SimCLR embeddings

Once the embeddings are trained, they can be used to train the classifier on top of them—either by fine-tuning the whole network or by freezing the base network with embeddings and learning linear classifiers on top of it—I show the latter below.

Save the weights of a neural network with embeddings

I save the whole network in a form on a checkpoint. Only the internal part of the network will be used later with the classifier (projection layers will be discarded).

checkpoint_file = "efficientnet-b0-stl10-embeddings.ckpt"
trainer.save_checkpoint(checkpoint_file)
trainer.logger.experiment.log_artifact(checkpoint_file, type="model")

Classifier module

Again, I define a custom module—this time it uses already existing embeddings and freezes the base model’s weights on demand. Note that SimCLRClassifier.embeddings are only the EfficientNet part of the whole network used before—the projection head is discarded.

class SimCLRClassifier(nn.Module):
    def __init__(self, n_classes, freeze_base, embeddings_model_path, hidden_size=512):
        super().__init__()
        
        base_model = ImageEmbeddingModule.load_from_checkpoint(embeddings_model_path).model
        
        self.embeddings = base_model.embedding
        
        if freeze_base:
            print("Freezing embeddings")
            for param in self.embeddings.parameters():
                param.requires_grad = False
                
        # Only linear projection on top of the embeddings should be enough
        self.classifier = nn.Linear(in_features=base_model.projection[0].in_features, 
                      out_features=n_classes if n_classes &gt; 2 else 1)

    
    def forward(self, X, *args):
        emb = self.embeddings(X)
        return self.classifier(emb)

Classifier training code

The classifier training code again uses PyTorch lightning, so I’ll skip the in-depth explanation.

from torch import nn
from torch.optim.lr_scheduler import CosineAnnealingLR
class SimCLRClassifierModule(pl.LightningModule):
    def __init__(self, hparams):
        super().__init__()
        hparams = Namespace(**hparams) if isinstance(hparams, dict) else hparams
        self.hparams = hparams
        self.model = SimCLRClassifier(hparams.n_classes, hparams.freeze_base, 
                                      hparams.embeddings_path,
                                      self.hparams.hidden_size)
        self.loss = nn.CrossEntropyLoss()
    
    def total_steps(self):
        return len(self.train_dataloader()) // self.hparams.epochs
    
    def preprocessing(seff):
        return transforms.Compose([
                transforms.ToTensor(),
                transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
        ])
    
    def get_dataloader(self, split):
        return DataLoader(STL10(".", split=split, transform=self.preprocessing()),
                          batch_size=self.hparams.batch_size, 
                          shuffle=split=="train",
                          num_workers=cpu_count(),
                         drop_last=False)
    
    def train_dataloader(self):
        return self.get_dataloader("train")
    
    def val_dataloader(self):
        return self.get_dataloader("test")
    
    def forward(self, X):
        return self.model(X)
    
    def step(self, batch, step_name = "train"):
        X, y = batch
        y_out = self.forward(X)
        loss = self.loss(y_out, y)
        loss_key = f"{step_name}_loss"
        tensorboard_logs = {loss_key: loss}

        return { ("loss" if step_name == "train" else loss_key): loss, 'log': tensorboard_logs,
                        "progress_bar": {loss_key: loss}}
    
    def training_step(self, batch, batch_idx):
        return self.step(batch, "train")
    
    def validation_step(self, batch, batch_idx):
        return self.step(batch, "val")
    
    def test_step(self, batch, batch_idx):
        return self.step(Batch, "test")
    
    def validation_end(self, outputs):
        if len(outputs) == 0:
            return {"val_loss": torch.tensor(0)}
        else:
            loss = torch.stack([x["val_loss"] for x in outputs]).mean()
            return {"val_loss": loss, "log": {"val_loss": loss}}

    def configure_optimizers(self):
        optimizer = RMSprop(self.model.parameters(), lr=self.hparams.lr)
        schedulers = [
            CosineAnnealingLR(optimizer, self.hparams.epochs)
        ] if self.hparams.epochs &gt; 1 else []
        return [optimizer], schedulers

It’s worth mentioning here that training with a frozen base model gives a great performance boost during training as the gradients need to be calculated only for a single layer. Additionally, by utilizing good embeddings, only a few epochs are required to reach a good quality classifier with a single linear projection.

hparams_cls = Namespace(
    lr=1e-3,
    epochs=5,
    batch_size=160,
    n_classes=10,
    freeze_base=True,
    embeddings_path="./efficientnet-b0-stl10-embeddings.ckpt",
    hidden_size=512
)
module = SimCLRClassifierModule(hparams_cls)
logger = WandbLogger(project="simclr-blogpost-classifier")
logger.watch(module, log="all", log_freq=10)
trainer = pl.Trainer(gpus=1, max_epochs=hparams_cls.epochs, logger=logger)
lr_find_cls = trainer.lr_find(module)
| Name  | Type             | Params
-------------------------------------------
0 | model | SimCLRClassifier | 4 M   
1 | loss  | CrossEntropyLoss | 0  
LR finder stopped early due to diverging loss.   
lr_find_cls.plot(show=False, suggest=True)
lr_find_cls.suggestion()
0.003981071705534969

hparams_cls = Namespace(
    lr=0.003981071705534969,
    epochs=5,
    batch_size=160,
    n_classes=10,
    freeze_base=True,
    embeddings_path="./efficientnet-b0-stl10-embeddings.ckpt",
    hidden_size=512
)
module = SimCLRClassifierModule(hparams_cls)
trainer.fit(module)
| Name  | Type             | Params
-------------------------------------------
0 | model | SimCLRClassifier | 4 M   
1 | loss  | CrossEntropyLoss | 0    

Evaluation

Here I define a utility function for evaluating the model using the provided data loader. Note that the transfer between GPU and CPU and storing all results in-memory might be not effective for large datasets!

from sklearn.metrics import classification_report
def evaluate(data_loader, module):
    with torch.no_grad():
        progress = ["/", "-", "\\", "|", "/", "-", "\\", "|"]
        module.eval().cuda()
        true_y, pred_y = [], []
        for i, batch_ in enumerate(data_loader):
            X, y = batch_
            print(progress[i % len(progress)], end="\r")
            y_pred = torch.argmax(module(X.cuda()), dim=1)
            true_y.extend(y.cpu())
            pred_y.extend(y_pred.cpu())
        print(classification_report(true_y, pred_y, digits=3))
        return true_y, pred_y
_ = evaluate(module.val_dataloader(), module)
   precision    recall  f1-score   support
    
               0      0.856     0.864     0.860       800
               1      0.714     0.701     0.707       800
               2      0.903     0.919     0.911       800
               3      0.678     0.599     0.636       800
               4      0.665     0.746     0.703       800
               5      0.633     0.564     0.597       800
               6      0.729     0.781     0.754       800
               7      0.678     0.709     0.693       800
               8      0.868     0.910     0.888       800
               9      0.862     0.801     0.830       800
    
        accuracy                          0.759      8000
       macro avg      0.759     0.759     0.758      8000
    weighted avg      0.759     0.759     0.758      8000

Summary

I hope you find my explanation of the SimCLR framework useful. Feel free to ask about anything, comment, play with the notebook, and share this post with the ML community!

Additional links and resources

Comments are closed.