News
An extension of this work has been released on https://github.com/Yunfan-Li/Twin-Contrastive-Learning. It significantly improves the clustering performance and supports multi-gpu training.
Contrastive Clustering (CC)
This is the code for the paper "Contrastive Clustering" (AAAI 2021)
Dependency
- python>=3.7
- pytorch>=1.6.0
- torchvision>=0.8.1
- munkres>=1.1.4
- numpy>=1.19.2
- opencv-python>=4.4.0.46
- pyyaml>=5.3.1
- scikit-learn>=0.23.2
- cudatoolkit>=11.0
Usage
Configuration
There is a configuration file "config/config.yaml", where one can edit both the training and test options.
Training
After setting the configuration, to start training, simply run
python train.py
Since the traning strategy for STL-10 is slightly different from others (unlabeled data is used on ICH only while training and test split are used on both ICH and CCH), to start training on STL-10, run
python train_STL10.py
Test
Once the training is completed, there will be a saved model in the "model_path" specified in the configuration file. To test the trained model, run
python cluster.py
We uploaded the pretrained model which achieves the performance reported in the paper to the "save" folder for reference.
Dataset
CIFAR-10, CIFAR-100, STL-10 will be automatically downloaded by Pytorch. Tiny-ImageNet can be downloaded from http://cs231n.stanford.edu/tiny-imagenet-200.zip. For ImageNet-10 and ImageNet-dogs, we provided their description in the "dataset" folder.
Citation
If you find CC useful in your research, please consider citing:
@inproceedings{li2021contrastive,
title={Contrastive clustering},
author={Li, Yunfan and Hu, Peng and Liu, Zitao and Peng, Dezhong and Zhou, Joey Tianyi and Peng, Xi},
booktitle={Proceedings of the AAAI Conference on Artificial Intelligence},
volume={35},
number={10},
pages={8547--8555},
year={2021}
}