Package: simclr
SimCLRPyTorch implementation of SimCLR: A Simple Framework for Contrastive Learning of Visual Representations by T. Chen et al.
Including support for:Distributed data parallel trainingGlobal batch normalizationLARS (Layer-wise Adaptive Rate Scaling) optimizer.Link to paperOpen SimCLR in Google Colab Notebook (with TPU support)Open SimCLR results comparison on tensorboard.dev:Quickstart (fine-tune linear classifier)This downloads a pre-trained model and trains the linear classifier, which should receive an accuracy of ±82.9%on the STL-10 test set.git clone https://github.com/spijkervet/SimCLR.git && cd SimCLR
wget https://github.com/Spijkervet/SimCLR/releases/download/1.2/checkpoint_100.tar
sh setup.sh || python3 -m pip install -r requirements.txt || exit 1
conda activate simclr
python linear_evaluation.py --dataset=STL10 --model_path=. --epoch_num=100 --resnet resnet50CPUwget https://github.com/Spijkervet/SimCLR/releases/download/1.1/checkpoint_100.tar -O checkpoint_100.tar
python linear_evaluation.py --model_path=. --epoch_num=100 --resnet=resnet18 --logistic_batch_size=32simclrpackageSimCLR for PyTorch is now available as a Python package! Simply run and use it in your project:pip install simclrYou can then simply import SimCLR:from simclr import SimCLR

encoder = ResNet(...)
projection_dim = 64
n_features = encoder.fc.in_features  # get dimensions of last fully-connected layer
model = SimCLR(encoder, projection_dim, n_features)Training ResNet encoder:Simply run the following to pre-train a ResNet encoder using SimCLR on the CIFAR-10 dataset:python main.py --dataset CIFAR10Distributed TrainingWith distributed data parallel (DDP) training:CUDA_VISIBLE_DEVICES=0 python main.py --nodes 2 --nr 0
CUDA_VISIBLE_DEVICES=1 python main.py --nodes 2 --nr 1
CUDA_VISIBLE_DEVICES=2 python main.py --nodes 2 --nr 2
CUDA_VISIBLE_DEVICES=N python main.py --nodes 2 --nr 3ResultsThese are the top-1 accuracy of linear classifiers trained on the (frozen) representations learned by SimCLR:MethodBatch SizeResNetProjection output dimensionalityEpochsOptimizerSTL-10CIFAR-10SimCLR + Linear eval.256ResNet5064100Adam0.8290.833SimCLR + Linear eval.256ResNet5064100LARS0.783-SimCLR + Linear eval.256ResNet1864100Adam0.765-SimCLR + Linear eval.256ResNet186440Adam0.719-SimCLR + Linear eval.512ResNet186440Adam0.71-Logistic Regression---40Adam0.3580.389Pre-trained modelsResNet (batch_size, epochs)OptimizerSTL-10 Top-1ResNet50 (256, 100)Adam0.829ResNet18 (256, 100)Adam0.765ResNet18 (256, 40)Adam0.719python linear_evaluation.py --model_path=. --epoch_num=100LARS optimizerThe LARS optimizer is implemented inmodules/lars.py. It can be activated by adjusting theconfig/config.yamloptimizer setting to:optimizer: "LARS". It is still experimental and has not been thoroughly tested.What is SimCLR?SimCLR is a "simple framework for contrastive learning of visual representations". The contrastive prediction task is defined on pairs of augmented examples, resulting in 2N examples per minibatch. Two augmented versions of an image are considered as a correlated, "positive" pair (x_i and x_j). The remaining 2(N - 1) augmented examples are considered negative examples. The contrastive prediction task aims to identify x_j in the set of negative examples for a given x_i.UsageRun the following command to setup a conda environment:sh setup.sh
conda activate simclrOr alternatively with pip:pip install -r requirements.txtThen, simply run for single GPU or CPU training:python main.pyFor distributed training (DDP), use for every process in nodes, in which N is the GPU number you would like to dedicate the process to:CUDA_VISIBLE_DEVICES=0 python main.py --nodes 2 --nr 0
CUDA_VISIBLE_DEVICES=1 python main.py --nodes 2 --nr 1
CUDA_VISIBLE_DEVICES=2 python main.py --nodes 2 --nr 2
CUDA_VISIBLE_DEVICES=N python main.py --nodes 2 --nr 3--nrcorresponds to the process number of the N nodes we make available for training.TestingTo test a trained model, make sure to set themodel_pathvariable in theconfig/config.yamlto the log ID of the training (e.g.logs/0).
Set theepoch_numto the epoch number you want to load the checkpoints from (e.g.40).python linear_evaluation.pyor in place:python linear_evaluation.py --model_path=./save --epoch_num=40ConfigurationThe configuration of training can be found in:config/config.yaml. I personally prefer to use files instead of long strings of arguments when configuring a run. An exampleconfig.yamlfile:# train options
batch_size: 256
workers: 16
start_epoch: 0
epochs: 40
dataset_dir: "./datasets"

# model options
resnet: "resnet18"
normalize: True
projection_dim: 64

# loss options
temperature: 0.5

# reload options
model_path: "logs/0" # set to the directory containing `checkpoint_##.tar` 
epoch_num: 40 # set to checkpoint number

# logistic regression options
logistic_batch_size: 256
logistic_epochs: 100Logging and TensorBoardTo view results in TensorBoard, run:tensorboard --logdir runsOptimizers and learning rate scheduleThis implementation features the Adam optimizer and the LARS optimizer, with the option to decay the learning rate using a cosine decay schedule. The optimizer and weight decay can be configured in theconfig/config.yamlfile.Dependenciestorch
torchvision
tensorboard
pyyaml