VOLO: Vision Outlooker for Visual Recognition 阅读笔记

该论文提出的方法目前在ImageNet分类任务中排名40+，但是在cityspaces-val语义分割任务中排名第5，说明该方法可以有效的应用于语义分割任务

该论文主要基于VIT进行进一步的研究和优化，VIT将transformer迁移到视觉任务的一个重要实践，在图像分类任务中，CNN一直占据主导优势地位，VIT的诞生以及后续的一系列的算法改进使得transformer逐渐在图像分类任务以及其他的一些任务中占据主导地位，反超CNN，成为了一个主流

• 首先作者发现限制 ViT 优于 CNN 的一个主要因素是它们在将精细特征和上下文编码为令牌表示时效率低下，这对于实现引人注目的视觉识别性能至关重要。 精细级别的信息可以通过更细粒度的图像标记化编码为标记，但这会导致更大长度的标记序列，从而二次增加 ViT 的自注意力机制的复杂性。

  • 由此提出了一种新的简单且轻量级的注意力机制，称为 Outlooker，以有效地利用精细信息丰富令牌表示。创新令牌聚合生成注意力的方式，并使模型能够有效地编码精细信息。 特别是，它直接通过有效的线性投影从锚标记特征中推断出聚合周围标记的机制，从而摆脱了昂贵的点积注意力计算。
  • 整体结构表达如下 ，可以看出来其实就是将原始的VIT中的self-attention替换成OutlookAttention
    

  方法提出的原因：1）每个空间位置的特征具有足够的代表性，可以生成注意力权重，用于局部聚合其相邻特征； 2）密集和局部空间聚合可以有效地编码精细信息。
  对于每个空间位置 (i, j)，outlook attention 计算其以 (i, j) 为中心的大小为 K × K 的局部窗口内的所有邻居的相似性。 与需要 Query-Key 矩阵乘法来计算注意力的 self-attention 不同，outlook attention 仅通过reshape操作来简化此过程。

  • Outlook attention：位置 (i, j) 的 Outlook 权重直接用作值聚合的注意力权重，然后是 Softmax 函数，表达式如下
    

  • Dense aggregation：Outlook attention 密集聚合映射值表示，将来自不同局部窗口的同一位置的不同加权值相加得到输出，表达式如下
    

  

  # Outlook attention 实现
  class OutlookAttention(nn.Module):
      """
      Implementation of outlook attention
      --dim: hidden dim
      --num_heads: number of heads
      --kernel_size: kernel size in each window for outlook attention
      return: token features after outlook attention
      """
  
      def __init__(self, dim, num_heads, kernel_size=3, padding=1, stride=1,
                   qkv_bias=False, qk_scale=None, attn_drop=0., proj_drop=0.):
          super().__init__()
          head_dim = dim // num_heads
          self.num_heads = num_heads
          self.kernel_size = kernel_size
          self.padding = padding
          self.stride = stride
          self.scale = qk_scale or head_dim**-0.5
  
          self.v = nn.Linear(dim, dim, bias=qkv_bias)
          self.attn = nn.Linear(dim, kernel_size**4 * num_heads)
  
          self.attn_drop = nn.Dropout(attn_drop)
          self.proj = nn.Linear(dim, dim)
          self.proj_drop = nn.Dropout(proj_drop)
  
          self.unfold = nn.Unfold(kernel_size=kernel_size, padding=padding, stride=stride)
          self.pool = nn.AvgPool2d(kernel_size=stride, stride=stride, ceil_mode=True)
  
      def forward(self, x):
          B, H, W, C = x.shape
  
          v = self.v(x).permute(0, 3, 1, 2)  # B, C, H, W
  
          h, w = math.ceil(H / self.stride), math.ceil(W / self.stride)
          v = self.unfold(v).reshape(B, self.num_heads, C // self.num_heads,
                                     self.kernel_size * self.kernel_size,
                                     h * w).permute(0, 1, 4, 3, 2)  # B,H,N,kxk,C/H
  
          attn = self.pool(x.permute(0, 3, 1, 2)).permute(0, 2, 3, 1)
          attn = self.attn(attn).reshape(
              B, h * w, self.num_heads, self.kernel_size * self.kernel_size,
              self.kernel_size * self.kernel_size).permute(0, 2, 1, 3, 4)  # B,H,N,kxk,kxk
          attn = attn * self.scale
          attn = attn.softmax(dim=-1)
          attn = self.attn_drop(attn)
  
          x = (attn @ v).permute(0, 1, 4, 3, 2).reshape(
              B, C * self.kernel_size * self.kernel_size, h * w)
          x = F.fold(x, output_size=(H, W), kernel_size=self.kernel_size,
                     padding=self.padding, stride=self.stride)
  
          x = self.proj(x.permute(0, 2, 3, 1))
          x = self.proj_drop(x)
  
          return x
  
  1
  2
  3
  4
  5
  6
  7
  8
  9
  10
  11
  12
  13
  14
  15
  16
  17
  18
  19
  20
  21
  22
  23
  24
  25
  26
  27
  28
  29
  30
  31
  32
  33
  34
  35
  36
  37
  38
  39
  40
  41
  42
  43
  44
  45
  46
  47
  48
  49
  50
  51
  52
  53
  54
  55
  56
  57
  58

• 优点：outlook attention 通过测量标记表示对之间的相似性来编码空间信息，这对于特征学习比卷积更具参数效率 ；outlook attention 采用滑动窗口机制对令牌表示进行局部编码，并在一定程度上保留了视觉任务的关键位置信息；注意力权重的生成方式简单高效。

• 为了利用精细级别的标记表示，在第一阶段，调整补丁嵌入模块，使图像在大小为 8 × 8 而不是 16 × 16 的小图像补丁上进行标记。一堆 Outlookers 用于生成精细级别的更具表现力的令牌表示。 在第二阶段，使用另一个补丁嵌入模块对令牌进行下采样。 然后采用一系列transformer对全局信息进行编码。Outlooker 和 Transformer 的比例保持在 1:3 左右。

实验部分可查看原文：VOLO: Vision Outlooker for Visual Recognition，源码：https://github.com/sail-sg/volo