Press "Enter" to skip to content

Swin Transformer V2:通向视觉大模型之路

本站内容均来自兴趣收集,如不慎侵害的您的相关权益,请留言告知,我们将尽快删除.谢谢.

码字不易,欢迎点赞!

 

继Swin Transformer之后,微软在去年11月份发布了Swin Transformer V2,目前模型的实现以及预训练模型已经开源。Swin Transformer V2的核心是 将模型扩展到更大的容量和分辨率 ,其中最大的模型 SwinV2-G参数量达到了30亿 ,在物体检测任务上 图像分辨率达到1536×1536 , 基于SwinV2-G的模型也在4个任务上达到了SOTA :在图像分类数据集ImageNet V2上达到了84.0%的top1准确度,在物体检测数据集COCO上达到了63.1/54.4的box/mask mAP,在语义分割数据集上ADE20K达到了59.9 mIoU,在视频分类数据集Kinetics-400上达到了86.8%,不过目前,只有COCO和ADE20K两个数据集上还保持SOTA。当将模型进行扩展时,作者发现面临两个主要的困难:一是 大模型的训练变得不稳定 ,二是 预训练模型增大窗口大小会出现性能下降 。作者通过相应的优化策略来重点解决这两个问题,本文将结合论文和源码来介绍这些优化方法。

 

Post norm

 

将模型变大变深后往往出现训练难题,作者发现将Swin模型扩展到大模型后,激活值随着层数的加深大幅增长,如下图所示,其中最大值与最小值比例可高达 ,对于参数量为658M的的Swin-H模型(左图最上面红色曲线),这种差异使得模型无法完成训练(图右蓝色线,训练loss出现上升)。

 

作者发现可以调整LayerNorm的位置来解决这个问题,vision transfomer默认采用的pre-norm的方式,即在Attention和FFN之前先进行LayerNorm,而 SwinV2采用post-norm:将LayerNorm调整到Attention和FFN之后 。采用post-norm,Attention或者FFN的输出经过LayerNorm之后才和shortcut相加,这样就尽量减少网络加深后的激活值积累,论文发现采用post-norm后,不同层的激活值确实也变得更平稳了,如上图所示。对于最大的模型Swin-G,还每6个transformer blocks之间再增加一个LayerNorm来进一步稳定训练和激活值。

 

关于深度ViT的训练,其实meta AI也在CaiT中提出了一个简单的策略 LayerScale ,如下图d所示,FFN和SA之后每个channel乘以一个训练的系数,这个系数往往初始化值较小比如1e-4,这样看来其实LayerScale和post-norm一样也是来约束FFN和SA的输出,防止层数加深后出现激活值的增加。

 

另外一点是,原始的Transformer模型采用的也是post-norm,不过和SwinV2里说的post-norm不一样,它是残差和shortcut相加后才执行LayerNorm,如下图a所示。论文 On Layer Normalization in the Transformer Architecture 也研究了pre-norm和post-norm在NLP任务上表现,发现同等条件下,pre-norm相比post-norm更容易训练,而post-norm需要一些额外的策略如warmup来达到较优的效果。不过也有论文如 Understanding the Difficulty of Training Transformers 指出post-norm往往可以比pre-norm获得更好的效果,关于这一点可以见博客 为什幺Pre Norm的效果不如Post Norm? 一个简单的解释是pre-norm由于LayerNorm在残差结构里面,并没有破坏短路连接,所以更容易训练,但是同时也可能将网络退化到一个“宽网络”,降低网络的表达力。

 

虽然post-norm或许效果较好,但是训练比较困难,特别是层数较深时,最近微软也提出了一个简单的优化策略来 训练1000层的Transformer模型 DeepNet,这个策略就是将shortcut乘以一个大于1的常量(依赖模型架构),再和残差部分相加后进行LN。

 

为了防止和原始Transformer的post-norm相混淆,论文的第2版将提出的post-norm重命名为 res-post-norm(residual post normalization) ,也增加了基于Swin-S和Swin-B两个模型的对比实验,相比来看,还是res-post-norm效果更好一点。不过这里没有对比更大的模型。

 

Scaled cosine attention

 

原始的self-attention采用scaled dot product来计算两个token的相似度:计算两个token特征的点积,然后除以特征维度大小的平方根。作者发现对于大模型,其学习到的attention map往往被较少的一部分token所主导,特别是采用post-norm后。为了解决这个问题,论文提出scaled cosine attention,即计算两个token的余弦相似度,然后除以一个可学习的标量 :

 

 

其中 是relative position bias,而 是一个训练参数,不同层和attention head不共享,并且限制大于0.01。具体的实现如下所示:

 

# 初始化\tau
self.logit_scale = nn.Parameter(torch.log(10 * torch.ones((num_heads, 1, 1))), requires_grad=True)
# cosine attention
attn = (F.normalize(q, dim=-1) @ F.normalize(k, dim=-1).transpose(-2, -1))
logit_scale = torch.clamp(self.logit_scale, max=torch.log(torch.tensor(1. / 0.01))).exp()
attn = attn * logit_scale

 

由于余弦相似度范围在0~1,所以它使得不同的token间的相似度分布在同样的尺度下,从而减少attention map的过于集中的问题。下表为不同的模型使用scaled cosine attention后的对比实验,可以看到,在使用res-post-norm之后再使用scaled cosine attention可以进一步提升模型的分类准确度。

 

Log-spaced continuous position bias

 

分类模型的图像分辨率低,但是对于下游任务如检测和分割往往要采用更大的分辨率。由于Swin采用window attention和relative position bias,当采用分类预训练模型迁移到下游任务时如果要采用更大的window size,那幺就需要对relative position bias进行插值。大的视觉模型往往同时需要提升图像的分辨率,那幺此时也最好同时增大window来提升感受野。作者发现直接插值的话往往效果会下降,如下表所示,采用8×8的window和256×256分辨率的Swin-T模型在ImageNet1K上能达到81.7,但如果将这个预训练模型在12×12的window和384×384分辨率下,效果只有79.4,但是finetune之后能达到82.7%。论文提出了一种新的策略 log-spaced continuous position bias (记为 Log-CPB )来解决这个问题,如下表所示,基于 Log-CPB 的模型直接在同样的场景下迁移效果能达到82.4%,超过原来的81.8%,而且finetune之后可以达到83.2%,在COCO和ADE20k数据集上也表现更好。

 

在介绍Log-CPB,我们先来回顾一下Swin所采用的relative position bias,它是在计算attention时加上一个bias项,如下所示:

 

 

这里的 就是relative position bias,它用来编码window中的各个tokens间的相对位置, 为window的大小,总共有 个tokens。而由于图像是2D结构,token在每个维度上的相对位置分布在 范围内,总共有 个值,那幺通过定义一个训练的bias矩阵 ,就能通过索引的方式得到 ,这样处理的好处是参数量比较小。当改变window大小时,就需要对 进行插值。而SwinT-V2的改进策略时 采用连续的relative position bias ,不再定义一个固定大小的 ,而是采用一个小网络来预测relative position bias,这里:

 

 

这里的G是一个包含2层的MLP模型,中间采用ReLU激活,采用网络G的好处是它可以生成任意相对位置的relative position bias,这样在迁移到更大的window时就不需要任何更改。 虽然网络G可以适应不同大小的window,但是当window size变化时,相对位置范围也发生了变化,这意味着网络G要接受和预训练模型不一样的输入范围,为了尽量减少输入范围的变化,论文进一步提出采用log空间下的坐标来替换原来的线性坐标,两者的转换公式如下所示:

 

 

采用对数坐标的好处是可以减少范围的变化,这对网络G的泛化要求就变低了。比如要将8×8大小的window迁移到16×16大小,对于线性坐标,其范围从原来的[-7, 7]到[-15, 15],此时相当于额外增加了(15 – 7)/7=1.14倍原来的坐标范围。如果采用对数坐标,其范围从原来的[-2.079, 2.079]到[-2.773, 2.773],此时只额外增加了(2.773 – 2.079)/2.079=0.33倍原来的范围。从上表的对比结果来看,Log-CPB效果也要优于Linear-CPB。Log-CPB的具体实现代码如下所示(和论文描述有部分细节不同):

 

# 网络G来生成连续的relative position bias
self.cpb_mlp = nn.Sequential(nn.Linear(2, 512, bias=True),
                             nn.ReLU(inplace=True),
                             nn.Linear(512, num_heads, bias=False))
# 获取relative_coords_table: (2M-1)x(2M-1)作为网络G的输入
relative_coords_h = torch.arange(-(self.window_size[0] - 1), self.window_size[0], dtype=torch.float32)
relative_coords_w = torch.arange(-(self.window_size[1] - 1), self.window_size[1], dtype=torch.float32)
relative_coords_table = torch.stack(
            torch.meshgrid([relative_coords_h,
                            relative_coords_w])).permute(1, 2, 0).contiguous().unsqueeze(0)  # 1, 2*Wh-1, 2*Ww-1, 2
# 归一化,迁移预训练模型时用原来的window size
if pretrained_window_size[0] > 0:
    relative_coords_table[:, :, :, 0] /= (pretrained_window_size[0] - 1)
    relative_coords_table[:, :, :, 1] /= (pretrained_window_size[1] - 1)
else:
    relative_coords_table[:, :, :, 0] /= (self.window_size[0] - 1)
    relative_coords_table[:, :, :, 1] /= (self.window_size[1] - 1)
relative_coords_table *= 8  # 归一化到 -8, 8
relative_coords_table = torch.sign(relative_coords_table) * torch.log2(
            torch.abs(relative_coords_table) + 1.0) / np.log2(8) # 这里也同时除以np.log2(8)
self.register_buffer("relative_coords_table", relative_coords_table)
# 得到相对位置索引
coords_h = torch.arange(self.window_size[0])
coords_w = torch.arange(self.window_size[1])
coords = torch.stack(torch.meshgrid([coords_h, coords_w]))  # 2, Wh, Ww
coords_flatten = torch.flatten(coords, 1)  # 2, Wh*Ww
relative_coords = coords_flatten[:, :, None] - coords_flatten[:, None, :]  # 2, Wh*Ww, Wh*Ww
relative_coords = relative_coords.permute(1, 2, 0).contiguous()  # Wh*Ww, Wh*Ww, 2
relative_coords[:, :, 0] += self.window_size[0] - 1  # shift to start from 0
relative_coords[:, :, 1] += self.window_size[1] - 1
relative_coords[:, :, 0] *= 2 * self.window_size[1] - 1
relative_position_index = relative_coords.sum(-1)  # Wh*Ww, Wh*Ww
self.register_buffer("relative_position_index", relative_position_index)
# 预测时先用网络G产生\hat B,然后根据相对位置索引生成B
relative_position_bias_table = self.cpb_mlp(self.relative_coords_table).view(-1, self.num_heads)
relative_position_bias = relative_position_bias_table[self.relative_position_index.view(-1)].view(
            self.window_size[0] * self.window_size[1], self.window_size[0] * self.window_size[1], -1)  # Wh*Ww,Wh*Ww,nH
relative_position_bias = relative_position_bias.permute(2, 0, 1).contiguous()  # nH, Wh*Ww, Wh*Ww
relative_position_bias = 16 * torch.sigmoid(relative_position_bias) # 限制bias的大小范围?
attn = attn + relative_position_bias.unsqueeze(0)

 

论文还采用SwinV2-S和SwinV2-B对Log-CPB做了对比实验,可以看到使用Log-CPB后迁移到更大的window和更大的分辨率下,效果更好。

 

除此之外,论文还对不同方法得到的relative position bias矩阵(应该是relative position table,即 )进行了可视化,可以看到采用连续的relative position bias更光滑。

 

其它策略

 

对于大模型来说,显存消耗是一个大的挑战,论文也采用了一些优化策略来降低显存消耗,主要包括以下方法:

ZeRO优化器 :ZeRO是将模型参数和优化器状态参数均分到不同的GPU上,这样就可以降低显存使用,这里采用DeepSpeed库的ZeRO-1,它可以降低显存消耗但是几乎不影响训练速度。
Activation check-pointing :模型的中间特征也会消耗大量的GPU,可以采用 activation check-pointing 来缓解这个问题,目前PyTorch已经支持,具体见 https:// pytorch.org/docs/stable /checkpoint.html 。这个优化虽然降低了显存消耗,但是也将训练速度降低了30%。
Sequential self-attention computation :不再对self-attention进行batch运算,而是序列地逐个执行,从而降低显存使用,不过这个优化只在前两个stages使用,对训练速度只有少量的影响。

通过上述优化,可以用40GB的A100在COCO任务(1536×1536大小)或Kinetics-400任务(320x320x8大小)上训练最大的模型Swin-G。

 

大模型往往也需要更大的数据集来进行训练 ,作者将ImageNet-22K数据集扩展了5倍达到了70M,并 基于自监督学习方法 SimMIM 来进行预训练 。

 

模型设置

 

SwinV2和SwinV1一样,也首先包含4个不同大小的模型:

SwinV2-T :C = 96,block = {2, 2, 6, 2}
SwinV2-S/B/L :C=96/128/192,block={2, 2, 18, 2}

除了这个4个模型外,还包括2个更大的模型: SwinV2-H 和 SwinV2-G ,它们的参数量分别达到了658M和3B:

SwinV2-H : C = 352,block = {2, 2, 18, 2}
SwinV2-G : C = 512,block = {2, 2, 42, 4}

对于最大的两个模型,如前面所述,为了提升训练的稳定性,每6个transformer层间增加一个LayerNorm层。

 

实验结果

 

由于是研究大模型,所以论文的主要实验结果是基于最大的模型SwinV2-G,这里共对比了4个不同的任务: 图像分类ImageNet-1K (V1 and V2)数据集,物体检测COCO数据集,语义分割ADE20K数据集,以及视频理解Kinetics-400数据集 。

 

SwinV2-G采用2阶段的预训练策略来进行预训练 ,首先基于自监督方法SimMIM在70M的ImageNet-22K-ext数据集上训练20个epochs,然后在这个数据集上基于有监督再训练20个epochs。为了减少训练成本,采用192×192的分辨率进行预训练。

 

在 图像分类 任务上,SwinV2-G在ImageNet-1K上达到了90.17%,这也是继谷歌之后第一个超过90%的模型,而在ImageNet-1K-V2上达到了84.0%,是当时的SOTA,目前最好的结果是谷歌的Model soups。

 

在 物体检测 任务上,基于SwinV2-G的HTC++达到了63.1/54.4的box/mask mAP,对于实例分割目前是SOTA,但是检测效果低于DINO,这里的分辨率采用1536×1536,而窗口大小采用32×32。

 

在 语义分割 上,基于UperNet的分割模型在AD E20K上达到了59.9 mIoU。

 

在视频分类上,基于Video-SwinV2-G的模型分类准确度达到了86.8%。

 

小结

 

SwinV1是让我们看到Vision Transformer在下游分割和检测任务上应用的可能,而SwinV2让我们看到了Vision Transformer扩展到大模型的应用可能。

 

基于detectron2的版本见: GitHub – xiaohu2015/nndet2

 

参考

Swin Transformer V2: Scaling Up Capacity and Resolution
https:// github.com/microsoft/Sw in-Transformer
Swin Transformer: Hierarchical Vision Transformer using Shifted Windows
DeepNet: Scaling Transformers to 1,000 Layers
On Layer Normalization in the Transformer Architecture
Going deeper with Image Transformers

Be First to Comment

发表评论

您的电子邮箱地址不会被公开。