當(dāng)可變形注意力機(jī)制引入Vision Transformer

【GiantPandaCV導(dǎo)語】通過在Transformer基礎(chǔ)上引入Deformable CNN中的可變性能力,在降低模型參數(shù)量的同時提升獲取大感受野的能力,文內(nèi)附代碼解讀。

引言

Transformer由于其更大的感受野能夠讓其擁有更強(qiáng)的模型表征能力,性能上超越了很多CNN的模型。

然而單純增大感受野也會帶來其他問題,比如說ViT中大量使用密集的注意力,會導(dǎo)致需要額外的內(nèi)存和計算代價,特征很容易被無關(guān)的部分所影響。

而PVT或者Swin Transformer中使用的sparse attention是數(shù)據(jù)不可知的,會影響模型對長距離依賴的建模能力。

由此引入主角:Deformabel Attention Transformer的兩個特點(diǎn):

  • data-dependent: key和value對的位置上是依賴于數(shù)據(jù)的。
  • 結(jié)合Deformable 方式能夠有效降低計算代價,提升計算效率。

下圖展示了motivation:

image

圖中比較了幾種方法的感受野,其中紅色星星和藍(lán)色星星表示的是不同的query。而實(shí)線包裹起來的目標(biāo)則是對應(yīng)的query參與處理的區(qū)域。

(a) ViT對所有的query都一樣,由于使用的是全局的注意力,所以感受野覆蓋全圖。

(b) Swin Transformer中則使用了基于window劃分的注意力。不同query處理的位置是在一個window內(nèi)部完成的。

(c) DCN使用的是3x3卷積核基礎(chǔ)上增加一個偏移量,9個位置都學(xué)習(xí)到偏差。

(d) DAT是本文提出的方法,由于結(jié)合ViT和DCN,所有query的響應(yīng)區(qū)域是相同的,但同時這些區(qū)域也學(xué)習(xí)了偏移量。

方法

先回憶一下Deformable Convolution:

image

簡單來講是使用了額外的一個分支回歸offset,然后將其加載到坐標(biāo)之上得到合適的目標(biāo)。

在回憶一下ViT中的Multi-head Self-attention:

\begin{aligned} q&=x W_{q}, k=x W_{k}, v=x W_{v}, \\ z^{(m)}&=\sigma\left(q^{(m)} k^{(m) \top} / \sqrtu0z1t8os\right) v^{(m)}, m=1, \ldots, M, \\ z&=\text { Concat }\left(z^{(1)}, \ldots, z^{(M)}\right) W_{o}, \\ z_{l}^{\prime} &=\operatorname{MHSA}\left(\operatorname{LN}\left(z_{l-1}\right)\right)+z_{l-1}, \\ z_{l} &=\operatorname{MLP}\left(\operatorname{LN}\left(z_{l}^{\prime}\right)\right)+z_{l}^{\prime}, \end{aligned}

有了以上鋪墊,下圖就是本文最核心的模塊Deformable Attention。

image
  • 左邊這部分使用一組均勻分布在feature map上的參照點(diǎn)
  • 然后通過offset network學(xué)習(xí)偏置的值,將offset施加于參照點(diǎn)中。
  • 在得到參照點(diǎn)以后使用bilinear pooling操作將很小一部分特征圖摳出來,作為k和v的輸入
x_sampled = F.grid_sample(
input=x.reshape(B * self.n_groups, self.n_group_channels, H, W), 
grid=pos[..., (1, 0)], # y, x -> x, y
mode='bilinear', align_corners=True) # B * g, Cg, Hg, Wg
  • 之后將得到的Q,K,V執(zhí)行普通的self-attention, 并在其基礎(chǔ)上增加relative position bias offsets。

其中offset network構(gòu)建很簡單, 代碼和圖示如下:

  self.conv_offset = nn.Sequential(
      nn.Conv2d(self.n_group_channels, self.n_group_channels, kk, stride, kk//2, groups=self.n_group_channels),
      LayerNormProxy(self.n_group_channels),
      nn.GELU(),
      nn.Conv2d(self.n_group_channels, 2, 1, 1, 0, bias=False)
  )
image

最終網(wǎng)絡(luò)結(jié)構(gòu)為:

image

具體參數(shù)如下:

image

實(shí)驗

實(shí)驗配置:300epoch,batch size 1024, lr=1e-3,數(shù)據(jù)增強(qiáng)大部分follow DEIT

  • 分類結(jié)果:
image

目標(biāo)檢測數(shù)據(jù)集結(jié)果:

image

語義分割:

image
  • 消融實(shí)驗:
image
  • 可視化結(jié)果:COCO
image

這個可視化結(jié)果有點(diǎn)意思,如果是分布在背景上的點(diǎn)大部分變動不是很大,即offset不是很明顯,但是目標(biāo)附近的點(diǎn)會存在一定的集中趨勢(ps:這種趨勢沒有Deformable Conv中的可視化結(jié)果明顯)

代碼

  • 生成Q
  B, C, H, W = x.size()
  dtype, device = x.dtype, x.device
  
  q = self.proj_q(x)
  • offset network前向傳播得到offset
  q_off = einops.rearrange(q, 'b (g c) h w -> (b g) c h w', g=self.n_groups, c=self.n_group_channels)
  offset = self.conv_offset(q_off) # B * g 2 Hg Wg
  Hk, Wk = offset.size(2), offset.size(3)
  n_sample = Hk * Wk
  • 在參照點(diǎn)基礎(chǔ)上使用offset
offset = einops.rearrange(offset, 'b p h w -> b h w p')
reference = self._get_ref_points(Hk, Wk, B, dtype, device)
    
if self.no_off:
    offset = offset.fill(0.0)
    
if self.offset_range_factor >= 0:
    pos = offset + reference
else:
    pos = (offset + reference).tanh()
  • 使用bilinear pooling的方式將對應(yīng)feature map摳出來,等待作為k,v的輸入。
x_sampled = F.grid_sample(
    input=x.reshape(B * self.n_groups, self.n_group_channels, H, W), 
    grid=pos[..., (1, 0)], # y, x -> x, y
    mode='bilinear', align_corners=True) # B * g, Cg, Hg, Wg
    
x_sampled = x_sampled.reshape(B, C, 1, n_sample)

q = q.reshape(B * self.n_heads, self.n_head_channels, H * W)
k = self.proj_k(x_sampled).reshape(B * self.n_heads, self.n_head_channels, n_sample)
v = self.proj_v(x_sampled).reshape(B * self.n_heads, self.n_head_channels, n_sample)
  • 在positional encodding部分引入相對位置的偏置:
  rpe_table = self.rpe_table
  rpe_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
  ) # B * g, h_g, HW, Ns
  
  attn_bias = attn_bias.reshape(B * self.n_heads, H * W, n_sample)
  
  attn = attn + attn_bias

參考

https://github.com/LeapLabTHU/DAT

https://arxiv.org/pdf/2201.00520.pdf

?著作權(quán)歸作者所有,轉(zhuǎn)載或內(nèi)容合作請聯(lián)系作者
【社區(qū)內(nèi)容提示】社區(qū)部分內(nèi)容疑似由AI輔助生成,瀏覽時請結(jié)合常識與多方信息審慎甄別。
平臺聲明:文章內(nèi)容(如有圖片或視頻亦包括在內(nèi))由作者上傳并發(fā)布,文章內(nèi)容僅代表作者本人觀點(diǎn),簡書系信息發(fā)布平臺,僅提供信息存儲服務(wù)。

相關(guān)閱讀更多精彩內(nèi)容

友情鏈接更多精彩內(nèi)容