yuhuan-wu/P2T

[TPAMI22] Pyramid Pooling Transformer for Scene Understanding

This is the official repository for Pyramid Pooling Transformer (P2T). This repository contains:

• Full code for training/test
• Pretrained models in image classification, object detection, and semantic segmentation.

Related links: [Official PDF Download] [中译版全文] [5分钟中文解读]

Requirements:

• torch>=1.7+
• torchvision>=0.7.0+
• timm>=0.3.2

Validated on Torch 1.6/1.7/1.8, timm 0.3.2/0.4.12

Introduction

Pyramid pooling transformer (P2T) is a new generation backbone network,, benefiting many fundamental downstream vision tasks like object detection, semantic segmentation, and instance segmentation.

Although pyramid pooling has demonstrated its power on many downstream tasks such as object detection (SPP) and semantic segmentation (PSPNet), it has not been explored on the backbone network, which serves a cornerstone for many downstream vision tasks. P2T first bridges the gap between pyramid pooling and backbone network. The core idea is P2T is adapting pyramid pooling to the downsampling of the flatten sequences in computing the self-attention, simultaneously reducing the sequence length and capturing powerful multi-scale contextual features. Pyramid pooling is also very efficient and will only induce negligible computational cost.

In the experiments, P2T beats all the CNN/Transformer competitors such as ResNet, ResNeXt, Res2Net, PVT, Swin, Twins, and PVTv2, on image classification, semantic segmentation, object detection, and instance segmentation.

Image Classification

All models are trained on ImageNet1K dataset. You can see all weights/logs at this url: [Google Drive] BaiduPan, 提取码yhwu

Semantic Segmentation

ADE20K (val set)

The training and validation scripts can refer to the segmentation folder.

BaiduPan download link: BaiduPan, 提取码yhwu

Object Detection

Tested on the coco validation set

Use this address to access all pretrained weights and logs: [Google Drive]

BaiduPan download link: BaiduPan, 提取码yhwu

Instance Segmentation

Tested on the coco val set

APb denotes AP box metric, and APm is the AP mask metric.

Use this address to access all pretrained weights and logs: [Google Drive]

Train

Use the following commands to train P2T-Small for distributed learning with 8 GPUs:

python -m torch.distributed.launch --nproc_per_node=8 \
    --master_port=$((RANDOM+10000)) --use_env main.py --data-path ${YOUR_DATA_PATH} --batch-size 128 --model p2t_small --drop-path 0.1
# model names: --model p2t_tiny/p2t_small/p2t_base/p2t_large
# with --drop-path 0.1/0.1/0.3/0.3
# replace ${YOUR_DATA_PATH} with your data path that contains train/ val/ directory

Validate the performance

Download the pretrained weights to pretrained directory first. Then use the following commands to validate the performance:

python main.py --eval --resume pretrained/p2t_small.pth --model p2t_small

Citation

If you are using the code/model/data provided here in a publication, please consider citing our works:

@ARTICLE{wu2022p2t,
  author={Wu, Yu-Huan and Liu, Yun and Zhan, Xin and Cheng, Ming-Ming},
  journal={IEEE Transactions on Pattern Analysis and Machine Intelligence}, 
  title={{P2T}: Pyramid Pooling Transformer for Scene Understanding}, 
  year={2022},
  doi = {10.1109/tpami.2022.3202765},
}

Other Notes

If you meet any problems, please do not hesitate to contact us. Issues and discussions are welcome in the repository! You can also contact us via sending messages to this email: [email protected]

License

This code is released under the Creative Commons Attribution-NonCommercial-ShareAlike 4.0 International Public License for Non-Commercial use only. Any commercial use should get formal permission first.