Swin Transformer的繼任者（下）
點(diǎn)藍(lán)色字關(guān)注“機(jī)器學(xué)習(xí)算法工程師”
設(shè)為星標(biāo)，干貨直達(dá)！
GG-Transformer
上海交大提出的GG Transformer其主要思路是改變window的劃分方式，window不再局限于一個(gè)local region，而是來自全局。這里提出的一個(gè)操作是AdaptivelyDilatedSplitting，即window的token是通過以一定的adaptive dilation rate 來采樣獲得，下面是一個(gè)實(shí)例（2x2個(gè)windows）：
如果這樣劃分window，那么window attention將具有全局視野，但是相鄰的patchs之間缺乏交互，所以GG Transformer又增加了一個(gè)額外的Gaze分支：先將attention中的values進(jìn)行Merging操作，其實(shí)就是AdaptivelyDilatedSplitting的逆變換，那么將得到正常的tokens排列，然后通過一個(gè)depth-wise conv來提取局部信息，再通過AdaptivelyDilatedSplitting操作得到和attention一樣的windows，再加上attention后的特征即可：
論文里將這種結(jié)構(gòu)分成Glance和Gaze兩個(gè)分支，分別用來提取全局和局部信息，類比人類的Glance and Gaze行為。這里的AdaptivelyDilatedSplitting其實(shí)可以通過前面說的shuffle操作來實(shí)現(xiàn)，后面要講的Shuffle Transformer也是一樣的原理。論文中也沒有提到位置編碼，估計(jì)Gaze分支的卷積可以隱式地編碼位置信息。
在ImageNet上，GG-Transformer在同樣的參數(shù)和算力下，其模型效果要優(yōu)于Swin模型：
在COCO數(shù)據(jù)集上，基于Mask R-CNN，其模型效果也要優(yōu)于Swin：
Shuffle Transformer
騰訊提出的Shuffle Transformer其核心思路是通過spatial shuffle來建立cross-window之間聯(lián)系。這里的spatial shuffle和ShuffleNet中的channel shuffle類似，通過spatial shuffle可以將來自不同windows的token組成新的window：
這個(gè)實(shí)現(xiàn)上應(yīng)該是和AdaptivelyDilatedSplitting等價(jià)的，另外MSG Transfomer也是通過MSG tokens的channel shuffle來建立不同windows間的聯(lián)系。它們的實(shí)現(xiàn)都是類似的：reshape->transpose->reshape。開源代碼也給出了具體實(shí)現(xiàn)：
if self.shuffle:
q, k, v = rearrange(qkv, 'b (qkv h d) (ws1 hh) (ws2 ww) -> qkv (b hh ww) h (ws1 ws2) d', h=self.num_heads, qkv=3, ws1=self.ws, ws2=self.ws)
# 這里其實(shí)是三種操作
# reshape: qkv = qkv.reshape(b, 3, h, d, ws1, hh, ws2, ww)
# transpose：qkv = qkv.transpose(1, 0, 5, 7, 2, 4, 6, 3)
# reshape: q, k, v = qkv.reshape(3, b*hh*ww, h, ws1*ws2, d)
else:
q, k, v = rearrange(qkv, 'b (qkv h d) (hh ws1) (ww ws2) -> qkv (b hh ww) h (ws1 ws2) d', h=self.num_heads, qkv=3, ws1=self.ws, ws2=self.ws)
# 注意正常window split與shuffle版本的區(qū)別，第一步reshape有區(qū)別
與Swin Transformer模型類似，Shuffle Transformer交替地采用標(biāo)準(zhǔn)的WMSA和shuffle SWMSA：
可以看到，Shuffle Transformer在WMSA操作后增加了一個(gè)NWC操作，這個(gè)其實(shí)是一個(gè)depthwise conv，其kernel size和window size一樣，用于增強(qiáng)Neighbor-Window Connection。
class Block(nn.Module):
def __init__(self, dim, out_dim, num_heads, window_size=1, shuffle=False, mlp_ratio=4., qkv_bias=False, qk_scale=None, drop=0., attn_drop=0.,
drop_path=0., act_layer=nn.ReLU6, norm_layer=nn.BatchNorm2d, stride=False, relative_pos_embedding=False):
super().__init__()
self.norm1 = norm_layer(dim)
self.attn = Attention(
dim, num_heads=num_heads, window_size=window_size, shuffle=shuffle, qkv_bias=qkv_bias, qk_scale=qk_scale,
attn_drop=attn_drop, proj_drop=drop, relative_pos_embedding=relative_pos_embedding)
# NWC
self.local = nn.Conv2d(dim, dim, window_size, 1, window_size//2, groups=dim, bias=qkv_bias)
self.drop_path = DropPath(drop_path) if drop_path > 0. else nn.Identity()
self.norm2 = norm_layer(dim)
mlp_hidden_dim = int(dim * mlp_ratio)
self.mlp = Mlp(in_features=dim, hidden_features=mlp_hidden_dim, out_features=out_dim, act_layer=act_layer, drop=drop, stride=stride)
self.norm3 = norm_layer(dim)
print("input dim={}, output dim={}, stride={}, expand={}, num_heads={}".format(dim, out_dim, stride, shuffle, num_heads))
def forward(self, x):
x = x + self.drop_path(self.attn(self.norm1(x)))
x = x + self.local(self.norm2(x)) # local connection
x = x + self.drop_path(self.mlp(self.norm3(x)))
return x
從結(jié)構(gòu)上看，Shuffle Transformer幾乎和Swin Transformer一樣。在ImageNet數(shù)據(jù)集上，同等條件上Shuffle Transformer相比Swin提升明顯：
在COCO數(shù)據(jù)集上，基于Mask R-CNN，Shuffle Transformer和Swin性能不相上下：
后話
可以看到，這四個(gè)模型和Swin Transformer本質(zhì)上都是一種local attention，只不過它們從不同地方式來增強(qiáng)local attention的全局建模能力。而且，在相似的參數(shù)和計(jì)算量的條件下，5種模型在分類任務(wù)和dense任務(wù)上表現(xiàn)都是類似的。近期，微軟在論文Demystifying Local Vision Transformer: Sparse Connectivity, Weight Sharing, and Dynamic Weight上系統(tǒng)地總結(jié)了Local Vision Transformer的三大特性：
Sparse connectivity：每個(gè)token的輸出只依賴于其所在local window上tokens，而且各個(gè)channel之間是無聯(lián)系的；（這里忽略了attention中query，key和valude的linear projections，那么attention就其實(shí)可以看成在計(jì)算好的權(quán)重下對(duì)tokens的特征進(jìn)行加權(quán)求和，而且是channel-wise的） Weight sharing：權(quán)重對(duì)于各個(gè)channel是共享的； Dynamic weight：權(quán)重不是固定的，而是基于各個(gè)tokens動(dòng)態(tài)生成的。
那么local attention就和Depth-Wise Convolution就很相似，首先后者也具有Sparse connectivity：只在kernel size范圍內(nèi)，而且各個(gè)channel之間無連接。而Depth-Wise Convolution也具有weight sharing，但是卷積核是在所有的空間位置上共享的，但不同channle采用不同的卷積核。另外depth-wise convolution的卷積核是訓(xùn)練參數(shù)，一旦完成訓(xùn)練就是固定的，而不是固定的。另外local attention丟失了位置信息，需要位置編碼，但是depth-wise convolution不需要。下圖是不同操作的區(qū)別：
論文中也設(shè)計(jì)了基于depth-wise convolution的模型，和Swin模型結(jié)構(gòu)類似：
在ImageNet數(shù)據(jù)集上，DW-Conv模型效果和Swin模型相當(dāng)（這里D-DW-Conv增加了動(dòng)態(tài)權(quán)重的特性，類似SE模塊來動(dòng)態(tài)生成kernel weights）：
從這項(xiàng)研究來看，設(shè)計(jì)好的Conv模型在性能上也是可以和local attention模型匹敵的，也許local attention模型反而退化到了CNN模型。一點(diǎn)體外話是之前的CNN模型一般常采用3x3和1x1比較小的卷積核，但是這里采用7x7的卷積核反而大幅度提升模型效果（相比ResNet50），這里也值得深思。
參考
Shuffle Transformer: Rethinking Spatial Shuffle for Vision Transformer Twins: Revisiting the Design of Spatial Attention in Vision Transformers Glance-and-Gaze Vision Transformer MSG-Transformer: Exchanging Local Spatial Information by Manipulating Messenger Tokens Pyramid Vision Transformer: A Versatile Backbone for Dense Prediction without Convolutions Swin Transformer: Hierarchical Vision Transformer using Shifted Windows Demystifying Local Vision Transformer: Sparse Connectivity, Weight Sharing, and Dynamic Weight.
推薦閱讀
谷歌AI用30億數(shù)據(jù)訓(xùn)練了一個(gè)20億參數(shù)Vision Transformer模型，在ImageNet上達(dá)到新的SOTA！
"未來"的經(jīng)典之作ViT：transformer is all you need！
PVT：可用于密集任務(wù)backbone的金字塔視覺transformer！
漲點(diǎn)神器FixRes：兩次超越ImageNet數(shù)據(jù)集上的SOTA
不妨試試MoCo，來替換ImageNet上pretrain模型！
機(jī)器學(xué)習(xí)算法工程師
一個(gè)用心的公眾號(hào)
