作者：伝迹、谦言、临在

欢迎使用我们最近开源的 EasyCV，主要聚焦于最新的 Vision Transformer 模型，以及相关的下游 CV 任务

开源地址：https://github.com/alibaba/EasyCV

ViTDet 其实是恺明团队 MAE 和 ViT-based Mask R-CNN 两个工作的延续。MAE 提出了 ViT 的无监督训练方法，而 ViT-based Mask R-CNN 给出了用 ViT 作为 backbone 的 Mask R-CNN 的训练技巧，并证明了 MAE 预训练对下游检测任务的重要性。而 ViTDet 进一步改进了一些设计，证明了 ViT 作为 backone 的检测模型可以匹敌基于 FPN 的 backbone(如 SwinT 和 MViT)检测模型。

ViT 作为检测模型的 backbone 需要解决两个问题：

1. 如何提升计算效率？

2. 如何得到多尺度特征？

ViT-based Mask R-CNN 给出了初步的解决方案，ViTDet 在此基础上，对如何得到多尺度特征做了进一步的改进。

如何提升计算效率

ViT 采用的 global self-attention 和图像输入大小(HW)的平方成正比，对于检测模型，其输入分辨率往往较大，此时用 ViT 作为 backbone 在计算量和内存消耗上都是非常惊人的，比如输入尺寸为 1024x1024，采用 ViT-B 训练 Mask R-CNN 单 batch 就需要消耗约 20-30GB 显存。为了解决这个问题，ViT-based Mask R-CNN 将 ViT 分成 4 个 stage，每个 stage 的前几个 block 采用 windowed self-attention，最后一个 block 采用 global self-attention，比较 table 3 (2)和(3)显著降低显存消耗和训练时间，而且效果只有轻微下降。

ViTDet 进一步研究了如何做 window 的信息聚合，除了采用 4 个 global self-attention 以外，还可以采用 4 个 residual block。如下表(a)所示，采用 4 个 conv blocks 效果是最好的，并且 basic block 效果最好(b)。另外表(c)和表(d)表明每个 stage 的最后一个 block 使用信息聚合，速度和精度的平衡是最好的。

Backbone

根据 ViT-based Mask R-CNN 论文 table 4 (94)的结果，用预训练过的 pos embed 加上 BEiT 提出的 relative position bias 效果最好，其中将 pos embed 迁移到下游任务需要对 pos embed 的进行 resize 操作。

最开始实现了一版共享的 relational position bias，精度上不去，感觉是打开方式不对，后来参照 ViTAE 的不共享 relational paosition bias，能加快收敛速度，代码如下。

def calc_rel_pos_spatial(    attn,    q,    q_shape,    k_shape,    rel_pos_h,    rel_pos_w,):    """    Spatial Relative Positional Embeddings.    """    sp_idx = 0    q_h, q_w = q_shape    k_h, k_w = k_shape        # Scale up rel pos if shapes for q and k are different.    q_h_ratio = max(k_h / q_h, 1.0)    k_h_ratio = max(q_h / k_h, 1.0)    dist_h = (        torch.arange(q_h)[:, None] * q_h_ratio -        torch.arange(k_h)[None, :] * k_h_ratio)    dist_h += (k_h - 1) * k_h_ratio    q_w_ratio = max(k_w / q_w, 1.0)    k_w_ratio = max(q_w / k_w, 1.0)    dist_w = (        torch.arange(q_w)[:, None] * q_w_ratio -        torch.arange(k_w)[None, :] * k_w_ratio)    dist_w += (k_w - 1) * k_w_ratio        Rh = rel_pos_h[dist_h.long()]    Rw = rel_pos_w[dist_w.long()]        B, n_head, q_N, dim = q.shape        r_q = q[:, :, sp_idx:].reshape(B, n_head, q_h, q_w, dim)    rel_h = torch.einsum('byhwc,hkc->byhwk', r_q, Rh)    rel_w = torch.einsum('byhwc,wkc->byhwk', r_q, Rw)        attn[:, :, sp_idx:, sp_idx:] = (        attn[:, :, sp_idx:, sp_idx:].view(B, -1, q_h, q_w, k_h, k_w) +        rel_h[:, :, :, :, :, None] + rel_w[:, :, :, :, None, :]).view(        B, -1, q_h * q_w, k_h * k_w)        return

将 ViT 作为 ViTDet 的预训练需要对 foward 过程进行改造，通过 window_partition 和 window_reverse 两个操作，对输入 feature 反复进行切 window 和还原，这样子可以充分利用 ViT 的预训练模型，同时提高检测的计算效率，论文中描述如上。

ViT-based Mask R-CNN 和 ViTDet 提到的 window size 都是 14x14，但是在输入分辨率为 1024x1024 的情况下，先经过一个 patch_embed，就变成了 64x64 的分辨率，64 是不能整除 14 的，

这里有两种处理方式：

1.在 patch_embed 之后加一个插值变成 56x56，从 ViT 输出的时候再插值回 64x64。

2.在 patch_embed 之后 pad 成 70x70，恢复成原图的时候裁剪成 64x64。

两种都试了一下，发现第二种不会掉点，而第一种可能会导致 position embedding 的不对齐，代码如下。

x = F.pad(x, (0, 0, pad_l, pad_r, pad_t, pad_b))_, Hp, Wp, _ = x.shape......if pad_r > 0 or pad_b > 0:    x = x[:, :H, :W, :].contiguous()    x = x.view(B_, H * W, C)

如何得到多尺度特征

ViT 模型是同质结构，如果采用的 patch size 为 16x16，那么最终就得到一种 1/16 的尺度特征。但是常用的检测模型往往需要多尺度特征，大多数 CNN 和金字塔 ViT 都可以适应这种输出，比如 ResNet 从不同 stage 提取 1/4，1/8，1/16 和 1/32 的特征，并送入 FPN 进一步融合得到多尺度特征。ViT-based Mask R-CNN 借鉴了 XCiT 的解决方案，将 ViT 的 transformer blocks 均分成 4 个部分，然后从 d/4，2d/4，3d/4 和 d 的输出分别提取 1/4，1/8，1/16 和 1/32 的特征(分别采用 2 个 stride=2 的反卷积，一个 stride=2 的反卷积，identity，stride=2 的 max pooling)，然后送入 FPN。

而 ViTDet 进一步简化了这种策略，直接用最后的 1/16 特征通过上采样(stride=2 的反卷积)或者下采样(stride=2 的 max pooling)得到 4 个尺度的特征，而且也不再用 FPN 来进一步融合特征，如上图 c 所示。

比较 table 1 (a)(b)(c)这种设计不仅简单，而且效果是最好的。

Simple feature pyramid

为了方便起见，简写为 SFP。SFP 先用 ViT 的最后一层构建出多尺度特征，然后分别接 1 个 1x1conv 做通道数 reduce，再接一个 3x3conv，论文中的描述如上。

论文中说在 conv 之后使用 layernorm，那么就需要不断的进行 reshape 操作，实现起来会比较复杂冗余。为了实现更加简洁干净，复现采用了 groupnorm 等价 layernorm 的方式(只要把 group 数设置成 1 就可以了)。

按照 ViTDet 论文中的说法，应该是只有 4 层尺度特征，但是标准的 FPN 一般是 5 层，不清楚具体实现的时候是用的几层，本实现默认使用 5 层。

Mask RCNN

论文中对于 mask rcnn 的修改如上，总结一下：

rpn head 2conv + LN

roi head 4conv + 1fc，BN 替换成 LN

mask head BN 替换成 LN

数据增强

也就是说训练的时候，采用 large scale jitter，然后 padding 成 1024；推理的时候保持长宽比最长边不超过 1024，然后 padding 成 1024。

超参数

预训练默认使用 mae_vit-base-p16-1600e，使用 AdamW 优化器，并且用 step-wise lr，bs64，warmup 250 iter，lr 1e-4，weight decay 0.1，ViT-B 的 drop_path_rate 设置成 0.1。

ViTDet 文章中说是 layer-wise lr decay 可以涨点 0.3 左右，但是我的实现导致最开始收敛很慢，感觉不一定有效。本实现默认不使用 layer-wise lr decay。

复现 ViTDet 的过程中，让我惊叹的除了单尺度构建多尺度特征精度超过 FPN 之外，还有一点是从 ViT -> SFP -> RPN Head -> RoI Head -> Mask Head 的一整套流程中竟然没有使用一个 BN，所有的 norm 都用 LN 替换掉了，这不是完全跟 NLP 对齐了。

预训练对比实验

另外 ViTDet 还对有监督预训练和无监督预训练 MAE 做了对比实验，可以看到 MAE 可以大幅度提升 AP，尤其是 ViT-L，甚至超过了 IN-21k 有监督训练效果，如 table 4 所示。

和其他层次化的 backbone 相比，ViTDet 也取得了最好的效果，如 table 5 所示。

效果图

最终复现的基于 ViT-Base 的 ViTDet_MaskRCNN 精度为 50.6，比论文低 0.6，可能还有一点点小细节没有考虑到的。

Tutorial

接下来，我们将通过一个实际的例子介绍如何基于 EasyCV 进行 ViTDet 算法的训练，也可以在该链接查看详细步骤。

一、安装依赖包

如果是在本地开发环境运行，可以参考该链接安装环境。若使用 PAI-DSW 进行实验则无需安装相关依赖，在 PAI-DSW docker 中已内置相关环境。

二、数据准备

你可以下载COCO2017数据，也可以使用我们提供了示例 COCO 数据

wget http://pai-vision-data-hz.oss-cn-zhangjiakou.aliyuncs.com/data/small_coco_demo/small_coco_demo.tar.gz && tar -zxf small_coco_demo.tar.gz
mkdir -p data/  && mv small_coco_demo data/coco

data/coco 格式如下：

data/coco/├── annotations│   ├── instances_train2017.json│   └── instances_val2017.json├── train2017│   ├── 000000005802.jpg│   ├── 000000060623.jpg│   ├── 000000086408.jpg│   ├── 000000118113.jpg│   ├── 000000184613.jpg│   ├── 000000193271.jpg│   ├── 000000222564.jpg│       ...│   └── 000000574769.jpg└── val2017    ├── 000000006818.jpg    ├── 000000017627.jpg    ├── 000000037777.jpg    ├── 000000087038.jpg    ├── 000000174482.jpg    ├── 000000181666.jpg    ├── 000000184791.jpg    ├── 000000252219.jpg         ...    └── 000000522713.jpg

三、模型训练和评估

以 vitdet-base 为示例。在 EasyCV 中，使用配置文件的形式来实现对模型参数、数据输入及增广方式、训练策略的配置，仅通过修改配置文件中的参数设置，就可以完成实验配置进行训练。可以直接下载示例配置文件。

查看 easycv 安装位置

# 查看easycv安装位置import easycvprint(easycv.__file__)
export PYTHONPATH=$PYTHONPATH:root/EasyCV

执行训练命令

单机8卡：CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7 python -m torch.distributed.launch --nproc_per_node=8 --master_port=29500 tools/train.py configs/detection/vitdet/vitdet_100e.py --work_dir easycv/vitdet --launcher pytorch --fp168机8卡：cp  EasyCV/tools/launch.py ./ && cp EasyCV/tools/train.py ./ &&python -m launch --nproc_per_node=8 train configs/detection/vitdet_dlc/vitdet_100e.py --work_dir easycv/vitdet_100e --launcher pytorch --fp16

执行评估命令

CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7 python -m torch.distributed.launch --nproc_per_node=8 --master_port=29500 tools/train.py configs/detection/vitdet/vitdet_100e.py --work_dir easycv/vitdet --launcher pytorch --fp16 --eval

Reference

模型细节来源：

1. ViT-based Mask RCNN https://arxiv.org/abs/2111.11429
   

2. ViTDet https://arxiv.org/abs/2203.16527
   

代码实现：

https://github.com/alibaba/EasyCV/blob/master/easycv/models/backbones/vitdet.py

https://github.com/tuofeilunhifi/EasyCV/blob/master/easycv/models/detection/vitdet/sfp.py

EasyCV 往期分享

• MAE自监督算法介绍和基于EasyCV的复现
  

• EasyCV开源｜开箱即用的视觉自监督+Transformer算法库