GATCluster: Self-Supervised Gaussian-Attention Network for Image Clustering

By Chuang Niu, Jun Zhang, Ge Wang, and Jimin Liang

Introduction

This project is the Pytorch implementation of the paper at ECCV 2020.

Installation

Assuming Anaconda with python 3.6, the required packages for this project can be installed as:

conda install pytorch torchvision cudatoolkit=10.1 -c pytorch # The latest Pytorch version (1.5.1) has been tested.
conda install -c conda-forge addict
conda install matplotlib tqdm scikit-learn requests

Then, clone this repo

git clone https://github.com/niuchuangnn/GATCluster.git
cd GATCluster

Run

To train the model, simply run the following commands.

For STLl10,

python ./tools/cluster.py --config-file ./configs/stl10/gatcluster.py

For ImageNet10,

python ./tools/cluster.py --config-file ./configs/imagenet10/gatcluster.py

The dataset will be downloaded automatically and saved to ./datasets/ when missing. Our method is memory-efficient, a single GPU with 8G memory is enough for deep clustering.

Visualization of attention map

Before visualizing the attention maps, a model should be trained and set the corresponding path in ./tools/visualize_att_map.py. You can also use our trained model on STL10 at here, and place it in ./results/stl10/gatcluster/. This model is reproduced and its accuracy is slightly better than the best result reported in our paper.

Then, run

python ./tools/visualize_att_map.py

The results will be saved in ./results/stl10/att_maps/ as:

For ImageNet10, our trained model can be downloaded at here, and placed into ./results/imagenet10/gatcluster/ for visualization.

Then, run

python ./tools/visualize_att_map_imagenet10.py

The results will be saved in ./results/imagenet10/att_maps/ as:

Citation

@inproceedings{gatcluster2020,
  title={GATCluster: Self-Supervised Gaussian-Attention Network for Image Clustering},
  author={Niu, Chuang and Zhang, Jun and Wang, Ge and Liang, Jimin},
  booktitle={European Conference on Computer Vision (ECCV)},
  year={2020}
}