This repository contains PyTorch evaluation code, training code and pretrained models for PVT (Pyramid Vision Transformer).
Like ResNet, PVT is a pure transformer backbone that can be easily plugged in most downstream task models.
With a comparable number of parameters, PVT-Small RetinaNet achieves 40.4 AP on the COCO dataset, surpassing ResNet50 RetinNet (36.3 AP) by 4.1 AP.
Figure 1: Performance of RetinaNet 1x with different backbones.
This repository is developed on top of pytorch-image-models and deit.
For details see Pyramid Vision Transformer: A Versatile Backbone for Dense Prediction without Convolutions.
If you use this code for a paper please cite:
@misc{wang2021pyramid,
title={Pyramid Vision Transformer: A Versatile Backbone for Dense Prediction without Convolutions},
author={Wenhai Wang and Enze Xie and Xiang Li and Deng-Ping Fan and Kaitao Song and Ding Liang and Tong Lu and Ping Luo and Ling Shao},
year={2021},
eprint={2102.12122},
archivePrefix={arXiv},
primaryClass={cs.CV}
}
- PVT Semantic FPN configs & models
- PVT DETR/Sparse R-CNN config & models
- PVT Trans2Seg config & models
First, clone the repository locally:
git clone https://github.com/whai362/PVT.git
Then, install PyTorch 1.6.0 and torchvision 0.7.0 and pytorch-image-models 0.3.2:
conda install -c pytorch pytorch torchvision
pip install timm==0.3.2
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:
/path/to/imagenet/
train/
class1/
img1.jpeg
class2/
img2.jpeg
val/
class1/
img3.jpeg
class/2
img4.jpeg
Detection configs & models see here.
Method | Lr schd | box AP | mask AP | Config | Download |
---|---|---|---|---|---|
PVT-Tiny RetinaNet (800x) | 1x | 36.7 | - | config | Todo. |
PVT-Small RetinaNet (640x) | 1x | 38.7 | - | config | model |
PVT-Small RetinaNet (800x) | 1x | 40.4 | - | config | model |
R50 DETR | 50ep | 32.3 | - | config | Todo. |
PVT-Small DETR | 50ep | 34.7 | - | config | Todo. |
R50 DETR | 50ep | 32.3 | - | config | Todo. |
PVT-Tiny Mask RCNN | 1x | 36.7 | 35.1 | config | Todo. |
PVT-Small Mask RCNN | 1x | 40.4 | 37.8 | config | Todo. |
We provide baseline PVT models pretrained on ImageNet 2012.
name | acc@1 | #params (M) | url |
---|---|---|---|
PVT-Tiny | 75.1 | 13.2 | 51 M, PyTorch<=1.5 |
PVT-Small | 79.8 | 24.5 | 93 M, PyTorch<=1.5 |
PVT-Medium | 81.2 | 44.2 | 168M |
PVT-Large | 81.7 | 61.4 | 234M |
To evaluate a pre-trained PVT-Small on ImageNet val with a single GPU run:
sh dist_train.sh pvt_small 1 /path/to/checkpoint_root --data-path /path/to/imagenet --resume /path/to/checkpoint_file --eval
This should give
* Acc@1 79.764 Acc@5 94.950 loss 0.885
Accuracy of the network on the 50000 test images: 79.8%
To train PVT-Small on ImageNet on a single node with 8 gpus for 300 epochs run:
sh dist_train.sh pvt_small 8 /path/to/checkpoint_root --data-path /path/to/imagenet
This repository is released under the Apache 2.0 license as found in the LICENSE file.