关于机器学习:基于EasyCV复现ViTDet单层特征超越FPN

作者:伝迹、谦言、临在

欢送应用咱们最近开源的EasyCV,次要聚焦于最新的Vision Transformer模型,以及相干的上游CV工作

开源地址:​ ​https://github.com/alibaba/Ea…​​

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须要解决两个问题:

如何晋升计算效率?
如何失去多尺度特色?
ViT-based Mask R-CNN给出了初步的解决方案,ViTDet在此基础上,对如何失去多尺度特色做了进一步的改良。

如何晋升计算效率

ViT采纳的global self-attention和图像输出大小(HW)的平方成正比,对于检测模型,其输出分辨率往往较大,此时用ViT作为backbone在计算量和内存耗费上都是十分惊人的,比方输出尺寸为1024×1024,采纳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都是14×14,然而在输出分辨率为1024×1024的状况下,先通过一个patch_embed,就变成了64×64的分辨率,64是不能整除14的,

这里有两种解决形式:

  1. 在patch_embed之后加一个插值变成56×56,从ViT输入的时候再插值回64×64。
  2. 在patch_embed之后pad成70×70,复原成原图的时候裁剪成64×64。

两种都试了一下,发现第二种不会掉点,而第一种可能会导致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为16×16,那么最终就失去一种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,可能还有一点点小细节没有思考到的。

model ViTDet
base
cur E0
detials
box_AP 50.6
mask_AP 45.0
lr step
epoch 100
RunTime(hours) 50
bs(total=imgs/gpu x gpu_nums x cumulative_iters) 8node8bs=64
comments

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 easycv
print(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 --fp16
8机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
模型细节起源:

ViT-based Mask RCNN ​ ​https://arxiv.org/abs/2111.11429​​
ViTDet ​ ​https://arxiv.org/abs/2203.16527​​

代码实现:

​​ ​https://github.com/alibaba/Ea…​​

​ ​https://github.com/tuofeilunh…​​

EasyCV往期分享
MAE自监督算法介绍和基于EasyCV的复现
EasyCV开源|开箱即用的视觉自监督+Transformer算法库

评论

发表回复

您的邮箱地址不会被公开。 必填项已用 * 标注

这个站点使用 Akismet 来减少垃圾评论。了解你的评论数据如何被处理