This is the official pytorch implementation of the ClusTR:
Paper: ClusTR: Exploring Efficient Self-attention via Clustering for Vision Transformers.
The code is developed on top of pytorch-image-models and pvt.
CUDA 11.3
Python 3.10
Pytorch 1.11
Torchvision 0.12.0
mmcv 1.4.8
timm
Download and extract ImageNet train and val images from http://image-net.org/.
The directory structure is the standard layout for the torchvision datasets.ImageFolder
, and the training and validation data is expected to be in the train/
folder and val
folder respectively:
imagenet/
train/
class1/
img1.jpeg
class2/
img2.jpeg
val/
class1/
img3.jpeg
class/2
img4.jpeg
To train ClusTR-b1 on ImageNet on a single node with 8 gpus for 300 epochs run:
./dist_train.sh configs/ClusTR/ClusTR_b1.py 8 --data-path ImageNet --batch-size 128 --epochs 300
To evaluate a pre-trained ClusTR-b1 on ImageNet val with a single GPU run:
./dist_train.sh configs/ClusTR/ClusTR_b1.py 1 --data-path ImageNet --resume /path/to/checkpoint_file --eval
If you use this code for a paper please cite:
@article{xie2022clustr,
title={ClusTR: Exploring Efficient Self-attention via Clustering for Vision Transformers},
author={Xie, Yutong and Zhang, Jianpeng and Xia, Yong and Hengel, Anton van den and Wu, Qi},
journal={arXiv preprint arXiv:2208.13138},
year={2022}
}