狗都能看懂的Vision Transformer with Deformable Attention的讲解和代码实现

1、前言

在前面一篇博客介绍了可变形卷积，相比于普通的卷积，它自带的可学习偏移量使得模型能够关注感兴趣的区域，而不是固定的窗口。这个想法是在CNN中使用的，近些年来ViT的爆火，Self-Attention机制带来了更高的准确率。Transformer模型通过大接收场在视觉任务中展现了强大的表现力。但这种大接收场也带来了高昂的计算和内存成本，同时可能受到无关区域的干扰。随后Swin-Transformer的出现有效减少了计算量。不过其提出的MSA模块，虽然可以在一定程度上控制计算复杂度，但由于接收场的增长较慢，限制了对大物体建模的潜力。

原论文地址： https://arxiv.org/abs/2201.00520
官方开源代码：https://github.com/LeapLabTHU/DAT
Pytorch实现代码：https://github.com/Runist/DAT

2、Vision Transformer with Deformable Attention

作者参考Deformable Convolution提出了一种新颖的可变形自注意力模块（DMHA），这种模块的关键（key）和值（value）对的位置选择是数据相关的。通过这种灵活的机制，模型能够专注于重要的区域，从而捕捉到更多有用的特征，同时通过下采样，减少无关的计算。

文中会对Swin-Transformer和ViT进行比较，没有看过的同学建议先阅读一下。

3、模型结构

模型还是经典结构，由4个Stage组成，这里和ViT以及Swin-Transformer不同，每个Stage用到的Attention模块不同，后面会讲解原因。

• Patch Embedding用的是Conv + LayerNorm层，和ViT的token转换是一样的。这里就不多说了Vision Transformer的博客中有详细的讲解
• Stage 1 与 Stage 2先是一个W-MSA模块，再接另一个SW-MSA模块，这两个模块是成对出现的的，和Swin-Transformer是一摸一样。本文就不再展开讲解，请阅读Swin-Transformer的博客。
• Stage 3 和 Stage 4则是采用W-MSA + MDHA堆叠而成。这里MDHA就是DAT中提到Deformable attention module模块。
• 对于分类网络在代码中，还有LayerNorm、AvgPooling和一个全连接层组成，这个在图中没有体现。这个基本已经成了Transformer的不成文规矩了。

4、Deformable attention module详解

在Deformable Convolutional Networks中，每个特征图上的元素单独学习一个偏移量，空间复杂度为HWC，我们还需要考虑变形卷积的kernel大小，如果是3x3的卷积，空间复杂度需要再乘以3*3 = 9HWC。如果将这一机制直接应用于注意力模块，复杂度将急剧上升到$N_q{N}_{k}C$，其中$N_q$和$N_k$分别是query和key的数量，而这个数量同常是和特征图的尺寸一样，所以直接套用上去会造成4次方的复杂度。

Deformable DETR也应用了类似的方法，但他们通过设置较少的key（$N_k=4$）减少开销。但这只在检测头中表现良好，在主干网络中，因为key的数量过少而导致信息丢失问题比较严重。

本文提出的一个更简单的解决方案，为每个query共享key和value，如上图所示。

我们逐步讲解一下：

1. $x$为特征图，先生成Reference Points。图上只画了4个点（实际不止4个，为了简化）。
2. $x$重映射为$q$，通过Offset network生成偏移量，生成偏移量的数量和Reference Points一致。Offset network内可以通过控制downsample factor $r$，来控制生成的特征图大小。
3. Reference Points加上offsets得到中间的Features，为了减少计算量，我们会进行一次下采样。这里会根据通道数分成多个groups，类似Multi-Head Self-Attention的思想，增强不同groups之间特征多样性。
4. 生成对应的$v$和$k$，计算出attn值，得到output。

在实际代码中，步骤3与步骤4之间还有位置编码偏执，文章里只是简单提及了一下。源码中有具体的实现方法，简单来说是生成了一个位置编码表（可学习的），利用上图的Bilinear Interpolation进行下采样，得到和attn的shape一致的偏执。

看到这里，有的同学可能会问，这个offset的网络，如果在r=1的情况下，那岂不是计算量没有减少。并且本身ViT也是全局token之间互相计算attention，那这个Deformable attention module岂不是只增加了计算量？其实我刚开始看代码的时候也有这个疑问，但后面看了下生成的offset值，有些位置是重复的，有些位置是空的。也就意味着，模型会自己关注感兴趣的区域，重复的区域就代表权重大，位置为空，权重小。所以结合Deformable Convolution的思想，就有了Deformable multi-head attention(DMHA)。并且通过控制$r$，能够有效降低计算量，在高分辨率输入的任务比较有用。

具体代码如下：

    def forward(self, x):B, C, H, W = x.size()dtype, device = x.dtype, x.device# proj_q is weight_q, just convq = self.proj_q(x)      # [B, C, H, W] => [B, C, H, W]# 'b (g c) h w'：表示原始张量的维度。这里，(g c) 表示一个维度，它实际上是 g 和 c 这两个维度的乘积# (b g) c h w：将 b 和 g 维度合并成一个新的维度q_off = einops.rearrange(q, 'b (g c) h w -> (b g) c h w', g=self.n_groups, c=self.n_group_channels)     # [B, C, H, W] => [B * g, C / g, H, W]# Offset network in paperoffset = self.conv_offset(q_off)    # [B * g, C / g, H, W] => [B * g, 2, H, W]# H = Hk, W = WkHk, Wk = offset.size(2), offset.size(3)n_sample = Hk * Wkif self.offset_range_factor > 0:offset_range = torch.tensor([1.0 / Hk, 1.0 / Wk], device=device).reshape(1, 2, 1, 1)# tanh [-1, 1]# mul(offset_range) [-1 / Hk, 1 / Wk]# mul(offset_range) [-offset_range_factor / Hk, offset_range_factor / Wk]offset = offset.tanh().mul(offset_range).mul(self.offset_range_factor)offset = einops.rearrange(offset, 'b p h w -> b h w p')             # [B * g, 2, H, W] => [B * g, H, W, 2]reference = self._get_ref_points(Hk, Wk, B, dtype, device)          # [B * g, Hk, Wk, 2] if self.no_off:offset = offset.fill(0.0)if self.offset_range_factor >= 0:pos = offset + referenceelse:# To stabilize the training processpos = (offset + reference).tanh()# Bilinear Interpolation in paperx_sampled = F.grid_sample(input=x.reshape(B * self.n_groups, self.n_group_channels, H, W),  # [B, C, H, W] => [B * g, C / g, H, W]grid=pos[..., (1, 0)], # y, x -> x, ymode='bilinear', align_corners=True)x_sampled = x_sampled.reshape(B, C, 1, n_sample)    # [B * g, C / g, H, W] => [B, C, 1, H * W]q = q.reshape(B * self.n_heads, self.n_head_channels, H * W)                            # [B, C, 1, H * W] => [B * nh, C / nh, H * W]k = self.proj_k(x_sampled).reshape(B * self.n_heads, self.n_head_channels, n_sample)    # [B, C, 1, H * W] => [B * nh, C / nh, H * W]v = self.proj_v(x_sampled).reshape(B * self.n_heads, self.n_head_channels, n_sample)    # [B, C, 1, H * W] => [B * nh, C / nh, H * W]# m 和 n指代两个tensor的n_sampleattn = torch.einsum('b c m, b c n -> b m n', q, k) # [B * nh, H * W, H * W]attn = attn.mul(self.scale)if self.use_pe:if self.dwc_pe:# Depth-wise Convolutional Position Encodingresidual_lepe = self.rpe_table(q.reshape(B, C, H, W)).reshape(B * self.n_heads, self.n_head_channels, H * W)elif self.fixed_pe:# Fixed Position Encodingrpe_table = self.rpe_tableattn_bias = rpe_table[None, ...].expand(B, -1, -1, -1)attn = attn + attn_bias.reshape(B * self.n_heads, H * W, self.n_sample)else:# Relative Position Biasrpe_table = self.rpe_tablerpe_bias = rpe_table[None, ...].expand(B, -1, -1, -1)q_grid = self._get_ref_points(H, W, B, dtype, device)displacement = (q_grid.reshape(B * self.n_groups, H * W, 2).unsqueeze(2) - pos.reshape(B * self.n_groups, n_sample, 2).unsqueeze(1)).mul(0.5)attn_bias = F.grid_sample(input=rpe_bias.reshape(B * self.n_groups, self.n_group_heads, 2 * H - 1, 2 * W - 1),grid=displacement[..., (1, 0)],mode='bilinear', align_corners=True)attn_bias = attn_bias.reshape(B * self.n_heads, H * W, n_sample)attn = attn + attn_biasattn = F.softmax(attn, dim=2)attn = self.attn_drop(attn)# m 和 n都是n_sampleout = torch.einsum('b m n, b c n -> b c m', attn, v)    # [B * nh, C / nh, H * W]if self.use_pe and self.dwc_pe:out = out + residual_lepeout = out.reshape(B, C, H, W)        # [B * nh, C / nh, H * W] =>  [B, C, H, W]y = self.proj_drop(self.proj_out(out))return y, pos.reshape(B, self.n_groups, Hk, Wk, 2), reference.reshape(B, self.n_groups, Hk, Wk, 2)

5、Deformable attention module计算量

DMHA相比于Swin-Transformer和PVT只多了一个offset networks的计算。文章给出了复杂度的计算公式：
$$\Omega (DMHA)=2HW{N}_{S}C+2HW{C}^2+2N_s{C}^2+(k^2+2)N_{s}C$$

• H代表feature map的高度
• W代表feature map的宽度
• C代表feature map的通道数
• 其中$N_S=H_G{W}_G=HW/r^2$
• k为DWConv的卷积核数量

其中前三项是Self-attention的固定计算开销，最后一项是offset network的开销。

先看一下最后一项，比较简单。由于DWConv的处理的是下采样之后的Feature map，所以$N_{s}C$对应的其实是$HW/r^2\ast C$，$r=1$时两者相等，k对应的是DWConv的卷积核数量，这没什么好解释。+2对应的是1x1的卷积和LayerNorm层。

再来从通用Transformer推导前三项，首先看一下Self-Attention的公式：
$$Attention(Q,K,V)=SoftMax(\frac{Q{K}^T}{\sqrt{d}})V$$
对于feature map的每个像素（或称为token，patch），都要通过$W_q$，$W_k$，$W_v$生成对应qkv。这里假设q，k，v的向量长度与feature map的channel数量C保持一致。那么对应所有像素生成Q的过程如下：
$$X^{HW\times C}\cdot W_q^{C\times C}=Q^{HW\times C}$$

• $X^{HW\times C}$为所有token拼接一起得到的矩阵（一共有HW个像素，每个像素的深度为C）
• $W_q^{C\times C}$为生成query的变换矩阵
• $Q^{HW\times C}$为所有像素通过$W_q^{C\times C}$得到的query拼接后的矩阵

根据矩阵运算的计算量公式可以得到生成$Q$的计算量为$HW\times C\times C$，生成K和V的过程一样，同理都是$HW{C}^2$，那么总共是$3HW{C}^2$（但这里注意一下，生成q是不用下采样的，只有k和v才需要，所以对应的是公式中的$2N_s{C}^2$）。接下来$Q$和$K^T$相乘，对应计算量为$(HW{)}^{2}C$，由于有$r$的存在，所以实际是$HW{N}_{S}C$：
$$Q^{HW\times C}\cdot K^{T(C\times HW)}=X^{HW\times HW}$$
这里忽略除以$\sqrt{d}$以及softmax的计算量，假设得到$A^{HW\times HW}$，最后还要乘以$V$，这里对应的计算量是$(HW{)}^{2}C$，由于有$r$的存在，所以实际是$HW{N}_{S}C$：
$$A^{hw\times hw}\cdot V^{hw\times C)}=X^{hw\times C}$$
那么对应单头的Self-Attention模块，总共需要$3HW{C}^2+(HW{)}^{2}C+(HW{)}^{2}C=3HW{C}^2+2(HW{)}^{2}C$。而在实际使用过程中，使用的是多头的Multi-head Self-Attention模块，在之前的文章中有进行过实验对比，多头注意力模块相比单头注意力模块的计算量仅多了最后一个融合矩阵$W_O$的计算量$HW{C}^2$。

所以总共加起来是：$4HW{C}^2+2(HW{)}^{2}C$

由于下采样因子$r$的存在，DMHA的开销为：$HW{N}_{S}C+HW{N}_{S}C+2N_s{C}^2$ = $2HW{N}_{S}C+2N_s{C}^2$，再加上一个生成q的计算量$HW{C}^2$以及多头注意力模块的最后一个融合矩阵$W_O$的计算量$HW{C}^2$。
所以总共是$2HW{N}_{S}C+2N_s{C}^2+HW{C}^2+HW{C}^2=2HW{N}_{S}C+2HW{C}^2+2N_s{C}^2$

现在把$r=1$代入到$N_S=H_G{W}_G=HW/r^2$，就和$4HW{C}^2+2(HW{)}^{2}C$一样。

6、模型详细配置参数

下图（表1）是原论文中给出的关于不同DAT的配置，T(Tiny)，S(Small)，B(Base)，其中：

• win. sz. 7x7表示使用的窗口（Windows）的大小
• N表示堆叠的次数
• head表示多头注意力模块中head的个数
• groups表示DMHA中的分组注意力机制的数量

需要提一下，为什么只在最后两个Stage中用到Deformable attention module，作者在第三章节说了，是为了实现模型容量和计算负担之间的权衡。但在作者最新的代码中，模型结构已经不是这样了。且不再使用W-MSA模块，SW-MSA模块。笔者根据其以往提交的代码，重新修改后，已和原文对应。