Self-Supervised Learning for Large-Scale Unsupervised Image Clustering

This is code to run experiments for paper "Self-Supervised Learning for Large-Scale Unsupervised Image Clustering".

Running the code

For part of the models, you'll need to download the chekpoints manually:

• SimCLR models: https://github.com/google-research/simclr
• MoCo and InfoMin models: https://github.com/HobbitLong/PyContrast/blob/master/pycontrast/docs/MODEL_ZOO.md
• SwAV models: https://github.com/facebookresearch/swav

and put them in chekpoint folder.

For SimCLRv2, BigBiGAN as well as supervised models checkpoints are downloaded automatically.

Download the code and install dependencies

Remember to clone the submodules by running

git clone --recurse-submodules https://github.com/Randl/kmeans_selfsuper.git

during cloning the repo, or, if you forgot to do it, by running

git submodule update --init --recursive

in the repo folder.

You'll need to install dependencies, by running

pip install -r requirements.txt

Generating features

For SimCLRv2 and BigBiGAN run

python3 generate_prediction_tf.py --model resnet152_simclr2
python3 generate_prediction_tf.py --model resnet50_simclr2
python3 generate_prediction_tf.py --model resnet152x3_simclr2
python3 generate_prediction_tf.py --model resnet50_bigbigan
python3 generate_prediction_tf.py --model revnet50x4_bigbigan

For InfoMin, MoCo v2 and SwAV, run

python3 generate_prediction_pytorch.py --model resnext152_infomin
python3 generate_prediction_pytorch.py --model resnet50_infomin
python3 generate_prediction_pytorch.py --model resnet50_mocov2
python3 generate_prediction_pytorch.py --model resnet50_swav

Finally, for supervised models, run

python3 generate_prediction_pytorch_supervised.py --model tf_efficientnet_l2_ns_475
python3 generate_prediction_pytorch_supervised.py --model gluon_resnet152_v1s
python3 generate_prediction_pytorch_supervised.py --model ig_resnext101_32x48d

You'll need large amount of RAM since the script keeps features in memory. It was tested on machine with 128 GB RAM.

Running clustering

To run clustering, you need to run

python3 cluster.py --model resnet50_infomin

where the model name should fit the name in generating part. For overclustering, e.g., 1.25 times more clusters than classes, run

python3 cluster.py --model resnet152_simclr2 --over 1.25

For using smaller dimensions of features, e.g., 512, run

python3 cluster.py --model resnet152_simclr2 --n-components 512

Citing the paper

If you found the paper or the code useful, please cite it. You can use following bibtex entry:

@article{zheltonozhskii2020unsupervised,
  title = {Self-Supervised Learning for Large-Scale Unsupervised Image Clustering},
  author = {Zheltonozhskii, Evgenii and Baskin, Chaim and Bronstein, Alex M. and Mendelson, Avi},
  journal = {NeurIPS Self-Supervised Learning Workshop},
  year = {2020},
  month = aug,
  url = {https://arxiv.org/abs/2008.10312},
  code = {https://github.com/Randl/kmeans_selfsuper},
  arxiv = {2008.10312},
}