GroupMixFormer

GroupMixFormer: 利用群组混合注意力推进视觉Transformer

</div>

GroupMixFormer的官方PyTorch实现，对应论文:

🐱 摘要

<p style="text-align: justify;"> 我们提出了GroupMixFormer，它使用群组混合注意力(GMA)作为传统自注意力的一种先进替代方案。GMA旨在同时捕捉token之间以及不同token群组之间的相关性，适应多种群组大小。 </p> <details><summary><b>完整摘要</b></summary> <p style="text-align: justify;"> 视觉Transformer(ViTs)通过多头自注意力(MHSA)建模长程依赖关系来增强视觉识别能力，这通常以查询-键-值计算的形式实现。然而，由查询和键生成的注意力图只能在单一粒度上捕捉token到token的相关性。在本文中，我们认为自注意力应该有一个更全面的机制来捕捉token和群组(即多个相邻token)之间的相关性，以获得更高的表示能力。因此，我们提出了群组混合注意力(GMA)作为传统自注意力的先进替代方案，它可以同时捕捉token到token、token到群组以及群组到群组的相关性，适用于各种群组大小。为此，GMA将查询、键和值均匀分割成若干段，并执行不同的群组聚合以生成群组代理。注意力图基于token和群组代理的混合计算得出，并用于重新组合值中的token和群组。基于GMA，我们提出了一个强大的骨干网络，即GroupMixFormer，它在图像分类、目标检测和语义分割任务中取得了最先进的性能，且参数量少于现有模型。例如，GroupMixFormer-L(7030万参数，384^2输入)在不使用外部数据的情况下，在ImageNet-1K上达到了86.2%的Top-1准确率，而GroupMixFormer-B(4580万参数)在ADE20K上达到了51.2%的mIoU。 </p> </details>

🚩 更新

新功能

✅ 2023年10月18日。发布训练代码。
✅ 2023年10月18日。发布推理代码。
✅ 2023年10月18日。发布分类任务的预训练模型。

⚙️ 使用方法

1 - 安装

创建一个新的conda虚拟环境

conda create -n groupmixformer python=3.8 -y
conda activate groupmixformer

按照官方说明安装Pytorch>=1.8.0, torchvision>=0.9.0。例如：

pip install torch==1.8.0+cu111 torchvision==0.9.0+cu111 -f https://download.pytorch.org/whl/torch_stable.html

克隆此仓库并安装所需的包：

git clone https://github.com/AILab-CVC/GroupMixFormer.git
pip install timm==0.4.12 tensorboardX six tensorboard ipdb yacs tqdm fvcore

论文中的结果是使用torch==1.8.0+cu111 torchvision==0.9.0+cu111 timm==0.4.12产生的。
其他依赖：mmdetection和mmsegmentation是下游迁移的可选依赖。

2 - 数据准备

从http://image-net.org/下载并解压ImageNet训练和验证图像。目录结构如下：

│path/to/imagenet/
├──train/
│  ├── n01440764
│  │   ├── n01440764_10026.JPEG
│  │   ├── n01440764_10027.JPEG
│  │   ├── ......
│  ├── ......
├──val/
│  ├── n01440764
│  │   ├── ILSVRC2012_val_00000293.JPEG
│  │   ├── ILSVRC2012_val_00002138.JPEG
│  │   ├── ......
│  ├── ......

3 - 训练脚本

要在单个节点上使用8个GPU训练GroupMixFormer-Small在ImageNet-1k上300个epoch，请运行：

python3 -m torch.distributed.launch --nproc_per_node 8 --nnodes 1 --use_env train.py \
  --data-path <你的数据路径> \
  --batch-size 64 \
  --output <你的目标输出路径> \
  --cfg ./configs/groupmixformer_small.yaml \
  --model-type groupmixformer \
  --model-file groupmixformer.py \
  --tag groupmixformer_small

或者你可以直接运行以下脚本：

bash launch_scripts/run_train.sh

对于多节点训练，请参考代码：multi_machine_start.py

4 - 推理脚本

要在单个节点上评估ImageNet-1k上的GroupMixFormer-Small，请确定预训练权重的路径并运行：

CUDA_VISIBLE_DEVICES=1 OMP_NUM_THREADS=1 python3 -m torch.distributed.launch --nproc_per_node 1 --nnodes 1 --use_env test.py \
  --data-path <你的数据路径> \
  --batch-size 64 \
  --output <你的目标输出路径> \
  --cfg ./configs/groupmixformer_small.yaml \
  --model-type groupmixformer \
  --model-file groupmixformer.py \
  --tag groupmixformer_small

或者你可以直接运行以下脚本：

bash launch_scripts/run_eval.sh

这应该会得到

* Acc@1 83.400 Acc@5 96.464

⏬ 模型库

我们提供了在ImageNet 2012上预训练的GroupMixFormer模型。你可以下载相应的预训练模型并将其移动到./pretrained文件夹。

名称	分辨率	acc@1	参数量	FLOPs	模型 - 配置
GroupMixFormer-M	224x224	79.6	5.7M	1.4G	模型 - 配置
GroupMixFormer-T	224x224	82.6	11.0M	3.7G	模型 - 配置
GroupMixFormer-S	224x224	83.4	22.4M	5.2G	模型 - 配置
GroupMixFormer-B	224x224	84.7	45.8M	17.6G	模型 - 配置
GroupMixFormer-L	224x224	85.0	70.3M	36.1G	模型 - 配置