msf-ssl's Introduction

MSF

Official code for "Mean Shift for Self-Supervised Learning" accepted as an oral presentation in ICCV 2021.

@InProceedings{Koohpayegani_2021_ICCV,
    author    = {Koohpayegani, Soroush Abbasi and Tejankar, Ajinkya and Pirsiavash, Hamed},
    title     = {Mean Shift for Self-Supervised Learning},
    booktitle = {Proceedings of the IEEE/CVF International Conference on Computer Vision (ICCV)},
    month     = {October},
    year      = {2021},
    pages     = {10326-10335}
}

@misc{koohpayegani2021mean,
      title={Mean Shift for Self-Supervised Learning}, 
      author={Soroush Abbasi Koohpayegani and Ajinkya Tejankar and Hamed Pirsiavash},
      year={2021},
      eprint={2105.07269},
      archivePrefix={arXiv},
      primaryClass={cs.CV}
}

Requirements

Python >= 3.7.6
PyTorch >= 1.4
torchvision >= 0.5.0
faiss-gpu >= 1.6.1

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

We train on 4 RTX6000 GPUs with 24GB of memory. But one can run our model with 4 RTX 2080Ti GPUs with 11GB of memory as well(with 128K memory bank). 200 Epochs of training with ResNet50 backbone will take approximately 140 hours to train.

Following command can be used to train the MSF

python train_msf.py \
  --cos \
  --weak_strong \
  --learning_rate 0.05 \
  --epochs 200 \
  --arch resnet50 \
  --topk 10 \
  --momentum 0.99 \
  --mem_bank_size 128000 \
  --checkpoint_path <CHECKPOINT PATH> \
  <DATASET PATH>