VIT讲解

AN IMAGE IS WORTH 16X16 WORDS: TRANSFORMERS FOR IMAGE RECOGNITION AT SCALE

会议：Computer Vision and Pattern Recognition         

时间：2020-10-22

Attention is all you need 一经发表，Transformer架构问世，如今Transformer架构已成为自然语言处理任务的事实标准，但其在计算机视觉中的应用仍然有限。在视觉领域，注意力机制通常与卷积网络结合使用，或者用来替代卷积网络的某些组件，但保持其整体结构。在本文中我们展示了对CNN的这种依赖是不必要的，直接应用于图像块序列的纯Transformer可以在图像分类任务中表现出色。

从图中可以看到VIT模型的总体预览，以下分别讨论其重点部分：

Patch Embedding：对于图像数据而言，其数据格式为[C, H, W]是三维矩阵,Transformer想要的输入形状为[L, D]是二维的，其中L是序列长度，D为每个token的维度。所以需要先通过一个Embedding层来对数据做个变换。如上图所示，首先将一张图片按给定大小分成一堆Patches，将输入图片(224x224)按照16x16大小的Patch进行划分，划分后会得到 N=(224/16)^2=196个Patches。每一个patch视为一个token，故patches数量就相当于序列长度。接着通过线性映射将每个Patch映射到一维向量中。在代码实现中，直接通过一个卷积层来实现。直接使用一个卷积核大小为16x16，步距为16，卷积核个数为768的卷积来实现，得到无重复重叠的Patches，即[3, 224, 224] -> [768, 14, 14]，然后把H以及W两个维度展平后再改变形状得到[768 , 14, 14] -> [768, 196]->[196, 768]，此时变成了一个二维矩阵[N, D]，正是Transformer想要的。（ViT-B/16为例）

class PatchEmbed(nn.Module):
    """
    2D Image to Patch Embedding
    """
    def __init__(self, img_size=224, patch_size=16, in_c=3, embed_dim=768, norm_layer=None):
        super().__init__()
        img_size = (img_size, img_size)
        patch_size = (patch_size, patch_size)
        self.img_size = img_size
        self.patch_size = patch_size
        self.grid_size = (img_size[0] // patch_size[0], img_size[1] // patch_size[1])
        # 将图片分割成无重叠的16X16的patches,每一张patch视为一个token,patches数量就是序列长度
        self.num_patches = self.grid_size[0] * self.grid_size[1] 

        self.proj = nn.Conv2d(in_c, embed_dim, kernel_size=patch_size, stride=patch_size)  # [3, 224, 224] -> [768, 14, 14]
        self.norm = norm_layer(embed_dim) if norm_layer else nn.Identity()

    def forward(self, x):
        B, C, H, W = x.shape
        assert H == self.img_size[0] and W == self.img_size[1], \
            f"Input image size ({H}*{W}) doesn't match model ({self.img_size[0]}*{self.img_size[1]})."

        # flatten: [B, C, H, W] -> [B, C, HW]
        # transpose: [B, C, HW] -> [B, HW, C]
        x = self.proj(x).flatten(2).transpose(1, 2)   # [B, 768, 14, 14] ->[B, 768,196]-> [B, 196, 768]
        x = self.norm(x)
        return x

[class] Embedding / Position Embedding：在输入Transformer Encoder之前注意需要加上[class]token以及Position Embedding。 在原论文中，作者说参考BERT，在刚刚得到的一堆tokens中插入一个专门用于分类的[class]token，这个[class] token是一个可训练的参数，数据格式和其他token一样都是一个向量，就是一个长度为768的向量，与之前从图片中生成的patch tokens拼接在一起，Cat([1, 768], [196, 768]) -> [197, 768]，最后的输出结果用来预测类别。这样一来，Transformer相当于一共处理了 N+1 个维度为 D 的token，并且只有最后一个token的输出用来预测类别。这种体系结构迫使patch token和class token之间传播信息。Position Embedding的作用是由于Transformer对输入的批量数据进行并行操作，失去了RNN固有的时间顺序特性，故需要加入位置编码来保证图片的位置信息不发生错乱。ViT 中的位置编码没有采用原版 Transformer 中的 sincos编码，而是直接设置为可学习的 Positional Encoding，采用的是一个可训练的参数（1D Pos. Emb.），是直接叠加在tokens上的（add），所以shape要一样。刚刚拼接[class]token后shape是[197, 768]，那么这里的Position Embedding的shape也是[197, 768]。

Transformer Encoder：在本文中采用纯Transformer结构，这里使用了Transformer中的Encoder(编码器），其中主要有两大组成部分，即Multi-Head Attention 和MLP。这里我们重点讲解MSA。

所谓自注意力，也就是说我们有一个序列X，然后我们想要算出X对X自己的注意力，也即X中的每个时间点与其余时间点的相关性（在注意力机制中表现为相似性），从而得到一个注意力矩阵。算出注意力矩阵后再将之用在序列X上来得到各个时间点的加权和，也即将其余各个时间点的信息都融合到了各个时间点中。

注意力机制通过注意力汇聚将查询（自主性提示）和键（非自主性提示）结合在一起，实现对值（感官输入）的选择倾向。注意力Attention机制的最核心的公式为：

其中，Q为Query、K为Key、V为Value。Q、K、V是从哪儿来的呢？Q、K、V其实都是从同样的输入矩阵X线性变换而来的。所谓的QKV本质上是代表了三个独立的矩阵，他们都是我们原本的序列X做了不同的线性变换之后的结果，都可以作为X的代表。然后为了得到注意力矩阵，我们肯定是先需要两个X的代表矩阵，因为计算相似度至少要有两个向量才可以计算，因此Q和K就是用于这个工作的。计算相似度有很多方式，可以直接做点乘（一般采用的都是Scaled Dot-Product Attention），也可以用个MLP来计算。然后再将结果经过一个softmax来保证输出的是attention weight，也即保证和为1，否则用这个attention matrix的时候数据的scale会不断的变化。得到了attention matrix之后要做的就是将这个weight用在X上，也即我们剩下的那个vector V身上。做法就是将attention matrix和V相乘求和即可得到最终结果。

以下是将图像作为输入，MSA中各部分形状变化：

这里的C对应论文下图中的D，N（num_patches+1)对应HW，C/n对应Dh。

Transformer最开始应用在NLP方向，其要求的输入形状为[batch_size, sequence_length, embedding_dim]，这里的输入的token为一张patches，故序列长度为patches数量，因为class token的存在，序列长度为patches+1。本文中先把q,k,v看成一个整体放在矩阵中，这里的qkv直接通过一个Linear层得到，输出维度为3*dim，所以其权重矩阵为U[dim , 3*dim]。后续通过改变形状得到对应的qkv[3, batch_size, num_heads, num_patches + 1, embed_dim_per_head]，因为q,k,v形状相同，故通过索引分别取出每一个[batch_size, num_heads, num_patches + 1, embed_dim_per_head]即为Q，K，V。其余实现和原Transformer中的MSA没有变化。

class MultiHeadAttention(nn.Module):
    def __init__(self,
                 dim,   # 输入token的dim
                 num_heads=8,
                 qkv_bias=False,
                 qk_scale=None,
                 attn_drop_ratio=0.,
                 proj_drop_ratio=0.):
        super(Attention, self).__init__()
        self.num_heads = num_heads
        head_dim = dim // num_heads
        self.scale = qk_scale or head_dim ** -0.5      # d**0.5
        self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias)
        self.attn_drop = nn.Dropout(attn_drop_ratio)
        self.proj = nn.Linear(dim, dim)
        self.proj_drop = nn.Dropout(proj_drop_ratio)

    def forward(self, x):
        # [batch_size, num_patches + 1, total_embed_dim]
        B, N, C = x.shape
        # qkv(): -> [batch_size, num_patches + 1, 3 * total_embed_dim]
        # reshape: -> [batch_size, num_patches + 1, 3, num_heads, embed_dim_per_head]
        # permute: -> [3, batch_size, num_heads, num_patches + 1, embed_dim_per_head]
        qkv = self.qkv(x).reshape(B, N, 3, self.num_heads, C // self.num_heads).permute(2, 0, 3, 1, 4)
        # [batch_size, num_heads, num_patches + 1, embed_dim_per_head]
        q, k, v = qkv[0], qkv[1], qkv[2]  # make torchscript happy (cannot use tensor as tuple)

        # transpose: -> [batch_size, num_heads, embed_dim_per_head, num_patches + 1]
        # @: multiply -> [batch_size, num_heads, num_patches + 1, num_patches + 1]
        attn = (q @ k.transpose(-2, -1)) * self.scale
        attn = attn.softmax(dim=-1)
        attn = self.attn_drop(attn)

        # @: multiply -> [batch_size, num_heads, num_patches + 1, embed_dim_per_head]
        # transpose: -> [batch_size, num_patches + 1, num_heads, embed_dim_per_head]
        # reshape: -> [batch_size, num_patches + 1, total_embed_dim]
        x = (attn @ v).transpose(1, 2).reshape(B, N, C)
        x = self.proj(x)
        x = self.proj_drop(x)
        return x

实验

数据集：  the ILSVRC-2012 ImageNet dataset with 1k classes                  

                ImageNet-21k with 21k classes                  

                JFT with 18k classes

在论文的Table1中有给出三个模型（Base/ Large/ Huge）的参数。其中的Layers就是Transformer Encoder中重复堆叠Encoder Block的次数，Hidden Size就是对应通过Embedding层后每个token的dim（向量的长度），MLP size是Transformer Encoder中MLP Block第一个全连接的节点个数（是Hidden Size的四倍），Heads代表Transformer中Multi-Head Attention的heads数。

Table 2 表明在JFT-300M上预训练的较小的VIT-L/16模型在所有任务上都优于BiT-L(在相同的数据集上预先训练)，同时需要更少的计算资源来训练。较大的模型 ViT-H/14 进一步提高了性能，尤其是在更具挑战性的数据集——ImageNet、CIFAR-100 和 VTAB 。

Table 6 是论文用来对比ViT，Resnet（使用的卷积层和Norm层都进行了修改）以及Hybrid模型的效果。通过对比发现，在训练epoch较少时Hybrid优于ViT，但当epoch增大后ViT优于Hybrid。

我们探索了Transformer在图像识别中的直接应用。与在计算机视觉中使用自注意力的先前工作不同，除了初始提取patch步骤之外，我们不会将特定于图像的归纳偏差引入架构中。相反，我们将图像解释为一系列patches，并通过 NLP 中使用的标准 Transformer 编码器对其进行处理。当与大型数据集的预训练相结合时，这种简单的但可扩展的策略效果很好。因此，Vision Transformer 在许多图像分类数据集上超过了最先进的技术，同时预训练成本相对较低。虽然这些初步结果令人鼓舞，但仍有许多挑战。一种是将ViT应用于其他计算机视觉任务，如检测和分割。我们的结果，加上Carion等人(2020)的结果，表明了这种方法的前景。另一个挑战是继续探索自监督的预训练方法。我们最初的实验表明自监督预训练有所提高，但自监督预训练和大规模监督预训练之间仍然存在很大差距。最后，ViT的进一步缩放可能会导致性能的提高。