【导读】随着ResMLP、MLP-Mixer等文章的提出,基于MLP的backbone重新回到了CV领域。在图像识别方面,基于MLP的结构具有较少的inductive bias,但是依旧能够达到与CNN和Vision Transformer相当的性能。那么,MLP在其它下游视觉任务的效果如何呢?自6,7月份以来,MLP正式进军下游视觉任务,在检测与分割领域纷纷推出了最强MLP架构,本文我们将对近期在检测与分割领域最新MLP架构进行梳理总结,主要包括:上科大&腾讯优图开源AS-MLP, 香港大学&商汤科技提出的CycleMLP, 百度提出的目前最强视觉MLP架构S2-MLP(V1-V2).
1
上科大&腾讯优图开源AS-MLP:一种轴向位移的MLP框架器
本文是上海科技大学和腾讯优图在MLP架构方面的探索,不同于MLP-Mixer通过矩阵转置+词混叠MLP进行全局空域特征编码,ASMLP在局部特征通信方向投入了更多的关注。主要设计了一种轴向移位操作以便于进行不同方向的空间信息流交互。首先通过在水平和垂直方向上空间移动特征,轴向位移可以将不同空间位置的特征排列在相同的位置。然后,使用 MLP 来组合这些功能,简单而有效。这种方法使模型能够获得更多的局部依赖,从而提高性能。我们还可以设计AS-MLP的感受野尺寸以及扩张因子。ASMLP凭借如此简单而有效的架构取得了优于其他MLP架构的性能,同时达到了与Transformer架构(比如Swin Transformer)相当的性能,并且只需要更少的FLOPs。比如,AS-MLP 在 ImageNet-1K 数据集中仅使用 88M 参数和 15.2 GFLOP 就获得了 83.3% Top-1 准确率,且无需额外训练数据。在架构方面,AS-MLP采用了类似PVT的分层架构,因为可以轻易的迁移到下游任务。所提方法在ImageNet数据集上取得了优于其他MLP架构的性能,AS-MLP也是首个用于下游任务(如目标检测、语义分割)的MLP架构。AS-MLP在COCO验证集上取得了51.5mAP指标,在ADE20K数据集上取得了49.5mIoU指标,具有与Transformer架构相当的性能。
class AxialShift(nn.Module):
r""" Axial shift
Args:
dim (int): Number of input channels.
shift_size (int): shift size .
as_bias (bool, optional): If True, add a learnable bias to as mlp. Default: True
proj_drop (float, optional): Dropout ratio of output. Default: 0.0
"""
def __init__(self, dim, shift_size, as_bias=True, proj_drop=0.):
super().__init__()
self.dim = dim
self.shift_size = shift_size
self.pad = shift_size // 2
self.conv1 = nn.Conv2d(dim, dim, 1, 1, 0, groups=1, bias=as_bias)
self.conv2_1 = nn.Conv2d(dim, dim, 1, 1, 0, groups=1, bias=as_bias)
self.conv2_2 = nn.Conv2d(dim, dim, 1, 1, 0, groups=1, bias=as_bias)
self.conv3 = nn.Conv2d(dim, dim, 1, 1, 0, groups=1, bias=as_bias)
self.actn = nn.GELU()
self.norm1 = MyNorm(dim)
self.norm2 = MyNorm(dim)
self.shift_dim2 = Shift(self.shift_size, 2)
self.shift_dim3 = Shift(self.shift_size, 3)
def forward(self, x):
"""
Args:
x: input features with shape of (num_windows*B, N, C)
mask: (0/-inf) mask with shape of (num_windows, Wh*Ww, Wh*Ww) or None
"""
B_, C, H, W = x.shape
x = self.conv1(x)
x = self.norm1(x)
x = self.actn(x)
'''
x = F.pad(x, (self.pad, self.pad, self.pad, self.pad) , "constant", 0)
xs = torch.chunk(x, self.shift_size, 1)
def shift(dim):
x_shift = [ torch.roll(x_c, shift, dim) for x_c, shift in zip(xs, range(-self.pad, self.pad+1))]
x_cat = torch.cat(x_shift, 1)
x_cat = torch.narrow(x_cat, 2, self.pad, H)
x_cat = torch.narrow(x_cat, 3, self.pad, W)
return x_cat
x_shift_lr = shift(3)
x_shift_td = shift(2)
'''
x_shift_lr = self.shift_dim3(x)
x_shift_td = self.shift_dim2(x)
x_lr = self.conv2_1(x_shift_lr)
x_td = self.conv2_2(x_shift_td)
x_lr = self.actn(x_lr)
x_td = self.actn(x_td)
x = x_lr + x_td
x = self.norm2(x)
x = self.conv3(x)
return x
def extra_repr(self) -> str:
return f'dim={self.dim}, shift_size={self.shift_size}'
def flops(self, N):
# calculate flops for 1 window with token length of N
flops = 0
# conv1
flops += N * self.dim * self.dim
# norm 1
flops += N * self.dim
# conv2_1 conv2_2
flops += N * self.dim * self.dim * 2
# x_lr + x_td
flops += N * self.dim
# norm2
flops += N * self.dim
# norm3
flops += N * self.dim * self.dim
return
2
CycleMLP:密集预测任务的通用主干,已开源
本文是香港大学&商汤科技在MLP架构方面的探索,主要提出了一个简单的MLP架构 CycleMLP,它是视觉识别和密集预测任务的通用主干网络。CycleMLP 旨在为 MLP 模型的目标检测、实例分割和语义分割提供一个有竞争力的基线。与其它MLP架构相比CycleMLP主要有两个优势:
CycleMLP的主要贡献点在于:
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,
):
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
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)
self.register_buffer('offset', self.gen_offset())
self.reset_parameters()
def reset_parameters(self) -> None:
init.kaiming_uniform_(self.weight, a=math.sqrt(5))
if self.bias is not None:
fan_in, _ = init._calculate_fan_in_and_fan_out(self.weight)
bound = 1 / math.sqrt(fan_in)
init.uniform_(self.bias, -bound, bound)
def gen_offset(self):
"""
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.
"""
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
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)
def extra_repr(self) -> str:
s = self.__class__.__name__ + '('
s += '{in_channels}'
s += ', {out_channels}'
s += ', kernel_size={kernel_size}'
s += ', stride={stride}'
s += ', padding={padding}' if self.padding != (0, 0) else ''
s += ', dilation={dilation}' if self.dilation != (1, 1) else ''
s += ', groups={groups}' if self.groups != 1 else ''
s += ', bias=False' if self.bias is None else ''
s += ')'
return s.format(**self.__dict__)
为了对不同的Patch进行交互,除了常规的channel-mixing MLP,还引入了额外的token-mixing(也就是将特征图转置下,再接MLP)。这些方法为了达到优异的性能,需要在大型数据集上预训练。为此我们分析了其中的token-mixing,发现其等价于一个拥有全局感受野,且空间特异性的Depthwise卷积,但这两个特性会使得token-mixing MLP容易过拟合。为此,我们提出了一种新颖的纯 MLP 架构:Spatial-shift MLP结构,我们抛弃了token-mixing的过程,仅包含channel-mixing MLP,同时引入了空间移位操作来增强各个Patch的通信,它具有局部感受野,并且是spatial-agnostic。
Spatial shift operation:首先我们给定一个输入X,其形状为W, H, C然后将该输入在通道维度上进行分组,这里我们只移动四个方向,因此分为四组,然后每一组为W, H, C/4。接着是对每组输入进行不同方向的移位操作,以第一组为例子,我们在W维度上移一格,第二组在W操作反着移一格。同理另外两组在H维度上进行相同操作。一段伪代码如下所示:
def spatial_shift(x):
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 x
spatial-shift MLP(S2-MLP)采用了空间移位操作,因此达到了比ResMLP、MLP-Mixer更好的性能。近期,采用了更小的patch和金字塔结构的Vision Permutator (ViP)和Global Filter Network (GFNet) 在性能上又一次超越了S2-MLP。
因此,在本文中,我们改进了空间移位 MLP(S2-MLP)并提出了 S2-MLPv2 模型。相比于S2-MLP,S2-MLPv2的改动主要有以下方面:
下图是将扩展后的特征图分割成三部分。它单独移动每个分割部分,然后通过split-attention融合分割特征图。
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 x
def 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 x
class 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
总结:与基于 Transformer 的方法相比,我们的 S2-MLPv2 模型在没有自注意力和更少参数的情况下实现了极具竞争力的效果。与 MLP-mixer、ResMLP 等基于 MLP 的开创性工作以及同期工作(包括 Vision Permutator 和 GFNet)相比,空间位移 MLP 的另一个重要优势是空间位移 MLP的可以保持输入图像大小形状不变。因此,S2MLPV2可以基于特定比例的输入图像得到相应的预训练权重,这样便可以很好地用于具有各种尺寸大小的输入图像的下游任务中。未来的工作将致力于不断提高空间位移 MLP 架构的图像识别精度。一个未来的方向是尝试更小尺寸的patch和更先进的四级金字塔,如 CycleMLP 和 AS-MLP,以进一步减少 FLOPs 并缩短基于 Transformer 的模型之间的识别差距。(目前百度这项工作没有完全开源,估计S^2MLP(V3)已经在路上了!)