Visoin MLP之CycleMLP A MLP-like Architecture for Dense Prediction
Visoin MLP之CycleMLP A MLP-like Architecture for Dense Prediction
從摘要讀文章
This paper presents a simple MLP-like architecture, CycleMLP, which is a versatile backbone for visual recognition and dense predictions, unlike modern MLP architectures, e.g., MLP-Mixer, ResMLP, and gMLP, whose architectures are correlated to image size and thus are infeasible in object detection and segmentation.
看來又是一個金字塔形狀的 MLP 架構，並且可以看到主要的工作必定又是圍繞空間 MLP 而展開的。因為這裡解除了架構對於輸入尺寸的依賴。
實際上在MLP-Mixer 原論文中，作者們其實也嘗試了金字塔結構，相較與固定解析度的形式，確實收斂更快。We tried using the token-mixing MLP to reduce the number of tokens by mapping from S input tokens to S'<S output tokens. While first experiments showed that on JFT-300M such models significantly reduced training time without losing much performance, we were unable to transfer these findings to ImageNet or ImageNet-21k.
CycleMLP has two advantages compared to modern approaches.
- (1) It can cope with various image sizes.
- (2) It achieves linear computational complexity to image size by using local windows.
這裡使用區域性視窗的操作實現了線性計算複雜度，這中固定區域性視窗也使得可以處理不同大小的影象。
In contrast, previous MLPs have quadratic computations because of their fully spatial connections.
指出了改進空間 MLP 的必要性。其計算複雜度太高。
We build a family of models that surpass existing MLPs and achieve a comparable accuracy (83.2%) on ImageNet-1K classification compared to the state-of-the-art Transformer such as Swin Transformer (83.3%) but using fewer parameters and FLOPs.
可以看到，本文的方法的效能是非常高的。不知道除了結構之外，是否使用了其他的策略。
We expand the MLP-like models' applicability, making them a versatile backbone for dense prediction tasks. CycleMLP aims to provide a competitive baseline on object detection, instance segmentation, and semantic segmentation for MLP models. In particular, CycleMLP achieves 45.1 mIoU on ADE20K val, comparable to Swin (45.2 mIOU). Code is available at this https URL.
主要內容
整體結構
可以看到，核心是在改進空間 MLP。
現有 MLP 結構的三點不足：
- 大多屬於單尺度結構，不便於遷移到其他的需要特徵金字塔的任務上，例如檢測、分割等。
- 空間 MLP 會連線輸入特徵空間上的所有點，這也限制了模型對於輸入尺寸的依賴。不利於多尺度訓練、多尺度測試、甚至訓練和測試解析度不一樣的情況。
- 空間 MLP 具有二次方量級的計算複雜度，使其不便於處理高解析度影象。
對此，作者們從兩個方面進行應對：
- 針對問題 1，設計了層級結構來生成特徵金字塔。
- 針對問題 2 和 3，實際上是針對空間 MLP 的問題，作者們設計了一種特殊的通道 MLP 來實現對於區域性空間的處理。由於是針對區域性空間的，所以說不再對於輸入尺寸有過強的依賴。並且仍然是通道 MLP（空間上共享的點操作），所以計算複雜度降低到了線性。
與 S2-MLP 的不同
雖然都是使用通道 MLP 替換了空間 MLP，但是具體方式和模型整體形式有所不同：
- S2-MLP 對特徵進行通道分組，不同組進行空間上不同方向的相對偏移。這在特徵圖上引入了額外的分組和偏移操作。而本文則是不需要改變特徵，僅僅是調整了通道 MLP 的運算形式。具有更良好的通用性和可拔插性。
- S2-MLP 依然是單尺度結構，本文引入了金字塔結構來更好的適應檢測分割等任務。
實際上，對於第 1 點，其實 S2-MLP 中給出了使用深度分離卷積實現的策略，即偏移可以通過特定形式的深度分離卷積核實現，對於輸入資料的分組和偏移都是可以直接通過對卷積核的操作來實現的。這裡的第 1 點並不成立。只能說，Cycle FC 的實現可能更加直接一些，不同於 S2-MLP 需要藉助一些與計算無關的處理操作。
更進一步，從實現上來講，Cycle FC 是否也是可以通過使用深度分離卷積實現便宜操作呢？回答是可以的，我會在後文的程式碼分析中提供一些簡單的嘗試。
模型細節
- patch embedding module: window size 7, stride 4,最終特徵下采樣四倍。
- 中間通過跨步卷積實現 2 倍下采樣。
- 最終輸出為下采樣 32 倍的特徵。
- 最終使用一個全連線層整合所有 token 的特徵。
核心操作——Cycle FC
論文提出 Cycle FC 的核心想法在於利用通道 MLP 的與特徵尺寸的無關性（減少了對輸入形狀的限制並且可以將計算複雜度降到線性），同時想辦法增大其感受野來更好的整合上下文特徵。
從圖中給出的形式可以看到，Cycle FC 實際上是一種在通道上進行特定位置的偏移（階梯狀取樣，stair-like style）的通道 MLP。所以對於輸入的形狀要求不會太嚴苛。當然，至少偏移位置不能超出 HW 上限定的核尺寸。
從程式碼中可以看到，這裡是限定了一個範圍，通過讓通道索引對其取模從而實現限定範圍內的迴圈偏移，這裡的實現很有意思，用到了可變形卷積來對核引數應用偏移。
具體而言，原始的通道 MLP 的計算方式為：\(Y_{i,j} = \sum_{c=0}^{C_i - 1} \mathcal{F}^{\top}_{j,c} \cdot X_{i,c}\)，其中的 \(\mathcal{F} \in \mathbb{R}^{C_i \times C_o}\) 表示通道 MLP 的可學習權重。其中的 \(i\&j\) 分別表示空間和通道的索引。
而對於本文提出的 Cycle FC，計算方式擴充套件成為：\(Y_{i, j} = \sum_{c=0}^{C_i-1} \mathcal{F}^{\top}_{j,c} \cdot X_{i+c\%S_{\mathcal{P}},c}\)，這引入了一個偏移範圍引數 \(S_{\mathcal{P}}\)，即 pseudo-kernel size，表示通道偏移後所有涉及到的計算位置在 HW 空間上的投影的面積，而另一個引數 \(i\) 表示偏移的起始位置。在程式碼中取值為 pseudo-kernel 矩形區域內部的中心相對座標（以區域內部左上角為起點，索引為 0，對矩形區域按照行主序進行排序索引），即start_idx=(self.kernel_size[0]*self.kernel_size[1])//2。
程式碼解析
這裡通過註釋的形式對部分核心程式碼分析。
from torchvision.ops.deform_conv import deform_conv2d as deform_conv2d_tv
class CycleFC(nn.Module):
def __init__(
self,
in_channels: int,
out_channels: int,
kernel_size, # re-defined kernel_size, represent the spatial area of staircase FC
stride: int = 1,
padding: int = 0,
dilation: int = 1,
groups: int = 1,
bias: bool = True,
):
"""
這裡的kernel_size實際使用的時候時3x1或者1x3
"""
super(CycleFC, self).__init__()
if in_channels % groups != 0:
raise ValueError('in_channels must be divisible by groups')
if out_channels % groups != 0:
raise ValueError('out_channels must be divisible by groups')
if stride != 1:
raise ValueError('stride must be 1')
if padding != 0:
raise ValueError('padding must be 0')
self.in_channels = in_channels
self.out_channels = out_channels
self.kernel_size = kernel_size
self.stride = _pair(stride)
self.padding = _pair(padding)
self.dilation = _pair(dilation)
self.groups = groups
# 被偏移調整的1x1卷積的權重，由於後面使用torchvision提供的可變形卷積的函式，所以權重需要自己構造
self.weight = nn.Parameter(torch.empty(out_channels, in_channels // groups, 1, 1))
# kernel size == 1
if bias:
self.bias = nn.Parameter(torch.empty(out_channels))
else:
self.register_parameter('bias', None)
# 要注意，這裡是在註冊一個buffer，是一個常量，不可學習，但是可以儲存到模型權重中。
self.register_buffer('offset', self.gen_offset())
def gen_offset(self):
"""
生成卷積核偏移量的核心操作。
要想理解這一函式的操作，需要首先理解後面使用的deform_conv2d_tv的具體用法。
具體可見：https://pytorch.org/vision/0.10/ops.html#torchvision.ops.deform_conv2d
這裡對於offset引數的要求是：
offset (Tensor[batch_size,
2 * offset_groups * kernel_height * kernel_width,
out_height,
out_width])
– offsets to be applied for each position in the convolution kernel.
也就是說，對於樣本s的輸出特徵圖的通道c中的位置(x,y)，這個函式會從offset中取出，形狀為
kernel_height*kernel_width的卷積核所對應的偏移引數為
offset[s, 0:2*offset_groups*kernel_height*kernel_width, x, y]
也就是這一系列引數都是對應樣本s的單個位置(x,y)的。
針對不同的位置可以有不同的offset，也可以有相同的（下面的實現就是後者）。
對於這2*offset_groups*kernel_height*kernel_width個數，涉及到對於輸入特徵通道的分組。
將其分成offset_groups組，每份單獨擁有一組對應於卷積核中心位置的相對偏移量，
共2*kernel_height*kernel_width個數。
對於每個核引數，使用兩個量來描述偏移，即h方向和w方向相對中心位置的偏移，
即下面程式碼中的減去kernel_height//2或者kernel_width//2。
需要注意的是，當偏移位置位於padding後的tensor邊界外，則是將網格使用0補齊。
如果網格上有邊界值，則使用邊界值和用0補齊的網格頂點來計算雙線性插值的結果。
"""
offset = torch.empty(1, self.in_channels*2, 1, 1)
start_idx = (self.kernel_size[0] * self.kernel_size[1]) // 2
assert self.kernel_size[0] == 1 or self.kernel_size[1] == 1, self.kernel_size
for i in range(self.in_channels):
if self.kernel_size[0] == 1:
offset[0, 2 * i + 0, 0, 0] = 0
# 這裡計算了一個相對偏移位置。
# deform_conv2d使用的以對應輸出位置為中心的偏移座標索引方式
offset[0, 2 * i + 1, 0, 0] = (
(i + start_idx) % self.kernel_size[1] - (self.kernel_size[1] // 2)
)
else:
offset[0, 2 * i + 0, 0, 0] = (
(i + start_idx) % self.kernel_size[0] - (self.kernel_size[0] // 2)
)
offset[0, 2 * i + 1, 0, 0] = 0
return offset
def forward(self, input: Tensor) -> Tensor:
"""
Args:
input (Tensor[batch_size, in_channels, in_height, in_width]): input tensor
"""
B, C, H, W = input.size()
return deform_conv2d_tv(input,
self.offset.expand(B, -1, H, W),
self.weight,
self.bias,
stride=self.stride,
padding=self.padding,
dilation=self.dilation)
由於這裡關於 offset 的生成有些讓人疑惑的地方，我也給作者提了 issue（https://github.com/ShoufaChen/CycleMLP/issues/11），同時附了一個關於 deform_conv2d 的小例子。
這裡作者給提供了一個與程式碼更貼合的圖示：
class CycleMLP(nn.Module):
def __init__(self, dim, qkv_bias=False, qk_scale=None, attn_drop=0., proj_drop=0.):
super().__init__()
self.mlp_c = nn.Linear(dim, dim, bias=qkv_bias)
self.sfc_h = CycleFC(dim, dim, (1, 3), 1, 0)
self.sfc_w = CycleFC(dim, dim, (3, 1), 1, 0)
self.reweight = Mlp(dim, dim // 4, dim * 3)
self.proj = nn.Linear(dim, dim)
self.proj_drop = nn.Dropout(proj_drop)
def forward(self, x):
B, H, W, C = x.shape
h = self.sfc_h(x.permute(0, 3, 1, 2)).permute(0, 2, 3, 1)
w = self.sfc_w(x.permute(0, 3, 1, 2)).permute(0, 2, 3, 1)
c = self.mlp_c(x)
a = (h + w + c).permute(0, 3, 1, 2).flatten(2).mean(2)
a = self.reweight(a).reshape(B, C, 3).permute(2, 0, 1).softmax(dim=0).unsqueeze(2).unsqueeze(2)
x = h * a[0] + w * a[1] + c * a[2]
x = self.proj(x)
x = self.proj_drop(x)
return x
從程式碼中可以看到，在實際使用的時候，是基於類似於 Inception V3 中分形卷積的形式，構建了 1x3 和 3x1 的兩組並行操作。另外也有一個普通的通道 MLP 來進行單個位置的處理。從而構建了一個三分支結構。
which is inspired by the factorization of convolution [47] and criss-cross attention [26].
實驗效果
和同期 MLP 方法的比較
這裡作者提到了 GFNet，它使用了 FFT 來學習空間特徵，並且計算量也更少，與 CycleMLP 也有著相近的效能。但是其受輸入解析度的制約，如果想要更改輸入解析度則需要使用引數插值。這可能會損害密集預測任務的效能。（這真的會有很大影響麼？）
另外，作者們也補充了關於不同測試解析度和不同分支的影響的消融實驗。
這兩個實驗展現出了一些有趣的現象。
- 首先看解析度的影響， 可以看到，最優的測試粉筆那縷可能並不是與原始訓練一致。結果中反映出了 CycleFC 對於測試尺寸的穩定性。
- 但是這裡要注意的是，這裡的實驗從程式碼中可以看到，測試中使用的是先放大到指定尺寸的 $$\frac{256}{224}$$ 倍後再從中心 crop 出對應的尺寸的操作方式。而其他形式的資料處理的對應的現象仍需進行更多的實驗驗證。
- 表 4 中關於多分支結構中的三個分支進行了消融實驗。可以看到，操作多樣性對於多分支結構是具有正向增益的。
也通過在檢測和分割任務上的表現展現出了提出結構（更靈活的輸入尺寸、更高效的空間計算）的有效性。