Updates | Introduction | Usage | Results&Pretrained Models | Statement |

Image Classification: Please see ViTAE-Transformer for image classification;

Object Detection: Please see ViTAE-Transformer for object detection;

Sementic Segmentation: Please see ViTAE-Transformer for semantic segmentation;

Animal Pose Estimation: Please see Usage for a quick start;

Matting: Please see ViTAE-Transformer for matting;

Remote Sensing: Please see ViTAE-Transformer for Remote Sensing;

24/03/2022 Initial code release for ViTAEv2 on Animal Pose Estimation. The code is built upon mmpose, thanks for their amazing work!

This repository contains the code, models, and logs demonstrated in ViTAEv2: Vision Transformer Advanced by Exploring Inductive Bias for Image Recognition and Beyond. It contains several reduction cells and normal cells to introduce scale-invariance and locality into vision transformers. Compared to the first version, we stack the two cells in a multi-stage manner and explore the benefits of another inductive bias, i.e., window-based attentions without shifts for a better trade-off between speed, memory footprint, and performance.

Here we provide the example of training models with the AP-10k dataset and ViTAEv2-S. We thanks the help of mmpose project.

• Clone this repo:

git clone https://github.com/ViTAE-Transformer/ViTAE-Transformer.git
cd ViTAE-Transformer/Animal-Pose-Estimation

• Create a conda virtual environment and activate it:

conda create -n vitae python=3.7 -y
conda activate vitae

conda install pytorch==1.8.1 torchvision==0.9.1 cudatoolkit=10.2 -c pytorch -c conda-forge

• Install timm==0.3.4:

pip install timm==0.3.4

• Install mmcv-full==1.3.9

cd ../..
git clone https://github.com/open-mmlab/mmcv.git
cd mmcv
git checkout v1.3.9
MMCV_WITH_OPS=1 pip install -e .

• Install this repo

cd ../ViTAE-Transformer/Animal-Pose-Estimation
pip install -v -e .

• Install other requirements:

pip install pyyaml ipdb

We use AP10K dataset, you can download it from https://github.com/AlexTheBad/AP-10K. The file structure should look like:

$ tree data
data
├── ap10k
│   ├── annotations
      │   │── ap10k-train-split1.json
      │   |── ap10k-train-split2.json
      │   |── ap10k-train-split3.json
      │   │── ap10k-val-split1.json
      │   |── ap10k-val-split2.json
      │   |── ap10k-val-split3.json
      │   |── ap10k-test-split1.json
      │   |── ap10k-test-split2.json
      │   |── ap10k-test-split3.json
│   ├── data
│   │   ├── 000000000001.jpg
│   │   └── ...

The pretrained models can be downloaded from ViTAEv2-S-AP10K.

Then, the evaluation can be conducted as

python tools/test.py configs/animal/2d_kpt_sview_rgb_img/topdown_heatmap/ap10k/vitae_ap10k_256x256.py <POSE_CKPT_PATH>

To train the ViTAEv2-S model on the AP-10K dataset, please run

bash tools/dist_train.sh configs/animal/2d_kpt_sview_rgb_img/topdown_heatmap/ap10k/vitae_ap10k_256x256.py <GPU_NUM> --cfg-options model.pretrained=<PRETRAINED_MODEL_PATH>

• The trained model file will be saved under the work_dirs folder;
• The weights for the pretrained backbone models can be downloaded from ViTAE for Image Classification.

This project is for research purpose only. For any other questions please contact yufei.xu at outlook.com qmzhangzz at hotmail.com .

@article{xu2021vitae,
  title={Vitae: Vision transformer advanced by exploring intrinsic inductive bias},
  author={Xu, Yufei and Zhang, Qiming and Zhang, Jing and Tao, Dacheng},
  journal={Advances in Neural Information Processing Systems},
  volume={34},
  year={2021}
}
@article{zhang2022vitaev2,
  title={ViTAEv2: Vision Transformer Advanced by Exploring Inductive Bias for Image Recognition and Beyond},
  author={Zhang, Qiming and Xu, Yufei and Zhang, Jing and Tao, Dacheng},
  journal={arXiv preprint arXiv:2202.10108},
  year={2022}
}

Image Classification: See ViTAE for Image Classification

Object Detection: See ViTAE for Object Detection.

Semantic Segmentation: See ViTAE for Semantic Segmentation.

Animal Pose Estimation: See ViTAE for Animal Pose Estimation.

Matting: See ViTAE for Matting.

Remote Sensing: See ViTAE for Remote Sensing.