<h1 align="center">[TPAMI 2023] 基于四边形注意力的视觉Transformer<a href="https://arxiv.org/abs/2303.15105"><img src="https://yellow-cdn.veclightyear.com/0a4dffa0/db7298d9-dffb-4785-87af-c54f2b312622.svg" ></a></h1> <p align="center"> <h4 align="center">这是论文<a href="https://arxiv.org/abs/2303.15105">基于四边形注意力的视觉Transformer</a>的官方代码库。</h4> <h5 align="center"><em>张启明、张静、徐宇飞、陶大程</em></h5> <p align="center"> <a href="#news">新闻</a> | <a href="#abstract">摘要</a> | <a href="#method">方法</a> | <a href="#usage">使用</a> | <a href="#results">结果</a> | <a href="#statement">声明</a> </p>

当前应用

分类: 分层模型已发布；简单模型即将发布。

目标检测: 即将发布；

语义分割: 即将发布；

人体姿态: 即将发布

新闻

2024年1月24日

分类任务的分层模型代码已发布。

2023年12月30日

论文被IEEE模式分析与机器智能汇刊(TPAMI)接收，影响因子24.314。

2023年3月27日

论文已发布在arxiv上！代码整理完成后将公开发布。

摘要

<p align="left">本代码库包含了论文<a href="https://arxiv.org/abs/2303.15105">基于四边形注意力的视觉Transformer</a>的代码、模型和测试结果，该论文是我们ECCV 2022论文<a href="https://arxiv.org/pdf/2204.08446.pdf">VSA</a>的实质性扩展。我们将基于窗口的注意力机制扩展为一般的四边形公式，并提出了一种新颖的四边形注意力机制。我们采用了端到端可学习的四边形回归模块，该模块预测一个变换矩阵，将默认窗口转换为目标四边形以进行token采样和注意力计算，使网络能够对不同形状和方向的各种目标进行建模，并捕获丰富的上下文信息。通过minor的代码修改和可忽略的额外计算成本，我们的QFormer在各种视觉任务上优于现有的代表性（分层和简单）视觉Transformer，包括分类、目标检测、语义分割和姿态估计。

方法

<figure> <img src="https://yellow-cdn.veclightyear.com/0a4dffa0/e0b984d3-a7a7-4f9b-af2e-3897f467a0eb.jpg"> <figcaption align = "center"><b>图1 - 当前设计（手工制作的窗口）与四边形注意力的比较。</b></figcaption> </figure> <figure> <img src="https://yellow-cdn.veclightyear.com/0a4dffa0/4ca9ab43-f087-4659-a204-0df88edf7678.jpg"> <figcaption align = "center"><b>图2 - 我们提出的四边形注意力（QA）的流程。</b></figcaption> </figure> <figure> <img src="https://yellow-cdn.veclightyear.com/0a4dffa0/880ba495-736b-4821-b025-1004f7514c82.jpg"> <figcaption align = "center"><b>图3 - 四边形注意力中的变换过程。</b></figcaption> </figure> <figure> <img src="https://yellow-cdn.veclightyear.com/0a4dffa0/e55487c1-ccf7-4376-8dcc-0d3d86beb7ad.jpg"> <figcaption align = "center"><b>图4 - 我们的简单QFormer<sub>p</sub> (a)和分层QFormer<sub>h</sub> (b)的架构。</b></figcaption> </figure>

使用

要求

PyTorch==1.7.1
torchvision==0.8.2
timm==0.3.2

Apex是可选的，用于更快的训练速度。

git clone https://github.com/NVIDIA/apex
cd apex
pip install -v --disable-pip-version-check --no-cache-dir --global-option="--cpp_ext" --global-option="--cuda_ext" ./

其他要求

pip install opencv-python==4.4.0.46 termcolor==1.1.0 yacs==0.1.8 timm==0.4.9
pip install einops

训练 & 评估

对于ImageNet-1K的分类任务，从头开始训练，运行：

CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7 \
python -m torch.distributed.launch \
  --nnodes ${NNODES} \
  --node_rank ${SLURM_NODEID} \
  --master_addr ${MHOST} \
  --master_port 25901 \
  --nproc_per_node 8 \
  ./main.py \
  --cfg configs/swin/qformer_tiny_patch4_window7_224.yaml \
  --data-path ${IMAGE_PATH} \
  --batch-size 128 \
  --tag 1024-dpr20-coords_lambda1e-1 \
  --distributed \
  --coords_lambda 1e-1 \
  --drop_path_rate 0.2 \

对于单GPU训练，运行

python ./main.py \
  --cfg configs/swin/qformer_tiny_patch4_window7_224.yaml \
  --data-path ${IMAGE_PATH} \
  --batch-size 128 \
  --tag 1024-dpr20-coords_lambda1e-1 \
  --coords_lambda 1e-1 \
  --drop_path_rate 0.2 \

要进行评估，运行：

CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7 \
python -m torch.distributed.launch \
  --nnodes ${NNODES} \
  --node_rank ${SLURM_NODEID} \
  --master_addr ${MHOST} \
  --master_port 25901 \
  --nproc_per_node 8 \
  ./main.py \
  --cfg configs/swin/qformer_tiny_patch4_window7_224.yaml \
  --data-path ${IMAGE_PATH} \
  --batch-size 128 \
  --tag eval \
  --distributed \
  --resume ${MODEL PATH} \
  --eval

对于单GPU评估，运行

python ./main.py \
  --cfg configs/swin/qformer_tiny_patch4_window7_224.yaml \
  --data-path ${IMAGE_PATH} \
  --batch-size 128 \
  --tag eval \
  --resume ${MODEL PATH} \
  --eval

结果

简单模型的结果

在ImageNet-1K上使用MAE预训练模型的分类结果

模型	分辨率	acc@1	权重 & 日志
ViT-B + 窗口注意力	224x224	81.2	\
ViT-B + 移位窗口	224x224	82.0	\
QFormer<sub>p</sub>-B	224x224	82.9	即将发布

在COCO上使用MAE预训练模型和Mask RCNN检测器的检测结果，遵循<a href="https://arxiv.org/abs/2203.16527">ViTDet</a>

模型	边界框平均精度	掩码平均精度	参数量	权重与日志
ViTDet-B	51.6	45.9	111M	\
QFormer<sub>p</sub>-B	52.3	46.6	111M	即将发布

使用MAE预训练模型和UPerNet分割器在ADE20k上的语义分割结果

模型	图像尺寸	mIoU	mIoU*	权重与日志
ViT-B + 窗口注意力	512x512	39.7	41.8	\
ViT-B + 移位窗口注意力	512x512	41.6	43.6	\
QFormer<sub>p</sub>-B	512x512	43.6	45.0	即将发布
ViT-B + 窗口注意力	640x640	40.2	41.5	\
ViT-B + 移位窗口注意力	640x640	42.3	43.5	\
QFormer<sub>p</sub>-B	640x640	44.9	46.0	即将发布

使用MAE预训练模型在COCO数据集上的人体姿态估计结果，遵循<a href="https://arxiv.org/abs/2204.12484">ViTPose</a>

注意力机制	模型	AP	AP<sub>50</sub>	AR	AR<sub>50</sub>	权重与日志
窗口	ViT-B	66.4	87.7	72.9	91.9	\
移位窗口	ViT-B	76.4	90.9	81.6	94.5	\
四边形	ViT-B	77.0	90.9	82.0	94.7	即将发布
窗口 + 全局	ViT-B	76.9	90.8	82.1	94.7	\
移位窗口 + 全局	ViT-B	77.2	90.9	82.2	94.7	\
四边形 + 全局	ViT-B	77.4	91.0	82.4	94.9	即将发布

层级模型结果

ImageNet-1K主要结果

名称	分辨率	Top-1准确率	Top-5准确率	真实Top-1准确率	权重与日志
Swin-T	224x224	81.2	\	\	\
DW-T	224x224	82.0	\	\	\
Focal-T	224x224	82.2	95.9
QFormer<sub>h</sub>-T	224x224	82.5	96.2	87.5	模型 & [日志](https://github.com/ViTAE-Transformer/QFormer/blob/main/logs/QFormer-T.txt
Swin-S	224x224	83.2	96.2	\	\
Focal-S	224x224	83.5	96.2	\	\
QFormer<sub>h</sub>-S	224x224	84.0	96.8	88.6	模型 & [日志](https://github.com/ViTAE-Transformer/QFormer/blob/main/logs/QFormer-S.txt
Swin-B	224x224	83.4	96.5	\	\
DW-B	224x224	83.4	\	\	\
Focal-B	224x224	83.8	96.5	\	\
QFormer<sub>h</sub>-B	224x224	84.1	96.8	88.7	模型 & [日志](https://github.com/ViTAE-Transformer/QFormer/blob/main/logs/QFormer-B.txt

目标检测结果

Mask R-CNN

骨干网络	预训练	学习率调度	边界框mAP	掩码mAP	参数量	配置文件	日志	模型
Swin-T	ImageNet-1K	1x	43.7	39.8	48M	\	\	\
DAT-T	ImageNet-1K	1x	44.4	40.4	48M	\	\	\
Focal-T	ImageNet-1K	1x	44.8	41.0	49M	\	\	\
QFormer<sub>h</sub>-T	ImageNet-1K	1x	45.9	41.5	49M	配置文件	日志	OneDrive
Swin-T	ImageNet-1K	3x	46.0	41.6	48M	\	\	\
DW-T	ImageNet-1K	3x	46.7	42.4	49M	\	\	\
DAT-T	ImageNet-1K	3x	47.1	42.4	48M	\	\	\
DAT-T	ImageNet-1K	3x	47.1	42.4	48M	\	\	\
QFormer<sub>h</sub>-T	ImageNet-1K	3x	47.5	42.7	49M	配置文件	日志	OneDrive
Swin-S	ImageNet-1K	3x	48.5	43.3	69M	\	\	\
Focal-S	ImageNet-1K	3x	48.8	43.8	71M	\	\	\
DAT-S	ImageNet-1K	3x	49.0	44.0	69M	\	\	\
QFormer<sub>h</sub>-S	ImageNet-1K	3x	49.5	44.2	70M	配置文件	日志	OneDrive

Cascade Mask R-CNN

骨干网络	预训练	学习率调度	边界框mAP	掩码mAP	参数量	配置文件	日志	模型
Swin-T	ImageNet-1K	1x	48.1	41.7	86M	\	\	\
DAT-T	ImageNet-1K	1x	49.1	42.5	86M	\	\	\
QFormer<sub>h</sub>-T	ImageNet-1K	1x	49.8	43.0	87M	配置文件	日志	OneDrive
Swin-T	ImageNet-1K	3x	50.2	43.7	86M	\	\	\
QFormer<sub>h</sub>-T	ImageNet-1K	3x	51.4	44.7	87M	配置文件	日志	OneDrive
Swin-S	ImageNet-1K	3x	51.9	45.0	107M	\	\	\
QFormer<sub>h</sub>-S	ImageNet-1K	3x	52.8	45.7	108M	配置文件	日志	OneDrive

ADE20k语义分割结果

UperNet

骨干网络	预训练	学习率调度	mIoU	mIoU*	参数量	配置	日志	模型
Swin-T	ImageNet-1k	160k	44.5	45.8	60M	\	\	\
DAT-T	ImageNet-1k	160k	45.5	46.4	60M	\	\	\
DW-T	ImageNet-1k	160k	45.7	46.9	61M	\	\	\
Focal-T	ImageNet-1k	160k	45.8	47.0	62M	\	\	\
QFormer<sub>h</sub>-T	ImageNet-1k	160k	46.9	48.1	61M	即将推出	即将推出	即将推出
Swin-S	ImageNet-1k	160k	47.6	49.5	81M	\	\	\
DAT-S	ImageNet-1k	160k	48.3	49.8	81M	\	\	\
Focal-S	ImageNet-1k	160k	48.0	50.0	61M	\	\	\
QFormer<sub>h</sub>-S	ImageNet-1k	160k	48.9	50.3	82M	即将推出	即将推出	即将推出
Swin-B	ImageNet-1k	160k	48.1	49.7	121M	\	\	\
DW-B	ImageNet-1k	160k	48.7	50.3	125M	\	\	\
Focal-B	ImageNet-1k	160k	49.0	50.5	126M	\	\	\
QFormer<sub>h</sub>-B	ImageNet-1k	160k	49.5	50.6	123M	即将推出	即将推出	即将推出

声明

本项目仅供研究使用。如有任何其他问题，请联系 qmzhangzz at hotmail.com。

代码基于 Swin 项目。

引用 QFormer、VSA 和 ViTAE

@article{zhang2023vision,
  title={Vision Transformer with Quadrangle Attention},
  author={Zhang, Qiming and Zhang, Jing and Xu, Yufei and Tao, Dacheng},
  journal={arXiv preprint arXiv:2303.15105},
  year={2023}
}
@inproceedings{zhang2022vsa,
  title={VSA: learning varied-size window attention in vision transformers},
  author={Zhang, Qiming and Xu, Yufei and Zhang, Jing and Tao, Dacheng},
  booktitle={Computer Vision--ECCV 2022: 17th European Conference, Tel Aviv, Israel, October 23--27, 2022, Proceedings, Part XXV},
  pages={466--483},
  year={2022},
  organization={Springer}
}
@article{zhang2023vitaev2,
  title={Vitaev2: Vision transformer advanced by exploring inductive bias for image recognition and beyond},
  author={Zhang, Qiming and Xu, Yufei and Zhang, Jing and Tao, Dacheng},
  journal={International Journal of Computer Vision},
  pages={1--22},
  year={2023},
  publisher={Springer}
}
@article{xu2021vitae,
  title={Vitae: Vision transformer advanced by exploring intrinsic inductive bias},
  author={Xu, Yufei and Zhang, Qiming and Zhang, Jing and Tao, Dacheng},
  journal={Advances in Neural Information Processing Systems},
  volume={34},
  year={2021}
}

我们的其他 Transformer 工作

ViTPose：请参见 <a href="https://github.com/ViTAE-Transformer/ViTPose">用于人体姿态估计的基线模型 ViTPose</a>；

VSA：请参见 <a href="https://github.com/ViTAE-Transformer/ViTAE-VSA">用于图像分类和目标检测的 ViTAE-Transformer</a>；

ViTAE & ViTAEv2：请参见 <a href="https://github.com/ViTAE-Transformer/ViTAE-Transformer">用于图像分类、目标检测和语义分割的 ViTAE-Transformer</a>；

抠图：请参见 <a href="https://github.com/ViTAE-Transformer/ViTAE-Transformer-Matting">用于抠图的 ViTAE-Transformer</a>；

遥感：请参见 <a href="https://github.com/ViTAE-Transformer/ViTAE-Transformer-Remote-Sensing">用于遥感的 ViTAE-Transformer</a>；<a href="https://github.com/ViTAE-Transformer/Remote-Sensing-RVSA">推进普通视觉 Transformer 向遥感基础模型发展</a>；