带有可变形注意力的视觉Transformer

本仓库包含论文《带有可变形注意力的视觉Transformer》（CVPR 2022，最佳论文入围者）[arXiv][视频][海报][CVPR页面]以及其扩展版本《DAT++：基于可变形注意力的空间动态视觉Transformer》[arXiv]的代码实现。

本仓库主要包含图像分类实验的实现。对于目标检测和实例分割任务，请参考DAT-Detection；对于语义分割任务，请参阅DAT-Segmentation以获取更多详细信息。

简介

动机

(a) 视觉Transformer (ViT) 凭借其大范围甚至全局的感受野，在众多任务中展现了优越性。然而，这种全局注意力机制带来了过高的计算开销。(b) Swin Transformer 提出了移位窗口注意力机制，这是一种更高效的稀疏注意力方法，具有线性的时间复杂度。尽管如此，这种人工设计的注意力模式可能会遗漏窗口之外的重要特征，而移位窗口的设计也限制了感受野的扩展，从而难以建模长距离依赖关系。(c) DCN 通过为每个查询学习偏移量来扩展标准卷积的感受野。然而，将这一技术直接应用于视觉Transformer并不容易，因为其空间复杂度呈二次方增长，且训练难度较大。(d) 可变形注意力 (DAT) 则是在特征图中重要区域的引导下，有效建模标记之间的关系。这种灵活的机制使自注意力模块能够聚焦于相关区域，捕捉更具信息量的特征。

方法

通过为网格参考点学习多组偏移量，从这些偏移后的位置采样变形的键和值。这种可变形注意力能够捕捉图像中最富信息的区域。在此基础上，我们提出了 可变形注意力Transformer (DAT) 和 DAT++，这是一种通用的骨干网络模型，适用于图像分类以及其他密集预测任务。

可视化结果

可视化结果显示，最重要的键用橙色圆圈标注，其中第三列中较大的圆圈表示更高的注意力得分。第四列和第五列展示了某些查询（红色星号）所关注的重要键（橙色圆圈）。这些重要键覆盖了物体的主要部分，这充分证明了DAT和DAT++的有效性。

依赖项

• NVIDIA GPU + CUDA 11.3
• Python 3.9
• PyTorch == 1.11.0
• torchvision == 0.12.0
• numpy == 1.20.3
• timm == 0.5.4
• einops == 0.6.1
• natten == 0.14.6
• PyYAML
• yacs
• termcolor

在ImageNet-1K分类数据集上评估预训练模型

我们提供了DAT++的tiny、small和base三个版本的预训练模型，具体如下表所示。

模型	分辨率	top-1准确率	配置文件	预训练权重
DAT-T++	224x224	83.9%	config	OneDrive / TsinghuaCloud
DAT-S++	224x224	84.6%	config	OneDrive / TsinghuaCloud
DAT-B++	224x224	84.9%	config	OneDrive / TsinghuaCloud
DAT-B++	384x384	85.9%	config	OneDrive / TsinghuaCloud

要评估某个模型，请先将预训练权重下载到本地机器，然后运行以下evaluate.sh脚本：

请注意：在进行训练或评估之前，请务必在train.sh或evaluate.sh脚本中将--data-path参数设置为ImageNet-1K数据存储的路径。

bash evaluate.sh <gpu数量> <配置文件路径> <预训练权重路径>

例如，假设使用8张GPU卡评估DAT-Tiny模型（dat_pp_tiny_in1k_224.pth），命令应为：

bash evaluate.sh 8 configs/dat_tiny.yaml dat_pp_tiny_in1k_224.pth

评估结果应显示：

[2023-09-04 17:18:15 dat_plus_plus] (main.py 301): INFO  * Acc@1 83.864 Acc@5 96.734
[2023-09-04 17:18:15 dat_plus_plus] (main.py 179): INFO 网络在50000张测试图像上的准确率为：83.9%

从头开始训练模型

为了从头开始训练模型，我们提供了一个简单的train.sh脚本。例如，要在单节点上使用8张GPU卡训练模型，可以使用以下命令：

bash train.sh 8 <配置文件路径> <实验标签>

此外，我们还提供了一个用于多机训练的train_slurm.sh脚本，支持更大的批量大小，如4096。

bash train_slurm.sh 32 <配置文件路径> <slurm作业名称>

请记得将脚本文件中的<path-to-imagenet>替换为你自己的ImageNet数据目录路径。

未来更新

• 分类任务的预训练模型。
• 目标检测的代码库及模型。
• 语义分割的代码库及模型。
• DAT-B++和DAT-L++在ImageNet-22K数据集上的预训练。
• DINO / Mask2Former用于系统级的目标检测/分割。
• CUDA / CUTLASS加速（可能）。

致谢

本代码是在Swin Transformer的基础上开发的，我们感谢他们高效且整洁的代码库。本研究的计算资源由杭州高飞人工智能基础研究有限公司提供。

引用

如果您认为我们的工作对您的研究有所帮助，请考虑引用以下文献：

@article{xia2023dat,
    title={DAT++：具有可变形注意力的空间动态视觉Transformer}, 
    author={Zhuofan Xia 和 Xuran Pan 和 Shiji Song 和 Li Erran Li 和 Gao Huang},
    year={2023},
    journal={arXiv 预印本 arXiv:2309.01430},
}

@InProceedings{Xia_2022_CVPR,
    author    = {Xia, Zhuofan 和 Pan, Xuran 和 Song, Shiji 和 Li, Li Erran 和 Huang, Gao},
    title     = {具有可变形注意力的视觉Transformer},
    booktitle = {IEEE/CVF 计算机视觉与模式识别会议（CVPR）论文集},
    month     = {6月},
    year      = {2022},
    pages     = {4794-4803}
}

联系方式

如果您有任何问题或疑虑，请发送邮件至 xzf23@mails.tsinghua.edu.cn。

DAT (Deformable Attention Transformer) 快速上手指南

DAT 是一种结合可变形注意力机制的视觉 Transformer 模型，旨在高效捕捉图像中的关键区域特征，在图像分类、目标检测和语义分割任务中表现优异。本指南将帮助你快速搭建环境并运行预训练模型。

环境准备

系统要求

• GPU: NVIDIA GPU
• CUDA: 11.3
• Python: 3.9
• PyTorch: 1.11.0
• torchvision: 0.12.0

前置依赖

请确保安装以下 Python 库：

numpy==1.20.3
timm==0.5.4
einops==0.6.1
natten==0.14.6
PyYAML
yacs
termcolor

提示：国内用户建议使用清华源或阿里源加速 pip 安装，例如：

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

安装步骤

1. 克隆仓库

   git clone https://github.com/LeapLabTHU/DAT.git
   cd DAT
   

2. 安装依赖 创建虚拟环境（推荐）并安装依赖：

   conda create -n dat python=3.9
   conda activate dat
   
   # 安装 PyTorch (根据 CUDA 11.3)
   pip install torch==1.11.0+cu113 torchvision==0.12.0+cu113 -f https://download.pytorch.org/whl/torch_stable.html
   
   # 安装其他依赖
   pip install -r requirements.txt
   

   注：若 requirements.txt 未直接提供，请手动安装上述“前置依赖”列表中列出的包。

3. 编译 NATTEN (如需) 如果 natten 未通过 pip 直接安装成功，可能需要从源码编译：

   pip install natten==0.14.6
   

基本使用

1. 准备数据集

下载 ImageNet-1K 数据集，并记录其本地路径（例如 /data/imagenet）。

2. 下载预训练模型

本项目提供 Tiny、Small 和 Base 版本的 DAT++ 预训练权重。国内用户推荐从 清华云盘 下载。

以 DAT-T++ 为例：

• 配置文件: configs/dat_tiny.yaml
• 权重下载: 清华云盘链接
• 将下载的权重文件重命名为 dat_pp_tiny_in1k_224.pth 并放在项目根目录。

3. 评估模型 (Evaluation)

在运行脚本前，请务必修改 evaluate.sh 或命令中的 --data-path 为你的 ImageNet 数据路径。

使用 8 张 GPU 评估 DAT-Tiny 模型的命令如下：

bash evaluate.sh 8 configs/dat_tiny.yaml dat_pp_tiny_in1k_224.pth

预期输出示例：

[2023-09-04 17:18:15 dat_plus_plus] (main.py 301): INFO  * Acc@1 83.864 Acc@5 96.734
[2023-09-04 17:18:15 dat_plus_plus] (main.py 179): INFO Accuracy of the network on the 50000 test images: 83.9%

4. 从头训练 (Training)

若需从头训练模型，请修改 train.sh 中的数据路径，然后运行：

# 单节点 8 卡训练示例
bash train.sh 8 configs/dat_tiny.yaml my_experiment_tag

注意：对于多机大规模训练（如 batch-size 4096），可使用 train_slurm.sh 脚本。