【原】登上更高峰！颜水成、程明明团队开源ViP，引入三维信息编码机制，无需卷积与注意力

极市平台 2021-06-24

展开全文

作者丨Happy

审稿丨邓富城

编辑丨极市平台

极市导读

本文从位置信息编码出发，引入了高-宽-通道三维信息编码机制。为进一步校正不同分支的作用，提出了加权融合方式。ViP在ImageNet上取得了83.2%的top1精度，代码已开源。 >>加入极市CV技术交流群，走在计算机视觉的最前沿

paper: https:///abs/2106.12368

code: https://github.com/Andrew-Qibin/VisionPermutator

本文是颜水成与程明明团队在MLP架构方面的一次探索，从位置信息编码出发，引入了高-宽-通道三维信息编码机制；为进一步校正不同分支的作用，提出了加权融合方式(即注意力机制)。该文的思路与一作之前的“TripletAttention”非常相似，区别在于前者作用于CNN，而后者作用于MLP架构。Whatever，ViP将类MLP架构的性能进一步向上推了一把，使其具有与CNN、Transformer相当的竞争力。

Abstract

本文提出一种概念简单、数据高效的类MLP架构Vision Permutator(ViP)用于视觉识别任务。通过对2D特征表达所携带的位置信息重要性的认知，ViP采用线性投影方式沿高与宽维度编码特征表达。这使的ViP能够沿单一空间维度捕获长距离依赖关系，同时沿另一个方向保持精确的位置信息，然后通过相互补充聚合方式产生位置敏感输出，进而形成关于目标区域的强有力表征。

作者通过实验表明：ViP具有与CNN、Transformer相当的竞争力。无需空域卷积或者注意力机制，无需额外大尺度训练数据，仅需25M可学习参数，ViP在ImageNet上取得了81.5%的top1精度，这比大部分CNN与Transformer都要优秀。当把模型参数提升到88M，模型精度可以进一步提升到83.2%。作者期望该工作能促进社区重新思考空间信息的编码并辅助类MLP方法的设计。

Method

上图给出了本文所提ViP整体架构示意图，ViP以的图像作为输入，并将其均匀的拆分为图像块()序列。所有图像块将采用共享线性层将输入映射为线性嵌入(也称之为词)；接下来，我们将所有词送入到Permutators序列以编码空间、通道信息。所得词最后沿空间维度进行全局池化，后接线性分类层进行类别预测。

Permutator

在前面Figure1的左上角，我们可以看到所提Permutator模块的示意图。我们可以看到：除了LayerNorms与跳过连接外，Permutator包含两个成分(1) Permute-MLP用于编码空间信息（2）Channel-MLP用于编码通道信息。其中，Channel-MLP具有与ViT中的前馈层类似的架构，包含两个全连接层并内置GELU激活；对于Permute-MLP，本文提出沿着高与宽维度分别处理词。从数学公式来看，给定C维词，Permutator的公式定义如下：

Permute-MLP 上图给出了Permute-MLP的结构示意图，它包含三个分支，每个分支用于编码不同的信息：高、宽、通道。通道信息的编码比较简单，我们仅需全连接层进行线性投影；重点是如何编码高与宽两个维度的空间信息。

假设隐层维度C维384，输入图像分辨率。为沿着高维度编码空间信息，我们首先进行高-通道维度置换。给定输入，我们首先将其沿通道维度拆分为S块，满足。此时，块尺寸为，N=16，。然后，我们执行一次高-通道置换操作得到；接下来，我们采用全连接层对期进行高信息混合。为复原原始维度信息，我们需要再执行一次高-通道置换操作并输出。类似的，对于第二个分支，我们执行与上述类似的宽-通道置换操作并生成。最后，我们将所有词表达相加送入到新的全连接层中得到Permute-MLP的输出，描述如下：

Weighted Channel-MLP

在上面的公式中，本文采用了简单的加法进行三分支融合。这里，本文通过重校正不同分支的重要性改进Permute-MLP并提出了Weighted Permute-MLP。直接看code吧，如下所示：

class WeightedPermuteMLP(nn.Module):
    def __init__(self, dim, segment_dim=8, qkv_bias=False, qk_scale=None, attn_drop=0., proj_drop=0.):
        super().__init__()
        self.segment_dim = segment_dim
        self.mlp_c = nn.Linear(dim, dim, bias=qkv_bias)
        self.mlp_h = nn.Linear(dim, dim, bias=qkv_bias)
        self.mlp_w = nn.Linear(dim, dim, bias=qkv_bias)
        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
        S = C // self.segment_dim
        h = x.reshape(B, H, W, self.segment_dim, S).permute(0, 3, 2, 1, 4).reshape(B, self.segment_dim, W, H*S)
        h = self.mlp_h(h).reshape(B, self.segment_dim, W, H, S).permute(0, 3, 2, 1, 4).reshape(B, H, W, C)

        w = x.reshape(B, H, W, self.segment_dim, S).permute(0, 1, 3, 2, 4).reshape(B, H, self.segment_dim, W*S)
        w = self.mlp_w(w).reshape(B, H, self.segment_dim, W, S).permute(0, 1, 3, 2, 4).reshape(B, H, W, C)

        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
    
class Mlp(nn.Module):
    def __init__(self, in_features, hidden_features=None, out_features=None, act_layer=nn.GELU, drop=0.):
        super().__init__()
        out_features = out_features or in_features
        hidden_features = hidden_features or in_features
        self.fc1 = nn.Linear(in_features, hidden_features)
        self.act = act_layer()
        self.fc2 = nn.Linear(hidden_features, out_features)
        self.drop = nn.Dropout(drop)

    def forward(self, x):
        x = self.fc1(x)
        x = self.act(x)
        x = self.drop(x)
        x = self.fc2(x)
        x = self.drop(x)
        return x