wangjingbo1219 / PVT

Geek Repo:Geek Repo

Github PK Tool:Github PK Tool

Pyramid Vision Transformer: A Versatile Backbone for Dense Prediction without Convolutions

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

Todo List

  • PVT + Semantic FPN configs & models
  • PVT + DETR/Sparse R-CNN config & models
  • PVT + Trans2Seg config & models

Usage

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

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:

/path/to/imagenet/
  train/
    class1/
      img1.jpeg
    class2/
      img2.jpeg
  val/
    class1/
      img3.jpeg
    class/2
      img4.jpeg

Model Zoo

Object Detection

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.

Image Classification

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

Evaluation

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%

Training

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

License

This repository is released under the Apache 2.0 license as found in the LICENSE file.

About

License:Apache License 2.0


Languages

Language:Python 99.0%Language:Shell 1.0%