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},
}

[comment]: <> (📋Optional: include a graphic explaining your approach/main result, bibtex entry, link to demos, blog posts and tutorials)

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)

[comment]: <> (📋Describe how to set up the environment, e.g. pip/conda/docker commands, download datasets, etc...)

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

[comment]: <> (📋Describe how to train the models, with example commands on how to train the models in your paper, including the full training procedure and appropriate hyperparameters.)

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:

License

This project is under the MIT license.