AN IMAGE WORTH 16×1616times 16 WORDS TRANSFORMERS FOR IMAGE RECOGNITION AT SCALE

AN IMAGE WORTH $16\times 16$ WORDS : TRANSFORMERS FOR IMAGE RECOGNITION AT SCALE.

论文地址：https://arxiv.org/abs/2010.11929

源码地址：https://github.com/lucidrains/vit-pytorch

引言：

基于自注意力机制的模型，特别是在$Transformer$中，已经成为在$NatureLanguageProcessing$(自然语言处理，下面用$NLP$代指)任务中的一个选择。主要的方法是在一个大型的文本语料库进行预训练，并在较小的数据集上，针对特定的任务进行微调。通过$Transformer$的计算效率和可扩展性，随着模型和数据集的增长，仍然没有性能饱和的迹象。但是在计算机视觉的应用仍然有限。在视觉上，注意力要么与卷积网络结合使用，要么用于替换卷积网络的某些组件，同时保持其整体结构不变。$ViT$表明传统的$CNN$方法在计算机视觉任务中并不是唯一选择，应用于图像块序列的纯变换可以很好地执行图像分类任务。在对大量数据进行预训练并传输到多个中小型图像识别基准($ ImageNet、CIFAR-100、VTAB$等)时，视觉转换器($ViT$​)与最先进的卷积网络相比取得了优异的结果，而训练所需的计算资源要少得多。

模型结构：

$PatchEmbedding$

$Transformer$在$NLP$任务中，将 $1D$ 序列作为一个$token$作为输入，而在原文中，为了处理$2D$数据，即将 $x\in R^{C\times H\times W}$ 的图像进行重塑，即将图像展平为 $x\in R^{N\times (P^2\cdot C)}$ ，其中对于 $C$ 代表输入图像的通道数，$(H,W)$代表输入图像的分辨率。重塑后 $P$ 代表输入图像分块后，每个块($patch$)中的图像分辨率，即为 $P\times P $ , $N$ 代表输入图像的总分块数，即为 $\frac{H\times W}{P\times P}$ 。在原论文中，作者说参考 $BERT$ ，将得到的$tokens$中插入一个专门用于分类的 $classtoken$ ，这个是一个可训练的参数，数据格式和其他 $token$ 一样都是一个向量，就是一个长度为 $768$ 的向量，与之前从图片中生成的 $ token$s拼接在一起，最后的输出结果用来预测类别。这样一来，$Transformer$相当于一共处理了 $N+1$ 个维度为 $D$ 的 $token$，并且只有最后一个 $token$ 的输出用来预测类别。(计算机视觉任务中，一个 $token$ 表示一个块 $patch$ )，以输入图像$3\times 224\times 224$为例，为了可以匹配$Transformer$ 输入($sequencelength,tokendim$)，原文中进行 $16\times 16$ 的大小进行分块，从而一共分块($Patchs$)中个数为 $\frac{224}{16}\times \frac{224}{16}=196$个，从而一共有 $196$ 个$tokens$，再加上我们训练得到的$classtoken$，从而一共得到 $ 197 $ 个$tokens$，即我们已经得到了$sequencelength=197$，而对于每一个 $token $ 的维度，即为$3\times 16\times 16=768$。这里的 $768$ 即对应原文中的 $D$ 维。

代码

``

import torch
import torch.nn as nn
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])
        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)
        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)
        x = self.norm(x)
        return x

这里的卷积操作即为分块操作，利用尺寸为$16\times 16$ 的卷积核，步长设置为 $16$ , $0$ 填充为默认值。

$PostionEmbedding$

对于位置嵌入，它与块嵌入一起作为 $Transformer$ 的输入，在原论文中，对于这样的输入的刻画是这样的，
$$Z_0=[x_{class};x_p^{1}E;x_p^{2}E;x_p^{3}E;x_p^{4}E;\dots ;x_p^{N}E;]+E_{pos},E\in R^{(P^2\times C)\times D},E_{POS}=R^{(N+1)\times D}$$
这里作者利用 $1D$ 的向量作为位置嵌入，原文中也同时讨论了关于位置嵌入的多种情况：无位置嵌入，即直接将分块作为输入；$2D $ 位置嵌入：将输入视为二维的网格。在这种情况下，学习了两组嵌入，每组用于一个轴，X嵌入和Y嵌入，每个都具有大小$\frac{D}{2}$。然后，基于输入中路径上的坐标，我们将$X$和$Y$嵌入连接起来，以获得该补丁的最终位置嵌入；相对位置嵌入。通过实验证明：

从上表中，我们可以看到除了在有无位置嵌入时，模型的性能发生了较大的变化，在对于其他位置编码的选择时没有较大差别。对于位置编码的刻画为一个可学习的参数张量，通过我们在训练阶段设计的优化器进行参数的更新。

self.pos_embed = nn.Parameter(torch.zeros(1, num_patches + self.num_tokens, embed_dim))

这里我们初始化了一个全为 $0$ 的张量。第一个$1$代表$batchsize$，第二个为所有 $token$ 的数量，第三个代表每个块($patch$)的维度。我们现在回想在上一小节中结尾得到的 $patchembedding$ 的形状，即为$(B,197,768)$​ ，从而我们可以看到两种嵌入得到的形状是一样的，这样就可以得到论文中的公式：
$$Z_0=[x_{class};x_p^{1}E;x_p^{2}E;x_p^{3}E;x_p^{4}E;\dots ;x_p^{N}E;]+E_{pos},E\in R^{(P^2\times C)\times D},E_{POS}=R^{(N+1)\times D}$$
其中 $E$​ 代表一个全连接层。

$Encoderblock$

终于到了我们这个模型最重要的结构，我们在上面得到了 $z_0$​ ，先进行随机失活($dropout/droppath$)后，然后作为这个编码器的输入，现在我们详细详细讲解一下这个编码器的结构。

图源：http://t.csdnimg.cn/l7gL8

如上图，我们可以看到 $EncoderBlock$ 的结构，主要由三个部分组分：层归一化($LayerNorm$)，多头注意力机制($Multi-HeadAttention$)，随机失活($dropout/droppath$​)。

我们先从随机失活开始解读：

$1.$ $dropout$ :在机器学习中，我们最怕遇到的问题之一就是过拟合，随着神经网络中参数越多，而我们的训练样本又很少时，此时就会很容易出现过拟合的问题。这里不过多讲解过拟合，$dropout$ 的提出会在训练深度网络模型中缓解过拟合的问题。在2012年，Hinton在其论文Improving neural networks by preventing co-adaptation of feature detectors中提出$Dropout$​​。

图源：深度学习中Dropout原理解析 - Microstrong的文章 - 知乎
https://zhuanlan.zhihu.com/p/38200980

如上图，左侧为标准的神经网络，而使用 $dropout$ 算法后，如右图，我们可以看到中间三层的神经元部分失活，使得上一层的输出不能通过。即我们在前向传播的时候，让某个神经元的激活值以一定的概率 $p$ 停止工作，可以防止模型过多去学习局部的特征，以此来提高模型的泛化能力。

$2.$ $droppath:$ 这一个正则化手段类似于 $dropout$,通过对于深度学习模型的多个分支的子路径进行随机失活。它的提出为了有效的防止过拟合现象，以及网络退化问题。在计算机视觉中运用该算法，就是把得到的特征图的像素点随机设置为0，即代表一种失活。我们可以通过源代码来更好的理解：

``

def drop_path(x, drop_prob: float = 0., training: bool = False):

    if drop_prob == 0. or not training:
        return x
    keep_prob = 1 - drop_prob
    shape = (x.shape[0],) + (1,) * (x.ndim - 1)  # work with diff dim tensors, not just 2D ConvNets
    random_tensor = keep_prob + torch.rand(shape, dtype=x.dtype, device=x.device)
    random_tensor.floor_()  # binarize
    output = x.div(keep_prob) * random_tensor
    return output


class DropPath(nn.Module):
    """
    Drop paths (Stochastic Depth) per sample  (when applied in main path of residual blocks).
    """

    def __init__(self, drop_prob=None):
        super(DropPath, self).__init__()
        self.drop_prob = drop_prob

    def forward(self, x):
        return drop_path(x, self.drop_prob, self.training)

这里来解释一下这段代码，帮助我们更好的理解，首先先来看

``

shape = (x.shape[0],) + (1,) * (x.ndim - 1)  # work with diff dim tensors, not just 2D ConvNets
random_tensor = keep_prob + torch.rand(shape, dtype=x.dtype, device=x.device)
random_tensor.floor_()  # binarize
output = x.div(keep_prob) * random_tensor

首先定义了一个元组($tuple$)，其维度与输入 $x$ 一致，但是除了第一维以外，其余都设为 $1$ 。其中$x.ndim$表示 $x$ 的维数。举个例子假设我们输入的 $x$ 的形状为，此时 $x.ndim=4$，那么我们得到的将会是 $(B,1,1,1)$。再下一步，通过激活率加上随机生成的形如 $(B,1,1,1)$ 的随机张量，然后我们进行向下取整，这保证了张量中所有元素，取值只有 $0$ 或者 $1$ 。最后先让 $x$ 中的元素除以 $keepprob$ 之后,这里其实相当于一个正则化手段，再进行广播乘以生成的随机张量，特征图通过$droppath$ 之后，部分特征图失活。这段代码写的很巧妙，因为他可以适应不同维数的输入，比如我们的原始图像形状为$(B,3,256,256) $，而在 $transformer$ ，可以接受的输入形状为就是($B,sequencelength,tokendim$)。

下面我们来介绍这个多头注意力机制，和标准的 $Transformer$​ 类似，原文作者没有进行大的改动，

``

class Attention(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
        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

代码没有什么可以讲的，就是普通的多头注意力机制的代码。

最后是 $MLP$ :

``

class Mlp(nn.Module):
    """
    MLP as used in Vision Transformer, MLP-Mixer and related networks
    """

    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

实验过程：

作者评估了$ViT,ResNet,Hybird$ 三种不同的模型，并在不同的数据集上进行评估。当考虑到预训练模型的计算成本时，$ViT$ 表现非常有利，以较低的预训练成本在大多数识别基准上达到最先进的水平。

可以看到在不同的数据集上，在 $JFT-300M$ 上预训练的较小的 $ViT-L/16$ 模型在所有任务上都优于$BiT-L$（在相同数据集上预训练），同时需要更少的计算资源来训练。更大的模型 $ViT-H/14$ 进一步提高了性能，特别是在更具挑战性的数据集上-$ImageNet，CIFAR-100$ 和 $VTAB$​ 数据集上。

论文中，基于$BERT$ 在自然语言处理的成功经验，引申出三种不同规模的 $ViT$ 模型，如上图。我们还使用了与$BERT$ 相似的训练策略和 $Adam$ 优化器进行模型的预训练和微调，并使用学习率的线性预热和衰减。我们还使用了权重衰减($weightdecay$)和 $dropout$​ 等正则化技术来提高模型的泛化能力。

下面我们将利用这三种 $ViT$ 变种模型 与 $Sota$ 模型进行对比：

作者注意到，预训练效率可能不仅受到架构选择的影响，还受到其他参数的影响，例如优化器，权重衰减等因素，从而作者针对不同的架构提供了性能与计算成本的对照研究。其中包括不同深度的 $ResNet，ViT$ 的变种，还有混合模型。

在计算效率上，论文中提到实际执行时间$(wall-clocktime)$下，不同输入大小的模型结构执行推理的速度比较，如左图。可以看到 $ViT$ 的速度和 $ResNet$ 大致相同。而对于右图，关于一个内核对于不同模型来说，能接受的最大输入批次。从图中可以看出大型的 $ViT$ 模型在内存效率上优于 $ResNet $ 的。

$Natural,Specialized,andStructured$ 是 $VTAB（VisionTransformerAssessmentBenchmark）$任务组的分类。这些任务组用于评估图像分类模型在不同类型的任务上的性能。

$Natural$（自然）任务组包括7个任务，涵盖了一般的自然图像分类，例如在ImageNet上的分类任务。$Specialized$（专业）任务组包括4个任务，涵盖了一些特定领域的图像分类，例如医学图像分类或车辆分类。$Structured$（结构化）任务组包括8个任务，涵盖了对结构化图像的分类，例如场景分类或物体检测。 这些任务组的目的是测试模型在不同类型的图像分类任务上的泛化能力和性能。

现在我们了解这几种任务之后，$ViT-H$ 的分类性能普遍优于其他三种模型，只在专业任务组时与 $BiT-L$​ 性能相当。

随后作者评估了数据集大小对于模型性能的影响。首先作者在规模逐渐增大的不同数据集上预训练 $ViT$​ 模型，如图

在三种不同规模的数据集上进行预训练，并在小数据集上进行推理之后，我们可以看到当预训练数据集较小时，$BiT$ 的准确率是要高于 $ViT$ 的变种模型组的，然后随着预训练数据集增大，$ViT$ 变种组合开始追上$BiT$ 。最后到了 $JFT-300M$ ,我们可以看到已经完全超过了$BiT$，并且在变种组合中，最大的 $ViT$ 得到的准确率最高。

这里论文中给到了更为完整的结果

最后为了了解 $ViT$ 如何处理图像数据，作者对模型内部的输出做了可视化。

但是我们需要思考一个问题就是，这个 $Self-Attention$​ 关于整合信息在模型中起了多大作用？

为此作者根据注意力权重计算图像空间中信息被整合的平均距离，这种距离类似与卷积神经网络中的感受野大小，一些头部注意到已经在最低层的图像的大部分，这表明模型确实使用了全局整合信息的能力，如下图。其他注意头在低层中始终具有较小的注意距离。在$Transformer$之前应用$ResNet$的混合模型中，这种高度局部化的注意力不太明显，这表明它可能与卷积神经网络中的早期卷积层具有类似的功能。此外，注意距离随着网络深度的增加而增加。在全局上，我们发现该模型关注与分类语义相关的图像区域。

实验结果：

作者已经探索了 $Transformer$ 在图像识别中的直接应用。与先前在计算机视觉中使用自注意的作品不同，除了初始补丁提取步骤之外，我们没有将图像特定的归纳偏差引入到架构中。相反，我们将图像解释为一系列补丁，并通过 $NLP$ 中使用的标准 $Transformer$ 编码器对其进行处理。这种简单但可扩展的策略在与大型数据集的预训练相结合时效果惊人。因此，$ViT$ 在许多图像分类数据集上匹配或超过了最新技术水平，同时预训练相对便宜。

注记：

归纳偏好：根据我们已有的经验和观察到的样本来做出一般性的结论。这种偏差可能导致我们忽视了一些特殊情况或者没有考虑到其他可能的解释。 举个例子来说明归纳偏差：假设你在一个新的城市旅行，你发现这个城市的人们都非常友好和热情。你在几个餐馆用餐时，服务员都非常友好地接待你，给你提供了很好的服务。基于这些观察，你可能会得出结论，这个城市的所有人都非常友好。 然而，这个结论可能存在归纳偏差。因为你只是在几个餐馆遇到了友好的服务员，并不能代表整个城市的人都是友好的。可能其他地方或者其他情况下，你会遇到不太友好的人。这种归纳偏差导致你根据有限的观察得出了一个过于一般化的结论。 因此，归纳偏差提醒我们在做出一般性结论时要谨慎，需要考虑更多的证据和情况，避免过度归纳。

者没有考虑到其他可能的解释。 举个例子来说明归纳偏差：假设你在一个新的城市旅行，你发现这个城市的人们都非常友好和热情。你在几个餐馆用餐时，服务员都非常友好地接待你，给你提供了很好的服务。基于这些观察，你可能会得出结论，这个城市的所有人都非常友好。 然而，这个结论可能存在归纳偏差。因为你只是在几个餐馆遇到了友好的服务员，并不能代表整个城市的人都是友好的。可能其他地方或者其他情况下，你会遇到不太友好的人。这种归纳偏差导致你根据有限的观察得出了一个过于一般化的结论。 因此，归纳偏差提醒我们在做出一般性结论时要谨慎，需要考虑更多的证据和情况，避免过度归纳。

受到$Transformer$缩放($scaling$)的启发，这里的缩放指的是在$NLP$中，通过增加$Transformer$的层数和隐藏单元的维度，来显著提高模型的性能。从而在$ViT$中，作者通过增加模型的深度，可以提高模型对图像特征的建模能力。而增加隐藏单元的维度，则可以增加模型的表达能力。