深度學習之圖像分類（二十五）S2MLPv2 網絡詳解

目錄

• • 深度學習之圖像分類（二十五）S2MLPv2 網絡詳解
  • • 1. 前言
    • 2. S2MLPv2
    • • 2.1 S2MLPv2 Block
      • 2.2 Spatial-shift 與感受野反思
    • 3. 總結
    • 4. 代碼

經過 S2MLP 和 Vision Permutator 的沉淀，為此本節我們便來學習學習 S2MLPv2 的基本思想。

1. 前言

S2MLPv2 依是百度提出的用于視覺的空間位移 MLP 架構，其作者以及順序與 S2MLP 一模一樣，其論文為 S2-MLPv2: Improved Spatial-Shift MLP Architecture for Vision。S2MLPv2 的修改點主要在于三處：金字塔結構（參考 ViP）、分三類情況進行考慮（參考 ViP）、使用 Split Attention（參考 ViP 和 ResNeSt）。總結而言就是把 ViP 中的 Permute-MLP layer 中別人沿著不同方向進行交互替換為了 Spatial-shift 操作。在參數量基本一致的情況下，其性能優于 ViP。

2. S2MLPv2

2.1 S2MLPv2 Block

S2MLPv2 和 S2MLPv1 類似，整體網絡結構不做過多贅述，主要講解一下 S2MLPv2 Block 的細節（建議大家先回顧之前的章節 S2MLP 以及 ViP）：

• 首先是特征圖輸入后，對 Channel 進行一個全連接，這里是對于特定位置信息進行交流，其實也就是 $1\times 1$ 卷積，只不過這里將維度變為了原來的 3 倍。然后經過一個 GELU 激活函數。
• 將特征圖均分為 3 等分，分別用于后續三個 Spatial-shift 分支的輸入。
  • 第一個分支進行與 S2MLPv1 一樣的 Spatial-shift 操作，即右-左-下-上移動。
  • 第二個分支進行與第一個分支反對稱的 Spatial-shift 操作，即下-上-右-左移動。
  • 第三個分支保持不變
• 之后將三個分支的結果通過 Split Attention 結合起來。這樣不同位置的信息就被加到同一個通道上對齊了。
• 再經過一個 MLP 進行不同位置的信息整合，然后經過 LN 激活函數。（看了這么多網絡，其實激活函數在前在后都可以）

def spatial_shift1(x):b,w,h,c = x.size()x[:,1:,:,:c/4] = x[:,:w-1,:,:c/4]x[:,:w-1,:,c/4:c/2] = x[:,1:,:,c/4:c/2]x[:,:,1:,c/2:c*3/4] = x[:,:,:h-1,c/2:c*3/4]x[:,:,:h-1,3*c/4:] = x[:,:,1:,3*c/4:]return xdef spatial_shift2(x):b,w,h,c = x.size()x[:,:,1:,:c/4] = x[:,:,:h-1,:c/4]x[:,:,:h-1,c/4:c/2] = x[:,:,1:,c/4:c/2]x[:,1:,:,c/2:c*3/4] = x[:,:w-1,:,c/2:c*3/4]x[:,:w-1,:,3*c/4:] = x[:,1:,:,3*c/4:]return xclass S2-MLPv2(nn.Module):def __init__(self, channels):super().__init__()self.mlp1 = nn.Linear(channels,channels * 3)self.mlp2 = nn.Linear(channels,channels)self.split_attention = SplitAttention()def forward(self, x):b,w,h,c = x.size()x = self.mlp1(x)x1 = spatial_shift1(x[:,:,:,:c/3])x2 = spatial_shift2(x[:,:,:,c/3:c/3*2])x3 = x[:,:,:,c/3*2:]a = self.split_attention(x1,x2,x3)x = self.mlp2(a)return x

• 接下來的就是和 MLP-Mixer 中的 Channel-mixing MLP 一致。

2.2 Spatial-shift 與感受野反思

三組 Spatial-shift （包括恒等）與一組相比有什么進步和問題呢？

• 傳統計算機視覺感受野以及近期 ViP 工作等等，都提倡奇數和中心概念，即在某中心卷積核大小是奇數的，一左一右一上一下是對稱的。原始的一組 Spatial-shift 其實是一個菱形感受野且不包括中心。現在有恒等之后則是菱形感受野且包括中心了，這是一個進步。
• 但是第二組設計為與第一組反對稱的結構，但是這個反沒有反徹底。其實這三組 Spatial-shift 也可看作是人精心設計構造的。那么我們仔細看一下，其實沒有實現完全的互補。讓我們把目光放到 Split Attention 之后，輸出的特征圖其實也可被看作四個部分，每部分對應著：左上中相加，右下中相加，上左中相加，下右中相加。為了更好的方便大家理解這句話，我們不妨先忽略 Split Attention 給出的權重，并將經過 Spatial-shift 操作前的三部分特征圖分別記錄為 $f,g,h$，輸出記錄為 $z$。則有如下公式，其中下標表示不同的旋轉部分。
  • 如果從強迫癥的觀點看：第一組 Spatial-shift 是 右-左-下-上，則第二組 Spatial-shift 應該是 上-下-右-左 才對。
  • 如果從感受野完整性的觀點看：第一組 Spatial-shift 是 右-左-下-上，則第二組 Spatial-shift 應該是 左上-左下-右上-右下 才對。

$$\begin{gathered} z_1(x,y)=f_1(x?1,y)+g_1(x,y?1)+h_1(x,y) \\ {z}_2(x,y)=f_2(x+1,y)+g_2(x,y+1)+h_2(x,y) \\ {z}_3(x,y)=f_3(x,y?1)+g_3(x?1,y)+h_3(x,y) \\ {z}_4(x,y)=f_4(x,y+1)+g_4(x+1,y)+h_4(x,y) \end{gathered}$$

關于 Split 的消融實驗，作者分別移除了第二部分和第三部分，發現移除第二部分損失的性能還比第三部分（恒等）的多，但是就差 0.1%，這個消融實驗其實很難解釋三部分怎么相互作用的，至少從計算機視覺感受野的角度不太說得清楚。或許 MLP 結構就不太適合用感受野來分析吧…

3. 總結

相比于現有的 MLP 的結構，S2-MLP 的一個重要優勢是僅僅使用通道方向的全連接（$1\times 1$ 卷積）是可以作為 Backbone 的，期待該團隊后續的進展。S2-MLPv2 其實是通過 Spatial-shift 和 Split Attention 代替原有的 $N\times N$ 卷積，本質上并沒有延續 MLP-Mixer 架構中長距離依賴的思想。S2-MLPv2 中也并沒有長距離依賴的使用。S2-MLPv2 雖然性能提升了，但是還沒有開源，本身自己的貢獻點其實不太足，這樣做的理論性也不足。

延續我一貫的認識，如何在 MLP 架構中如何結合圖像局部性和長距離依賴依然是值得探討的點。

4. 代碼

代碼并沒有開源，非官發復現的代碼詳見 此處。

import torch from torch import nn from einops.layers.torch import Reduce from .utils import pairclass PreNormResidual(nn.Module):def __init__(self, dim, fn):super().__init__()self.fn = fnself.norm = nn.LayerNorm(dim)def forward(self, x):return self.fn(self.norm(x)) + xdef spatial_shift1(x):b,w,h,c = x.size()x[:,1:,:,:c//4] = x[:,:w-1,:,:c//4]x[:,:w-1,:,c//4:c//2] = x[:,1:,:,c//4:c//2]x[:,:,1:,c//2:c*3//4] = x[:,:,:h-1,c//2:c*3//4]x[:,:,:h-1,3*c//4:] = x[:,:,1:,3*c//4:]return xdef spatial_shift2(x):b,w,h,c = x.size()x[:,:,1:,:c//4] = x[:,:,:h-1,:c//4]x[:,:,:h-1,c//4:c//2] = x[:,:,1:,c//4:c//2]x[:,1:,:,c//2:c*3//4] = x[:,:w-1,:,c//2:c*3//4]x[:,:w-1,:,3*c//4:] = x[:,1:,:,3*c//4:]return xclass SplitAttention(nn.Module):def __init__(self, channel = 512, k = 3):super().__init__()self.channel = channelself.k = kself.mlp1 = nn.Linear(channel, channel, bias = False)self.gelu = nn.GELU()self.mlp2 = nn.Linear(channel, channel * k, bias = False)self.softmax = nn.Softmax(1)def forward(self,x_all):b, k, h, w, c = x_all.shapex_all = x_all.reshape(b, k, -1, c) #bs,k,n,ca = torch.sum(torch.sum(x_all, 1), 1) #bs,chat_a = self.mlp2(self.gelu(self.mlp1(a))) #bs,kchat_a = hat_a.reshape(b, self.k, c) #bs,k,cbar_a = self.softmax(hat_a) #bs,k,cattention = bar_a.unsqueeze(-2) # #bs,k,1,cout = attention * x_all # #bs,k,n,cout = torch.sum(out, 1).reshape(b, h, w, c)return outclass S2Attention(nn.Module):def __init__(self, channels=512):super().__init__()self.mlp1 = nn.Linear(channels, channels * 3)self.mlp2 = nn.Linear(channels, channels)self.split_attention = SplitAttention(channels)def forward(self, x):b, h, w, c = x.size()x = self.mlp1(x)x1 = spatial_shift1(x[:,:,:,:c])x2 = spatial_shift2(x[:,:,:,c:c*2])x3 = x[:,:,:,c*2:]x_all = torch.stack([x1, x2, x3], 1)a = self.split_attention(x_all)x = self.mlp2(a)return xclass S2Block(nn.Module):def __init__(self, d_model, depth, expansion_factor = 4, dropout = 0.):super().__init__()self.model = nn.Sequential(*[nn.Sequential(PreNormResidual(d_model, S2Attention(d_model)),PreNormResidual(d_model, nn.Sequential(nn.Linear(d_model, d_model * expansion_factor),nn.GELU(),nn.Dropout(dropout),nn.Linear(d_model * expansion_factor, d_model),nn.Dropout(dropout)))) for _ in range(depth)])def forward(self, x):x = x.permute(0, 2, 3, 1)x = self.model(x)x = x.permute(0, 3, 1, 2)return xclass S2MLPv2(nn.Module):def __init__(self,image_size=224,patch_size=[7, 2],in_channels=3,num_classes=1000,d_model=[192, 384],depth=[4, 14],expansion_factor = [3, 3],):image_size = pair(image_size)oldps = [1, 1]for ps in patch_size:ps = pair(ps)assert (image_size[0] % (ps[0] * oldps[0])) == 0, 'image must be divisible by patch size'assert (image_size[1] % (ps[1] * oldps[1])) == 0, 'image must be divisible by patch size'oldps[0] = oldps[0] * ps[0]oldps[1] = oldps[1] * ps[1]assert (len(patch_size) == len(depth) == len(d_model) == len(expansion_factor)), 'patch_size/depth/d_model/expansion_factor must be a list'super().__init__()self.stage = len(patch_size)self.stages = nn.Sequential(*[nn.Sequential(nn.Conv2d(in_channels if i == 0 else d_model[i - 1], d_model[i], kernel_size=patch_size[i], stride=patch_size[i]),S2Block(d_model[i], depth[i], expansion_factor[i], dropout = 0.)) for i in range(self.stage)])self.mlp_head = nn.Sequential(Reduce('b c h w -> b c', 'mean'),nn.Linear(d_model[-1], num_classes))def forward(self, x):embedding = self.stages(x)out = self.mlp_head(embedding)return out

總結

以上是生活随笔為你收集整理的深度学习之图像分类（二十五）-- S2MLPv2 网络详解的全部內容，希望文章能夠幫你解決所遇到的問題。

如果覺得生活随笔網站內容還不錯，歡迎將生活随笔推薦給好友。