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.