Awesome
Contrastive-Hierarchical-Clustering
This is the official code for the paper Contrastive Hierarchical Clustering (ECML PKDD 2023). Preprint.
Our repository is build on top of the SimCLR implementation done by leftthomas.
Environment setup
Our code runs on a single GPU. It does not support multi-GPUs.
The code is compatible with Pytorch >= 1.7
. See environment.yml for all prerequisites, and you can also install them using the following command.
conda env create -f environment.yml
Alternative instalation
conda create -n chc python=3.7.6
conda activate chc
conda install pytorch==1.7.0 torchvision==0.8.0 pytorch-cuda=11.7 -c pytorch -c nvidia
conda install tensorboard
conda install pandas
conda install -c anaconda scikit-learn
pip install omegaconf
pip install thop
conda install -c conda-forge tqdm
Training
For every dataset configuration file cfg/<datasetname>
specifies model architectures and hyperparameters for the training.
python main.py --dataset-name cifar10 # Start training for CIFAR10
python main.py --dataset-name cifar100 # Start training for CIFAR100
python main.py --dataset-name stl10 # Start training for STL10
python main.py --dataset-name imagenet10 # Start training for ImageNet10
python main.py --dataset-name imagenetdogs # Start training for ImageNetDogs
Evaluation and pretrained models
After the training, model can be evaluate with evaulate.py
to generate NMI, ARI and ACC metrics.
Pretrained models are available in pre-trained
folder.
python evaluate.py --save_point ./pre-trained/CIFAR10/models/ --dataset-name cifar10 # Generate metrics for CIFAR10
python evaluate.py --save_point ./pre-trained/STL10/models/ --dataset-name stl10 # Generate metrics for STL10
python evaluate.py --save_point ./pre-trained/ImageNet10/models/ --dataset-name imagenet10 # Generate metrics for ImageNet10
python evaluate.py --save_point ./pre-trained/ImageNetDogs/models/ --dataset-name imagenetdogs # Generate metrics for ImageNetDogs