The clustr from ytongxie

clustr's Introduction

ClusTR: Exploring Efficient Self-attention via Clustering for Vision Transformers

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.

Requirements

CUDA 11.3
Python 3.10
Pytorch 1.11
Torchvision 0.12.0
mmcv 1.4.8
timm

Data preparation

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

Training

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

Evaluation

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}
}

Recommend Projects