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},
}
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)
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
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>
"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 |
This project is under the MIT license.