ViTAE: Vision Transformer Advanced by Exploring Intrinsic Inductive Bias

Current applications

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

Object Detection: Please see Usage for a quick start;

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

Animal Pose Estimation: Please see ViTAE-Transformer for animal pose estimation;

Matting: Please see ViTAE-Transformer for matting;

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

Updates

24/03/2022 ViTAEv2 with various downstream tasks released!

Introduction

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.

This repository contains the code, models, test results for the paper ViTAE: Vision Transformer Advanced by Exploring Intrinsic Inductive Bias. It contains several reduction cells and normal cells to introduce scale-invariance and locality into vision transformers.

Results and Models

Mask R-CNN

Backbone	Pretrain	Lr Schd	box mAP	mask mAP	#params	config	log	model
Swin-T	ImageNet-1K	1x	43.7	39.8	48M	config	github/baidu	github/baidu
ViTAEv2-S	ImageNet-1K	1x	46.3	41.8	37M	config	github	Coming Soon
Swin-T	ImageNet-1K	3x	46.0	41.6	48M	config	github/baidu	github/baidu
ViTAEv2-S	ImageNet-1K	3x	47.8	42.6	37M	config	github	Coming Soon

Cascade Mask R-CNN

Backbone	Pretrain	Lr Schd	box mAP	mask mAP	#params	config	log	model
Swin-T	ImageNet-1K	1x	48.1	41.5	86M	config	github/baidu	github/baidu
ViTAEv2-S	ImageNet-1K	1x	50.6	43.6	75M	config	github	Coming Soon
Swin-T	ImageNet-1K	3x	50.2	43.5	86M	config	github/baidu	github/baidu
ViTAEv2-S	ImageNet-1K	3x	51.4	44.5	75M	config	github	Coming Soon

Notes:

The drop path rate needs to be tuned for various models and heads.

Usage

Installation

Please refer to get_started.md for installation and dataset preparation.

Inference

# single-gpu testing
python tools/test.py <CONFIG_FILE> <DET_CHECKPOINT_FILE> --eval bbox segm

# multi-gpu testing
tools/dist_test.sh <CONFIG_FILE> <DET_CHECKPOINT_FILE> <GPU_NUM> --eval bbox segm

Training

To train a detector with pre-trained models, run:

# single-gpu training
python tools/train.py <CONFIG_FILE> --cfg-options model.pretrained=<PRETRAIN_MODEL> [model.backbone.use_checkpoint=True] [other optional arguments]

# multi-gpu training
tools/dist_train.sh <CONFIG_FILE> <GPU_NUM> --cfg-options model.pretrained=<PRETRAIN_MODEL> [model.backbone.use_checkpoint=True] [other optional arguments]

For example, to train a Mask R-CNN model with a ViTAEv2-S backbone and 8 gpus, run:

tools/dist_train.sh configs/vitaev2/mask_rcnn_vitaev2_s_mstrain_480-800_adamw_1x_coco.py 8 --cfg-options model.pretrained=<PRETRAIN_MODEL>

Note: use_checkpoint is used to save GPU memory. Please refer to this page for more details.

Apex (optional):

We use apex for mixed precision training by default. To install apex, run:

git clone https://github.com/NVIDIA/apex
cd apex
pip install -v --disable-pip-version-check --no-cache-dir --global-option="--cpp_ext" --global-option="--cuda_ext" ./

If you would like to disable apex, modify the type of runner as EpochBasedRunner and comment out the following code block in the configuration files:

# do not use mmdet version fp16
fp16 = None
optimizer_config = dict(
    type="DistOptimizerHook",
    update_interval=1,
    grad_clip=None,
    coalesce=True,
    bucket_size_mb=-1,
    use_fp16=True,
)

Citing ViTAE and ViTAEv2

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

Other Links

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.

Statement

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