PCL

Prototypical Contrastive Learning of Unsupervised Representations (Salesforce Research)

This is a PyTorch implementation of the PCL paper:

@inproceedings{PCL, title={Prototypical Contrastive Learning of Unsupervised Representations}, author={Junnan Li and Pan Zhou and Caiming Xiong and Steven C.H. Hoi}, booktitle={ICLR}, year={2021} }

Requirements:

• ImageNet dataset
• Python ≥ 3.6
• PyTorch ≥ 1.4
• faiss-gpu: pip install faiss-gpu
• pip install tqdm

Unsupervised Training:

This implementation only supports multi-gpu, DistributedDataParallel training, which is faster and simpler; single-gpu or DataParallel training is not supported.

To perform unsupervised training of a ResNet-50 model on ImageNet using a 4-gpu or 8-gpu machine, run: python main_pcl.py \ -a resnet50 \ --lr 0.03
--batch-size 256
--temperature 0.2
--mlp --aug-plus --cos (only activated for PCL v2) \ --dist-url 'tcp://localhost:10001' --multiprocessing-distributed --world-size 1 --rank 0
--exp-dir experiment_pcl [Imagenet dataset folder]

Download Pre-trained Models

Linear SVM Evaluation on VOC

To train a linear SVM classifier on VOC dataset, using frozen representations from a pre-trained model, run: python eval_svm_voc.py --pretrained [your pretrained model]
-a resnet50 \ --low-shot (only for low-shot evaluation, otherwise the entire dataset is used)
[VOC2007 dataset folder]

Linear SVM classification result on VOC, using ResNet-50 pretrained with PCL for 200 epochs:

k is the number of training samples per class.

Linear Classifier Evaluation on ImageNet

Requirement: pip install tensorboard_logger
To train a logistic regression classifier on ImageNet, using frozen representations from a pre-trained model, run: python eval_cls_imagenet.py --pretrained [your pretrained model]
-a resnet50 \ --lr 5
--batch-size 256
--id ImageNet_linear \ --dist-url 'tcp://localhost:10001' --multiprocessing-distributed --world-size 1 --rank 0
[Imagenet dataset folder]

Linear classification result on ImageNet, using ResNet-50 pretrained with PCL for 200 epochs: