更新

• (2022/08/09) 聚生子分割（polyp-pvt）和视觉语言建模的应用示例。
• (2020/06/21) PVTv2 的代码已发布！PVTv2 在 PVTv1 的基础上进行了大幅改进，并且在 ImageNet-1K 预训练的情况下，性能优于 Swin Transformer。

梯度金字塔视觉Transformer

图片来自《变形金刚：复仇之战》。

本仓库包含 PVTv1 和 PVTv2 在图像分类、目标检测和语义分割任务中的官方实现。

模型库

图像分类

分类配置及权重请见 >>>这里<<<。

• PVTv2 在 ImageNet-1K 上的表现

方法	尺寸	Acc@1	参数量 (M)
PVTv2-B0	224	70.5	3.7
PVTv2-B1	224	78.7	14.0
PVTv2-B2-Linear	224	82.1	22.6
PVTv2-B2	224	82.0	25.4
PVTv2-B3	224	83.1	45.2
PVTv2-B4	224	83.6	62.6
PVTv2-B5	224	83.8	82.0

• PVTv1 在 ImageNet-1K 上的表现

方法	尺寸	Acc@1	参数量 (M)
PVT-Tiny	224	75.1	13.2
PVT-Small	224	79.8	24.5
PVT-Medium	224	81.2	44.2
PVT-Large	224	81.7	61.4

目标检测

检测配置及权重请见 >>>这里<<<。

• PVTv2 在 COCO 上的表现

基线检测器

方法	主干网络	预训练	学习率调度	数据增强	box AP	mask AP
RetinaNet	PVTv2-b0	ImageNet-1K	1x	无	37.2	-
RetinaNet	PVTv2-b1	ImageNet-1K	1x	无	41.2	-
RetinaNet	PVTv2-b2	ImageNet-1K	1x	无	44.6	-
RetinaNet	PVTv2-b3	ImageNet-1K	1x	无	45.9	-
RetinaNet	PVTv2-b4	ImageNet-1K	1x	无	46.1	-
RetinaNet	PVTv2-b5	ImageNet-1K	1x	无	46.2	-
Mask R-CNN	PVTv2-b0	ImageNet-1K	1x	无	38.2	36.2
Mask R-CNN	PVTv2-b1	ImageNet-1K	1x	无	41.8	38.8
Mask R-CNN	PVTv2-b2	ImageNet-1K	1x	无	45.3	41.2
Mask R-CNN	PVTv2-b3	ImageNet-1K	1x	无	47.0	42.5
Mask R-CNN	PVTv2-b4	ImageNet-1K	1x	无	47.5	42.7
Mask R-CNN	PVTv2-b5	ImageNet-1K	1x	无	47.4	42.5

进阶检测器

方法	主干网络	预训练	学习率调度	数据增强	box AP	mask AP
Cascade Mask R-CNN	PVTv2-b2-Linear	ImageNet-1K	3x	是	50.9	44.0
Cascade Mask R-CNN	PVTv2-b2	ImageNet-1K	3x	是	51.1	44.4
ATSS	PVTv2-b2-Linear	ImageNet-1K	3x	是	48.9	-
ATSS	PVTv2-b2	ImageNet-1K	3x	是	49.9	-
GFL	PVTv2-b2-Linear	ImageNet-1K	3x	是	49.2	-
GFL	PVTv2-b2	ImageNet-1K	3x	是	50.2	-
Sparse R-CNN	PVTv2-b2-Linear	ImageNet-1K	3x	是	48.9	-
Sparse R-CNN	PVTv2-b2	ImageNet-1K	3x	是	50.1	-

• PVTv1 在 COCO 上的表现

检测器	主干网络	预训练	学习率调度	box AP	mask AP
RetinaNet	PVT-Tiny	ImageNet-1K	1x	36.7	-
RetinaNet	PVT-Small	ImageNet-1K	1x	40.4	-
Mask RCNN	PVT-Tiny	ImageNet-1K	1x	36.7	35.1
Mask RCNN	PVT-Small	ImageNet-1K	1x	40.4	37.8
DETR	PVT-Small	ImageNet-1K	50ep	34.7	-

语义分割

分割配置及权重请见 >>>这里<<<。

PVT-v2 + 语义分割请见 >>>这里<<<。

• PVTv1 在 ADE20K 上的表现

方法	主干网络	预训练	迭代次数	mIoU
Semantic FPN	PVT-Tiny	ImageNet-1K	40K	35.7
Semantic FPN	PVT-Small	ImageNet-1K	40K	39.8
Semantic FPN	PVT-Medium	ImageNet-1K	40K	41.6
Semantic FPN	PVT-Large	ImageNet-1K	40K	42.1

聚生子分割

Polyp-PVT：基于梯度金字塔视觉Transformer的聚生子分割。pdf | 代码

视觉语言建模

时尚中的掩码视觉语言Transformer。pdf | 代码

许可证

本仓库根据 LICENSE 文件中的 Apache 2.0 许可证发布。

引用

如果您在论文中使用此代码，请引用：

PVTv1

@inproceedings{wang2021pyramid,
  title={Pyramid vision transformer: A versatile backbone for dense prediction without convolutions},
  author={Wang, Wenhai and Xie, Enze and Li, Xiang and Fan, Deng-Ping and Song, Kaitao and Liang, Ding and Lu, Tong and Luo, Ping and Shao, Ling},
  booktitle={Proceedings of the IEEE/CVF International Conference on Computer Vision},
  pages={568--578},
  year={2021}
}

PVTv2

@article{wang2021pvtv2,
  title={Pvtv2: Improved baselines with pyramid vision transformer},
  author={Wang, Wenhai and Xie, Enze and Li, Xiang and Fan, Deng-Ping and Song, Kaitao and Liang, Ding and Lu, Tong and Luo, Ping and Shao, Ling},
  journal={Computational Visual Media},
  volume={8},
  number={3},
  pages={1--10},
  year={2022},
  publisher={Springer}
}

联系方式

本仓库目前由 Wenhai Wang (@whai362)、Enze Xie (@xieenze) 和 Zhe Chen (@czczup) 维护。

PVT (Pyramid Vision Transformer) 快速上手指南

PVT 是一种专为密集预测任务（如目标检测、语义分割）设计的金字塔视觉 Transformer 骨干网络。本指南基于官方仓库，帮助您快速搭建环境并运行 PVTv2 模型。

环境准备

在开始之前，请确保您的系统满足以下要求：

• 操作系统: Linux (推荐 Ubuntu 18.04+)
• Python: 3.7 或更高版本
• GPU: NVIDIA GPU (支持 CUDA)，显存建议 8GB 以上（根据模型大小调整）
• 核心依赖:
  • PyTorch >= 1.8
  • torchvision
  • mmcv-full (用于检测和分割任务)
  • timm

前置检查： 请确保已安装与您的 CUDA 版本匹配的 PyTorch。

安装步骤

1. 克隆代码库

git clone https://github.com/whai362/PVT.git
cd PVT

2. 安装 Python 依赖

建议使用国内镜像源加速安装过程：

pip install -r requirements.txt -i https://pypi.tuna.tsinghua.edu.cn/simple

注意：如果 requirements.txt 未包含所有必要组件，请手动安装关键依赖：

pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu118
pip install timm mmcv-full -i https://pypi.tuna.tsinghua.edu.cn/simple

(注：mmcv-full 的版本需与您安装的 PyTorch/CUDA 版本严格对应，具体安装命令请参考 MMCV 官方文档)

3. 编译扩展模块 (如需)

如果涉及自定义算子编译，请运行：

python setup.py develop

基本使用

以下示例展示如何加载预训练的 PVTv2-B2 模型进行图像分类推理。

1. 下载预训练权重

从官方提供的链接下载权重文件（以 pvt_v2_b2.pth 为例），或访问 Model Zoo 获取其他型号权重。

2. 编写推理脚本

创建文件 demo_inference.py，输入以下代码：

import torch
from PIL import Image
from torchvision import transforms
from models.pvt_v2 import pvt_v2_b2

# 1. 配置设备
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

# 2. 初始化模型
model = pvt_v2_b2()
checkpoint = torch.load('pvt_v2_b2.pth', map_location=device)
model.load_state_dict(checkpoint)
model.to(device)
model.eval()

# 3. 数据预处理
transform = transforms.Compose([
    transforms.Resize((224, 224)),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
])

# 4. 加载图像并推理
image_path = 'your_image.jpg'  # 替换为您的图片路径
image = Image.open(image_path).convert('RGB')
input_tensor = transform(image).unsqueeze(0).to(device)

with torch.no_grad():
    output = model(input_tensor)
    prediction = torch.argmax(output, dim=1)

print(f"Predicted class index: {prediction.item()}")

3. 运行脚本

python demo_inference.py

---

提示：

• 目标检测与分割：请分别参考 detection/ 和 segmentation/ 目录下的配置文件，配合 MMDetection 或 MMSegmentation 框架使用。
• PVTv2-Seg：专门的分割任务实现请访问 PVTv2-Seg 仓库。