CompRess: Self-Supervised Learning by Compressing Representations

This repository is the official implementation of CompRess: Self-Supervised Learning by Compressing Representations

Project webpage. https://umbcvision.github.io/CompRess/

@Article{abbasi2020compress,
  author  = {Koohpayegani, Soroush Abbasi and Tejankar, Ajinkya and Pirsiavash, Hamed},
  title   = {CompRess: Self-Supervised Learning by Compressing Representations},
  journal = {Advances in neural information processing systems},
  year    = {2020},
}

Requirements

Install PyTorch and ImageNet dataset following the official PyTorch ImageNet training code. We used Python 3.7 for our experiments.

• Install PyTorch (pytorch.org)

To run NN and Cluster Alignment, you require to install FAISS.

FAISS:

• Install FAISS (https://github.com/facebookresearch/faiss/blob/master/INSTALL.md)

Training

Our code is based on unofficial implementation of MoCo from https://github.com/HobbitLong/CMC.

To train the student(s) using pretrained teachers in the paper :

Download pretrained official MoCo ResNet50 model from https://github.com/facebookresearch/moco.

Then train the student using pretrained model:

python train_student.py \
    --teacher_arch resnet50 \ 
    --teacher <path_to_pretrained_model or cached_features> \
    --student_arch mobilenet \
    --checkpoint_path <path_to_checkpoint_folder> \
    <path_to_imagenet_data>

To train the student(s) using cached teachers in the paper :

We converted TensorFlow SimCLRv1 ResNet50x4(https://github.com/google-research/simclr) to PyTorch. Optionally, you can download pretrained SimCLR ResNet50x4 PyTorch model from here.

First, run this command to calculate and store cached features.

python cache_feats.py \ 
    --weight <path_to_pretrained_model> \
    --save <path_to_save_folder> \
    --arch resnet50x4 \ 
    --data_pre_processing SimCLR \ 
    <path_to_imagenet_data>

Then train the student using cached features:

python train_student.py \
    --cache_teacher \ 
    --teacher <path_to_pretrained_model or cached_features> \
    --student_arch mobilenet \
    --checkpoint_path <path_to_checkpoint_folder> \
    <path_to_imagenet_data>

To train the student(s) without Momentum framework execute train_student_without_momentum.py instead of train_student.py

Evaluation

To run Nearest Neighbor evaluation on ImageNet, run:

python eval_knn.py \
    --arch alexnet \
    --weights <path_to_pretrained_model> \
    --save <path_to_save_folder> \
    <path_to_imagenet_data>

Note that above execution will cache features too. After first execution, you can add "--load_cache" flag to load cached features from a file.

To run Cluster Alignment evaluation on ImageNet, run:

python eval_cluster_alignment.py  \
    --weights <path_to_pretrained_model> \
    --arch resnet18  \
    --save <path_to_save_folder> \ 
    --visualization \ 
    --confusion_matrix \ 
    <path_to_imagenet_data> 

To run Linear Classifier evaluation on ImageNet, run:

python eval_linear.py \
    --arch alexnet \
    --weights <path_to_pretrained_model> \
    --save <path_to_save_folder> \
    <path_to_imagenet_data>

Results

"SOTA Self-Supervised" refers to SimCLR for RexNet50x4 and MoCo for all other architectures.

Our model achieves the following performance on ImageNet:

Model name	Teacher	Top-1 Linear Classifier Accuracy	Top-1 Nearest Neighbor Accuracy	Top-1 Cluster Alignment Accuracy	Pre-trained
CompRess(Resnet50)	SimCLR ResNet50x4(cached)	71.6%	63.4%	42.0%	Pre-trained Resnet50
CompRess(Mobilenet)	MoCoV2 ResNet50	63.0%	54.4%	35.5%	Pre-trained Mobilenet
CompRess(Resnet18)	MoCoV2 ResNet50	61.7%	53.4%	34.7%	Pre-trained Resnet18
CompRess(Resnet18)	SwAV ResNet50	65.6%	56.0%	26.3%	Pre-trained Resnet18
CompRess(Alexnet)	SimCLR ResNet50x4(cached)	57.6%	52.3%	33.3%	Pre-trained Alexnet

License

This project is under the MIT license.