Vision Transformer的讲解和代码实现

1、ViT介绍

从深度学习暴发以来，CNN一直是CV领域的主流模型，而且取得了很好的效果，相比之下，基于self-attention结构的Transformer在NLP领域大放异彩。虽然Transformer结构已经成为NLP领域的标准，但在计算机视觉领域的应用还非常有限。

ViT（vision transformer）是Google在2020年提出的直接将Transformer应用在图像分类的模型，通过这篇文章的实验，给出的最佳模型在ImageNet1K上能够达到88.55%的准确率（先在Google自家的JFT数据集上进行了预训练），说明Transformer在CV领域确实是有效的，而且效果还挺惊人。

2、模型详解

在讲解ViT原理之前，读者需要有self-attention的相关知识，之前博文有讲过，这里就不复述了。

下图是Vision Transformer的结构，乍一看和self-attention的结构非常像。主要由三个模块组成：

• Linear Projection (Patch + Position 的Embedding层)
• Transformer Encoder（详细结构见图右边）
• MLP Head（分类层）
  

Embedding层详解

从图中看，Transformer Encoder读入的是一小块一小块这样的图片。这样做的原因是：将一个个小块图像视为token（token可以理解为NLP中的一个字或词），在Transformer中计算每个token之间的相关性。这一点就和卷积神经网络有很大区别了。以往的CNN，以卷积 + 池化的方式不断下采样，这样理论上模型可以通过加深模型深度，达到增大感受野的目的。不过这样会有两个缺点：

1. 实际结果中显示，CNN对边缘的响应很弱。这也非常好理解，越靠边缘的像素，因为被卷积次数少，自然在梯度更新时，贡献更少。
2. CNN只能和临近像素计算相关性。由于其滑窗卷积的特性，无法对非领域的像素共同计算，例如左上角的像素无法和右下角的像素联合卷积。这就导致了某些空间信息是无法利用的。同时根据MAE论文中所说的，自然图像具有冗余性，即相邻像素点代表的信息是差不多的，所以只计算领域像素无法最大化利用图像特征。

回到ViT中，仅仅把图像拆分成小块（patch）是不够的，Transformer Encoder需要的是一个向量，shape为[num_token, token_dim]。对于图片数据来说，shape为[H,W,C]是不符合要求的，所以就需要转换，要将图片数据通过这个Embedding层转换成token。以ViT-B/16为例：

假设输入图像为224x224x3，一个token原始图像shape为16x16x3，那这样就可以将图像拆分成$(224/16{)}^2=196$个patch，然后将每个patch线性映射至一维向量中，那么这个一维向量的长度即为$16\ast 16\ast 3=768$维。将196个token叠加在一起最后维度就是[196, 768]。

在代码实现中，patch的裁剪是用一个patch_size大小的卷积以image_size // patch_size的步长进行卷积实现的：

class PatchEmbed(nn.Module):
    """
    2D Image to Patch Embedding
    """

    def __init__(self, image_size=224, patch_size=16, in_c=3, embed_dim=768, norm_layer=None):
        """
        Map input tensor to patch.
        Args:
            image_size: input image size
            patch_size: patch size
            in_c: number of input channels
            embed_dim: embedding dimension. dimension = patch_size * patch_size * in_c
            norm_layer: The function of normalization
        """
        super().__init__()
        image_size = (image_size, image_size)
        patch_size = (patch_size, patch_size)
        self.image_size = image_size
        self.patch_size = patch_size
        self.grid_size = (image_size[0] // patch_size[0], image_size[1] // patch_size[1])
        self.num_patches = self.grid_size[0] * self.grid_size[1]

        # The input tensor is divided into patches using 16x16 convolution
        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.image_size[0] and W == self.image_size[1], \
            f"Input image size ({H}*{W}) doesn't match model ({self.image_size[0]}*{self.image_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

同时需要注意的是，token可以添加位置编码。由于self-attention的机制，在没有字符位置的情况下，“我爱你”和“你爱我”的计算结果是一样的。所以当添加了位置编码，分类准确率将会更高。所以迁移到ViT的Embeding层中，位置信息是以***Add***的方式融合在输入tensor中的。

另外，输入tensor还需要有一个class token（分类层），数据格式和其他token一样，长度为768的向量，与位置编码的融合方式不一样，这里做的是***Concat***，这样做是因为分类信息是在后面需要取出来单独做预测的，所以不能以***Add***方式融合，shape也就从[196, 768]变为[197, 768]。

Transformer Encoder详解

Transformer Encoder也就是堆叠encoder block几次，和常规的CNN类似，主要由几个部分组成：

• Layer Normalization
  相比于Batch Normalization来说，BN是针对所有的样本，对某一个特征图计算均值和方差，然后然后对这个特征图神经元做归一化。LN是对某一个样本，计算该样本所有特征图的均值和方差，然后对这个样本做归一化。BN适用于不同mini batch数据分布差异不大的情况，而且BN需要开辟变量存每个节点的均值和方差，空间消耗略大；而且 BN适用于有mini_batch的场景。LN只需要一个样本就可以做归一化，可以避免 BN 中受 mini-batch 数据分布影响的问题，也不需要开辟空间存每个节点的均值和方差。
• Multi-Head Attention
  这个结构是self-attention中的一种，在之前的博客中有详细的讲解。
• DropOut/DropPath
  论文中描述是使用了DropOut，但实际别人复现的代码中使用的是DropPath。两者对最后的结果影响不大，这里也不展开详细说了。
• MLP Block
  如下图右侧所示，即一个全连接层 + GELU激活函数 + DropOut组成，采用的是倒瓶颈结构，输入特征层经过一次全连接之后，通道膨胀为原来的4倍，后一个全连接层再恢复成原来的数目。所以经过MLP Block之后tensor的shape是不变的。
  

Transformer Encoder代码实现：

class Block(nn.Module):
    def __init__(self,
                 dim,
                 num_heads,
                 mlp_ratio=4.,
                 qkv_bias=False,
                 qk_scale=None,
                 drop_ratio=0.,
                 attn_drop_ratio=0.,
                 drop_path_ratio=0.,
                 act_layer=nn.GELU,
                 norm_layer=nn.LayerNorm):
        super(Block, self).__init__()
        self.norm1 = norm_layer(dim)
        self.attn = Attention(dim, num_heads=num_heads, qkv_bias=qkv_bias, qk_scale=qk_scale,
                              attn_drop_ratio=attn_drop_ratio, proj_drop_ratio=drop_ratio)
        # NOTE: drop path for stochastic depth, we shall see if this is better than dropout here
        self.drop_path = DropPath(drop_path_ratio) if drop_path_ratio > 0. else nn.Identity()
        self.norm2 = norm_layer(dim)
        mlp_hidden_dim = int(dim * mlp_ratio)
        self.mlp = MLP(in_features=dim, hidden_features=mlp_hidden_dim, act_layer=act_layer, drop=drop_ratio)

    def forward(self, x):
        x = x + self.drop_path(self.attn(self.norm1(x)))
        x = x + self.drop_path(self.mlp(self.norm2(x)))

        return x

MLP Head详解

经过Transformer Encoder后输出tensor的shape和输出tensor的shape是一样的，以ViT-B/16为例，输入是[197, 768]，输出还是[197, 768]。下游任务如果是分类模型，需要提取出对应的class token获取分类结果。Vision Transformer中说MLP是由一个全连层 + Tanh激活 + 全连接层组成。但实际使用起来一层全连接层直接做分类即可。

完整模型（ViT-B/16为例）

具体代码如下：

class VisionTransformer(nn.Module):
    def __init__(self, image_size=224, patch_size=16, in_c=3, num_classes=1000,
                 embed_dim=768, depth=12, num_heads=12, mlp_ratio=4.0, qkv_bias=True,
                 qk_scale=None, representation_size=None, distilled=False, drop_ratio=0.,
                 attn_drop_ratio=0., drop_path_ratio=0.5, embed_layer=PatchEmbed, norm_layer=None,
                 act_layer=None):
        """
        Args:
            image_size (int, tuple): input image size
            patch_size (int, tuple): patch size
            in_c (int): number of input channels
            num_classes (int): number of classes for classification head
            embed_dim (int): embedding dimension, dim = patch_size * patch_size * in_c
            depth (int): depth of transformer
            num_heads (int): number of attention heads
            mlp_ratio (int): ratio of mlp hidden dim to embedding dim
            qkv_bias (bool): enable bias for qkv if True
            qk_scale (float): override default qk scale of head_dim ** -0.5 if set
            representation_size (Optional[int]): enable and set representation layer (pre-logits) to this value if set
            distilled (bool): model includes a distillation token and head as in DeiT models
            drop_ratio (float): dropout rate
            attn_drop_ratio (float): attention dropout rate
            drop_path_ratio (float): stochastic depth rate
            embed_layer (nn.Module): patch embedding layer
            norm_layer: (nn.Module): normalization layer
        """
        super(VisionTransformer, self).__init__()
        self.num_classes = num_classes
        self.num_features = self.embed_dim = embed_dim  # num_features for consistency with other models
        self.num_tokens = 2 if distilled else 1
        norm_layer = norm_layer or partial(nn.LayerNorm, eps=1e-6)  # partial类似闭包函数，第一个参数是函数对象，后面的参数是第一个函数对象的实参
        act_layer = act_layer or nn.GELU

        self.patch_embed = embed_layer(image_size=image_size, patch_size=patch_size, in_c=in_c, embed_dim=embed_dim)
        num_patches = self.patch_embed.num_patches

        self.cls_token = nn.Parameter(torch.zeros(1, 1, embed_dim))     # nn.Parameter转换函数，把参数转换为可训练变量
        self.dist_token = nn.Parameter(torch.zeros(1, 1, embed_dim)) if distilled else None
        self.pos_embed = nn.Parameter(torch.zeros(1, num_patches + self.num_tokens, embed_dim))
        self.pos_drop = nn.Dropout(p=drop_ratio)

        dpr = [x.item() for x in torch.linspace(0, drop_path_ratio, depth)]  # stochastic depth decay rule
        self.blocks = nn.Sequential(*[
            Block(dim=embed_dim, num_heads=num_heads, mlp_ratio=mlp_ratio, qkv_bias=qkv_bias, qk_scale=qk_scale,
                  drop_ratio=drop_ratio, attn_drop_ratio=attn_drop_ratio, drop_path_ratio=dpr[i],
                  norm_layer=norm_layer, act_layer=act_layer)
            for i in range(depth)
        ])
        self.norm = norm_layer(embed_dim)

        # Representation layer
        if representation_size and not distilled:
            self.has_logits = True
            self.num_features = representation_size
            self.pre_logits = nn.Sequential(OrderedDict([
                ("fc", nn.Linear(embed_dim, representation_size)),
                ("act", nn.Tanh())
            ]))
        else:
            self.has_logits = False
            self.pre_logits = nn.Identity()     # placeholder

        # Classifier head(s)
        self.head = nn.Linear(self.num_features, num_classes) if num_classes > 0 else nn.Identity()
        self.head_dist = None
        if distilled:
            self.head_dist = nn.Linear(self.embed_dim, self.num_classes) if num_classes > 0 else nn.Identity()

        # Weight init
        nn.init.trunc_normal_(self.pos_embed, std=0.02)
        if self.dist_token is not None:
            nn.init.trunc_normal_(self.dist_token, std=0.02)

        nn.init.trunc_normal_(self.cls_token, std=0.02)
        self.apply(_init_vit_weights)

    def forward_features(self, x):
        # [B, C, H, W] -> [B, num_patches, embed_dim]
        x = self.patch_embed(x)  # [B, 196, 768]
        # [1, 1, 768] -> [B, 1, 768]
        cls_token = self.cls_token.expand(x.shape[0], -1, -1)
        if self.dist_token is None:
            x = torch.cat((cls_token, x), dim=1)  # [B, 197, 768]
        else:
            x = torch.cat((cls_token, self.dist_token.expand(x.shape[0], -1, -1), x), dim=1)

        x = self.pos_drop(x + self.pos_embed)
        x = self.blocks(x)
        x = self.norm(x)
        if self.dist_token is None:
            return self.pre_logits(x[:, 0])
        else:
            return x[:, 0], x[:, 1]

    def forward(self, x):
        x = self.forward_features(x)
        if self.head_dist is not None:
            x, x_dist = self.head(x[0]), self.head_dist(x[1])
            if self.training and not torch.jit.is_scripting():
                # during inference, return the average of both classifier predictions
                return x, x_dist
            else:
                return (x + x_dist) / 2
        else:
            x = self.head(x)

        return x

3、总结

至此，Vision Transformer的原理及模型结构以及全部介绍完毕。从实验给出的结果，在当时也达到了SOTA，但是相比于CNN，它需要更多的数据集。在小数据集上训练出来的精度是不如CNN的，但在大数据集上ViT精度更高。一个直观的解释是：ViT因为self-attention独特的机制，更多的利用token与token跨像素之间的信息，而CNN只是对领域的像素进行计算，所以相同参数的情况下，ViT获得的信息更多，在某种程度上可以看成是模型深度更深。所以小数据集上ViT是欠拟合的。实际开发中的做法是：基于大数据集上训练，得到一个预训练权重，然后再在小数据集上Fine-Tune。

后续将会补充Vision Transformer实现的完整代码，先挖个坑。