InvCLR: Improving Transformation Invariance in Contrastive Representation Learning

This is a PyTorch implementation of the ICLR paper Improving Transformation Invariance in Contrastive Representation Learning (InvCLR):

@article{foster2020improving,
  title={Improving Transformation Invariance in Contrastive Representation Learning},
  author={Foster, Adam and Pukdee, Rattana and Rainforth, Tom},
  journal={arXiv preprint arXiv:2010.09515},
  year={2020}
}

Installation and datasets

Install PyTorch following the instructions here. Download the the CIFAR-10 and CIFAR-100 datasets. The Spirograph dataset is included in this code base. To use the Spirograph dataset on its own, see this standalone repo. To install the requisite packages for this project, use pip install -r requirements.txt. Note: to install torchlars it is necessary to set the environment variable CUDA_HOME.

Dataset locations

The dataset file locations should be specified in a JSON file of the following form

dataset-paths.json
{
    "cifar10": "/data/cifar10/",
    "cifar100": "/data/cifar100/",
    "stl10": "/data/stl10/",
    "imagenet": "/data/imagenet/2012/"
}

Training an encoder

We support multi-GPU DataParallel training. Use the following command to train an encoder from scratch on CIFAR-10. We ran this using 8 GPUs.

$ python3 invclr.py \
  --num-epochs 1000 \
  --cosine-anneal \
  --arch resnet50 \
  --dataset cifar10 \
  --lambda-gp 1e-1 \
  --filename cifar10_run

Set --dataset cifar100 to train on CIFAR-100. To train an encoder on the Spirograph dataset, use

$ python3 invclr.py \
  --num-epochs 50 \
  --cosine-anneal \
  --test-freq 0 \
  --save-freq 10 \
  --arch resnet18 \
  --dataset spirograph \
  --lambda-gp 1e-2 \
  --filename spirograph_run \
  --gp-upper-limit 1000

You can set --lambda-gp 0 to train an encoder with no gradient penalty.

Evaluating an encoder

Use the following command to evaluate the trained CIFAR-10 encoder on untransformed inputs with 50% of the training labels used for supervized training

$ python3 eval.py \
  --load-from cifar10_run_epoch999.pth \
  --untransformed \
  --proportion 0.5

for Spirograph, we used

$ python3 eval.py \
  --load-from spirograph_run_epoch049.pth \
  --reg-weight 1e-8 \
  --proportion 0.5

On Spirograph, you can view regression results for each of the separate tasks by adding --componentwise, otherwise the presented loss is the mean over all tasks.

Feature averaging

Use the following command to evaluate classification performance of feature averaging using an average of 100 samples

$ python3 eval.py \
  --load-from cifar10_run_epoch999.pth \
  --num-passes 100 

for Spirograph, run the following code

$ python3 eval.py \
  --load-from spirograph_run_epoch049.pth \
  --num-passes 30 \
  --reg-weight 1e-8

We obtained the following (best of 3 runs)

Dataset	Loss	Accuracy
CIFAR-10	0.154	94.9%
CIFAR-100	1.008	75.1%
Spirograph	0.000983	n/a