GPViT: A High Resolution Non-Hierarchical Vision Transformer with Group Propagation

This repository contains the official PyTorch implementation of GPViT, a high-resolution non-hierarchical vision transformer architecture designed for high-performing visual recognition, which is introduced in our paper:

GPViT: A High Resolution Non-Hierarchical Vision Transformer with Group Propagation, Chenhongyi Yang*, Jiarui Xu*, Shalini De Mello, Elliot J. Crowley, Xiaolong Wang, ICLR 2023

Usage

Environment Setup

Our code base is built upon the MM-series toolkits. Specifically, classification is based on MMClassification; object detection is based on MMDetection; and semantic segmentation is based on MMSegmentation. Users can follow the official site of those toolkit to set up their environments. We also provide a sample setting up script as following:

conda create -n gpvit python=3.7 -y
source activate gpvit
pip install torch==1.7.1+cu101 torchvision==0.8.2+cu101 -f https://download.pytorch.org/whl/torch_stable.html
pip install -U openmim
mim install mmcv-full==1.4.8
pip install timm
pip install lmdb # for ImageNet experiments
pip install -v -e .
cd downstream/mmdetection  # set up object detection and instance segmentation
pip install -v -e . 
cd ../mmsegmentation # set up semantic segmentation
pip install -v -e .

Data Preparation

Please follow MMClassification, MMDetection and MMSegmentation to set up the ImageNet, COCO and ADE20K datasets. For ImageNet experiment, we convert the dataset to LMDB format to accelerate training and testing. For example, you can convert you own dataset by running:

python tools/dataset_tools/create_lmdb_dataset.py \
       --train-img-dir data/imagenet/train \
       --train-out data/imagenet/imagenet_lmdb/train \
       --val-img-dir data/imagenet/val \
       --val-out data/imagenet/imagenet_lmdb/val

After setting up, the datasets file structure should be as follows:

GPViT
|-- data
|   |-- imagenet
|   |   |-- imagenet_lmdb
|   |   |   |-- train
|   |   |   |   |-- data.mdb
|   |   |   |   |__ lock.mdb
|   |   |   |-- val
|   |   |   |   |-- data.mdb
|   |   |   |   |__ lock.mdb 
|   |   |-- meta
|   |   |   |__ ...
|-- downstream 
|   |-- mmsegmentation
|   |   |-- data
|   |   |   |-- ade
|   |   |   |   |-- ADEChallengeData2016
|   |   |   |   |   |-- annotations
|   |   |   |   |   |   |__ ...
|   |   |   |   |   |-- images
|   |   |   |   |   |   |__ ...
|   |   |   |   |   |-- objectInfo150.txt
|   |   |   |   |   |__ sceneCategories.txt
|   |   |__ ...
|   |-- mmdetection
|   |   |-- data
|   |   |   |-- coco
|   |   |   |   |-- train2017
|   |   |   |   |   |-- ...
|   |   |   |   |-- val2017
|   |   |   |   |   |-- ...
|   |   |   |   |-- annotations
|   |   |   |   |   |-- instances_train2017.json
|   |   |   |   |   |-- instances_val2017.json
|   |   |   |   |   |__ ...
|   |   |__ ...
|__ ...

ImageNet Classification

Training GPViT

# Example: Training GPViT-L1 model
zsh tool/dist_train.sh configs/gpvit/gpvit_l1.py 16 

Testing GPViT

# Example: Testing GPViT-L1 model
zsh tool/dist_test.sh configs/gpvit/gpvit_l1.py work_dirs/gpvit_l1/epoch_300.pth 16 --metrics accuracy

COCO Object Detection and Instance Segmentation

Run cd downstream/mmdetection first.

Training GPViT based Mask R-CNN

# Example: Training GPViT-L1 models with 1x and 3x+MS schedules
zsh tools/dist_train.sh configs/gpvit/mask_rcnn/gpvit_l1_maskrcnn_1x.py 16
zsh tools/dist_train.sh configs/gpvit/mask_rcnn/gpvit_l1_maskrcnn_3x.py 16

Training GPViT based RetinaNet

# Example: Training GPViT-L1 models with 1x and 3x+MS schedules
zsh tools/dist_train.sh configs/gpvit/retinanet/gpvit_l1_retinanet_1x.py 16
zsh tools/dist_train.sh configs/gpvit/retinanet/gpvit_l1_retinanet_3x.py 16

Testing GPViT based Mask R-CNN

# Example: Testing GPViT-L1 Mask R-CNN 1x model
zsh tools/dist_test.sh configs/gpvit/mask_rcnn/gpvit_l1_maskrcnn_1x.py work_dirs/gpvit_l1_maskrcnn_1x/epoch_12.pth 16 --eval bbox segm

Testing GPViT based RetinaNet

# Example: Testing GPViT-L1 RetinaNet 1x model
zsh tools/dist_test.sh configs/gpvit/retinanet/gpvit_l1_retinanet_1x.py work_dirs/gpvit_l1_retinanet_1x/epoch_12.pth 16 --eval bbox

ADE20K Semantic Segmentation

Run cd downstream/mmsegmentation first.

Training GPViT based semantic segmentation models

# Example: Training GPViT-L1 based SegFormer and UperNet models
zsh tools/dist_train.sh configs/gpvit/gpvit_l1_segformer.py 16
zsh tools/dist_train.sh configs/gpvit/gpvit_l1_upernet.py 16

Testing GPViT based semantic segmentation models

# Example: Testing GPViT-L1 based SegFormer and UperNet models
zsh tools/dist_test.sh configs/gpvit/gpvit_l1_segformer.py work_dirs/gpvit_l1_segformer/iter_160000.pth 16 --eval mIoU
zsh tools/dist_test.sh configs/gpvit/gpvit_l1_upernet.py work_dirs/gpvit_l1_upernet/iter_160000.pth 16 --eval mIoU

Benchmark results

ImageNet-1k Classification

COCO Mask R-CNN 1x Schedule

COCO Mask R-CNN 3x+MS Schedule

COCO RetinaNet 1x Schedule

COCO RetinaNet 3x+MS Schedule

ADE20K UperNet

ADE20K SegFormer

Citation

@InProceedings{yang2023gpvit,
      title={{GPViT: A High Resolution Non-Hierarchical Vision Transformer with Group Propagation}}, 
      author={Chenhongyi Yang and Jiarui Xu and Shalini De Mello and Elliot J. Crowley and Xiaolong Wang},
      journal={ICLR}
      year={2023},
}