In this project, we present an implementation of SimCLR in PyTorch. We also form probabilistic notion of the contrastive learning framework and derive a new loss function. The goal is to truly understand how a contrastive learning model (SimCLR) learns, how to interpret learned representations, and to quantify and interpret uncertainty.
To pretrain the model with gradient accumulation with batch size = n_accum * 64
, for a number of epochs =
num_of_epochs
, dataset = "cifar10"/"stl10"
, path for saving the model and checkpoints = "/path/for/saving/"
, use_new_loss, run
python3 pretrain.py --n_epoch=num_of_epochs --accum_steps=n_accum --dataset=dataset --path_for_saving="/path/for/saving/" --new_loss=use_new_loss
The CIFAR-10 dataset is used to build a prototype.
Below is a sample of augmented CIFAR-10 data:
Augmentation performed for this project followed exactly the same procedure as what was carried out in the paper.
For clarity, we list the steps here:
- Random cropping (inception-style: random crop size uniform from 0.08 to 1 in area and a random aspect ratio) and
resizing to original size with random flipping (p=50%);
torchvision.transforms.RandomResizedCrop
- Random color distortions (color jittering + color dropping)
transforms.ColorJitter
,transforms.RandomApply
,transforms.RandomGrayscale
- Random Gaussian blur (p=50%). Randomly sample volatility in [0.1, 2.0], and the kernel size is 10% of the image height/width.
Following the paper, we used a slightly modified ResNet50
as the encoder for CIFAR images.
We modify the original resnet module in pytorch by defining a wrapper on top of it in order to:
- replace the first 7x7 convolution layer of stride 2 (
conv1
) with a 3x3 convolution of stride 1, to adjust for the smaller resolution of images in CIFAR10. - remove the first maxpooling operation (
maxpool
) - remove the last fully-connected layer and take the output of the average pooling layer
The basic model for pretraining consists of:
- Encoder
f
(ResnetEncoder
) - Projection head
g
(2-layer (could adjust this in further experiments) MLP with a RELU activation and batch normalization)
- Currently using a
batch_size
of 512 and gradient accumulation to allow (relatively) larger batch training on a single Nvidia Tesla GPU with 12GB RAM. - Different from the paper, Adam
was used with a learning rate of 1e-3 and a
weight_decay
of 1e-6 in pretraining. Lars optimizer could be added later.
After pretraining the model, features are extracted with the model and used as input to a linear classifier trained
with CrossEntropyLoss
. This linear model is trained with the L-BFGS
optimizer as suggested in the paper.
During fine-tuning, we copy the pretrained model weights to a new model, remove the projection heads and attach a linear
classifier. This new model is trained with a fraction (e.g., 10%) of labelled data using SGD
with Nesterov momentum.
@article{chen2020simple,
title={A Simple Framework for Contrastive Learning of Visual Representations},
author={Chen, Ting and Kornblith, Simon and Norouzi, Mohammad and Hinton, Geoffrey},
journal={arXiv preprint arXiv:2002.05709},
year={2020}
}