目标检测中探索简单视觉Transformer主干网络的非官方PyTorch实现

结果 | 更新 | 使用 | 待办事项 | 致谢

本分支包含目标检测中探索简单视觉Transformer主干网络的非官方 PyTorch 实现。感谢他们的出色工作！

本仓库在COCO上的结果

模型在4台A100机器上训练，每GPU处理2张图像，因此训练时的批量大小为64。

模型	预训练	设备	框架	Box mAP	Mask mAP	配置文件	日志	权重
ViT-Base	IN1K+MAE	TPU	Mask RCNN	51.1	45.5	配置	日志	OneDrive
ViT-Base	IN1K+MAE	GPU	Mask RCNN	51.1	45.4	配置	日志	OneDrive
ViTAE-Base	IN1K+MAE	GPU	Mask RCNN	51.6	45.8	配置	日志	OneDrive
ViTAE-Small	IN1K+Sup	GPU	Mask RCNN	45.6	40.1	配置	日志	OneDrive

更新

[2022-04-18] 探索使用小型1K监督训练模型（20M参数）用于ViTDet（45.6 mAP）。采用多阶段结构后，Swin-T 的结果为 46.0 mAP，而 ViTAEv2-S 在COCO数据集上使用Mask RCNN时达到 47.8 mAP。

[2022-04-17] 发布了ViT-B和ViTAE-B在MS COCO上的预训练权重和日志。这些模型完全使用PyTorch在GPU上训练完成。

[2022-04-16] 发布了基于ViT-Base模型的ViTDet初始非官方实现！该模型在检测和分割任务上分别取得了51.1 mAP和45.5 mAP的成绩。权重和日志将很快上传。

ViTAE Transformer的应用包括：图像分类 | 目标检测 | 语义分割 | 动物姿态分割 | 遥感 | 抠图

使用方法

我们使用PyTorch 1.9.0或NGC Docker 21.06，以及mmcv 1.3.9进行实验。

git clone https://github.com/open-mmlab/mmcv.git
cd mmcv
git checkout v1.3.9
MMCV_WITH_OPS=1 pip install -e .
cd ..
git clone https://github.com/ViTAE-Transformer/ViTDet.git
cd ViTDet
pip install -v -e .

安装完这两个仓库后，再安装timm和einops，即：

pip install timm==0.4.9 einops

从MAE或ViTAE下载预训练模型，然后通过以下方式开展实验：

# 单机运行
bash tools/dist_train.sh <Config PATH> <NUM GPUs> --cfg-options model.pretrained=<Pretrained PATH>

# 多机运行
python -m torch.distributed.launch --nnodes <Num Machines> --node_rank <Rank of Machine> --nproc_per_node <GPUs Per Machine> --master_addr <Master Addr> --master_port <Master Port> tools/train.py <Config PATH> --cfg-options model.pretrained=<Pretrained PATH> --launcher pytorch

待办事项

当前本仓库包含以下修改：

• 在RPN和头部中使用LN替代卷积
• 使用大规模jittor进行数据增强
• 使用来自MViT的RPE
• 延长训练周期并采用1024的测试尺寸
• 使用全局注意力层

此外还有其他待办事项：

• 实现用于全局信息传递的卷积块

• 调优Cascade RCNN模型

• 训练ViT模型以适应LVIS数据集

• 使用ViTDet框架训练ViTAE模型

致谢

我们感谢mmdetection、MAE、MViT以及BeiT提供的优秀实现。

引用ViTDet

@article{Li2022ExploringPV,
  title={Exploring Plain Vision Transformer Backbones for Object Detection},
  author={Yanghao Li and Hanzi Mao and Ross B. Girshick and Kaiming He},
  journal={ArXiv},
  year={2022},
  volume={abs/2203.16527}
}

对于ViTAE和ViTAEv2，请参考：

@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}
}

@article{zhang2022vitaev2,
  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={arXiv preprint arXiv:2202.10108},
  year={2022}
}

ViTDet 快速上手指南

ViTDet 是一个基于纯 Vision Transformer (ViT) 主干网络进行目标检测的非官方 PyTorch 实现。本指南将帮助你快速搭建环境并运行模型。

环境准备

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

• 操作系统: Linux (推荐 Ubuntu)
• Python: 3.7+
• PyTorch: 1.9.0 或更高版本 (或使用 NGC docker 21.06)
• GPU: 支持 CUDA 的 NVIDIA GPU (训练示例基于 4 张 A100，但单卡也可运行)
• 编译器: GCC 用于编译 mmcv 算子

安装步骤

请依次执行以下命令来安装依赖库和 ViTDet 项目。

1. 安装 MMCV

ViTDet 依赖 mmcv-full 特定版本 (1.3.9)。

git clone https://github.com/open-mmlab/mmcv.git
cd mmcv
git checkout v1.3.9
MMCV_WITH_OPS=1 pip install -e .
cd ..

提示: 如果下载速度慢，可尝试使用国内镜像源，例如：pip install -e . -i https://pypi.tuna.tsinghua.edu.cn/simple

2. 安装 ViTDet

克隆本项目并进行安装。

git clone https://github.com/ViTAE-Transformer/ViTDet.git
cd ViTDet
pip install -v -e .

3. 安装其他依赖

安装 timm 和 einops 库。

pip install timm==0.4.9 einops

4. 下载预训练权重

你需要下载主干网络的预训练权重（如 MAE 或 ViTAE），可从以下地址获取：

• MAE: https://github.com/facebookresearch/mae
• ViTAE: https://github.com/ViTAE-Transformer/ViTAE-Transformer

下载后，记下权重文件的路径（例如 /path/to/pretrained.pth）。

基本使用

安装完成后，你可以使用提供的脚本进行训练。以下是最简单的单机器多卡训练示例。

单机器训练

假设你使用了 4 张 GPU，配置文件路径为 configs/ViTDet/ViTDet-ViT-Base-100e.py，预训练权重路径为 /path/to/mae_pretrain_vit_base.pth。

bash tools/dist_train.sh configs/ViTDet/ViTDet-ViT-Base-100e.py 4 --cfg-options model.pretrained=/path/to/mae_pretrain_vit_base.pth

多机器分布式训练

如果在多台机器上进行分布式训练，请使用以下命令格式：

python -m torch.distributed.launch --nnodes <Num Machines> --node_rank <Rank of Machine> --nproc_per_node <GPUs Per Machine> --master_addr <Master Addr> --master_port <Master Port> tools/train.py configs/ViTDet/ViTDet-ViT-Base-100e.py --cfg-options model.pretrained=/path/to/mae_pretrain_vit_base.pth --launcher pytorch

参数说明：

• <Config PATH>: 模型配置文件的路径。
• <NUM GPUs>: 当前机器使用的 GPU 数量。
• model.pretrained: 指定预训练权重文件的绝对路径。