PyTorch implementation of the following papers:

A. v. d. Oord, Y. Li, and O. Vinyals, Representation learning with contrastive predictive coding

O. J. H ́enaff, A. Srinivas, J. D. Fauw, A. Razavi, C. Doersch, S. M. A. Eslami, and A. van den Oord Data-Efficient Image Recognition with Contrastive Predictive Coding

• PyTorch (verified with version 1.6.0)
• tqdm
• numpy
• opencv-python (patch aug visualistaion only - not training CPC)

Included is environment.yml

There are two training functions, one for the unsupervised training and one for supervised training.

• Viewing all command-line options

  python train_classifier.py -h
  

  python train_CPC.py -h
  

• Training a fully supervised model

  python train_classifier.py --fully_supervised --dataset stl10 --encoder resnet18
  

• Training Resnet14 on STL10 with CPCV1 - Unsupervised Stage

  python  train_CPC.py --dataset stl10 --epochs 300 --crop 64-0 --encoder resnet14 --norm none --grid_size 7 --pred_steps 5 --pred_directions 1
  

• Training Wideresnet-28-2 on CIFAR10 with CPCV2 (and a smaller grid size) - Unsupervised Stage

  python train_CPC.py --dataset  cifar10 --epochs 500 --crop 30-2 --encoder wideresnet-28-2 --norm layer --grid_size 5 --pred_steps 3 --pred_directions 4 --patch_aug 
  

• Training Wideresnet-28-2 on CIFAR10 with CPCV2 (and a smaller grid size) - Supervised Stage with 10,000 labeled images

  python train_classifier.py --dataset cifar10 --train_size 10000 --epochs 100 --lr 0.1 --crop 30-2 --encoder wideresnet-28-2 --norm layer --grid_size 5 --pred_directions 4 --cpc_patch_aug --patch_aug --model_num 500    
  

Clone this repository on Euler and place the required training datasets in the data folder.

🆕 You can now specify the size of unsupervised dataset for training CPC. Pass the --unsupervised_size 10000 argument in train_CPC.py to only use 10000 images. Use the --t1 {color, rotate, cutout, crop} and --t2 {color, rotate, cutout, crop} flags to specify two augmentations applied one after the other on a patchwise level.

• Submitting an interactive job on Euler with GPU

  srun --gpus=1 --gres=gpumem:12g --ntasks=2 --mem-per-cpu=12G --pty python train_CPC.py --dataset stl10 --unsupervised_size 10000 --epochs 100 --crop 64-0 --encoder resnet14 --norm none --grid_size 7 --pred_steps 5 --pred_directions 1
  

• Submitting a batch job on Euler with GPU

  sbatch --time=4:00:00 --gpus=1 --gres=gpumem:12g --ntasks=2 --mem-per-cpu=12G --wrap="python train_CPC.py --dataset stl10 --unsupervised_size 10000 --epochs 100 --crop 64-0 --encoder resnet14 --norm none --grid_size 7 --pred_steps 5 --pred_directions 1"