Vision Transformer模型学习笔记

模型构成

根据原论文，ViT模型主要有以下三个部分组成：

• 输入图片—>Patch Embedding
• Transformer 编码器
• 用于分类的MLP Head
  

图1 原论文模型图

理解Patch Embedding

在深度学习中，图像数据的输入维度为[B, C, H, W],其中B表示batch size,C表示图片的通道数，H、W分别表示图片的高和宽。而标准的Transformer编码器要求的输入是二维的向量序列,其维度为[num_patches, patch_dim]因此，需要对图片进行预处理以满足标准Transformer编码器的输入要求。如图2所示，需要将输入的图片数据按照指定的patch尺寸将图片划分成一系列的patch。
假设输入一张图片，其大小为224x224，输入网络tensor的维度为[1, 3, 224,224]；每个patch的尺寸为16x16，那么可以将该图片划分成196张大小为16x16的小图片。由于我们的目的是要将图片数据转换成满足Transformer编码器要求的二维向量输入[num_patches, patch_dim],那么此时的patch_dim = 16x16x3(这里的3为图片的通道数)=768 ，而num_patches=14x14=196，这样我们就可以将每张小图片(patch)映射并展平成维度为768的一维向量。那么我们目前得到的转换后的tokens的维度为[1, num_patches, patch_dim ]。而这个过程主要通过一个卷积层和展平(flatten)操作来实现，卷积核大小为每个patch的尺寸，这里是(16,16)，同时设置stride与卷积核大小相同，代码如下：

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]})."
        # proj:[B, C, H, W] -> [B, embed_dim, grid_size[0], grid_size[1]]
        # flatten: [B, embed_dim, grid_size[0], grid_size[1]] -> [B, embed_dim, num_patches]
        # transpose: [B, embed_dim, num_patches] -> [B, num_patches, embed_dim]
        x = self.proj(x).flatten(2).transpose(1, 2)
        x = self.norm(x)
        return x

图2 patch embedding过程

紧接着，需要将上面转换到的tokens加上分类用的class token以及位置嵌入(Position embedding)等信息。这个class token是一个可学习的训练参数，会在模型的训练过程中同步更新迭代，它也是一个一维向量，其维度与经映射后的每个patch的维度保持一致，在上面的例子中就是patch_dim(768),总的维度为[1, 1,patch_dim]需要与tokens拼接到一起。加上class token 之后，tokens的维度将由[1, num_patches, patch_dim]变为[1, num_patches+1, patch_dim].
位置嵌入采用的是一维位置嵌入(Positional Encoding),它也是一个可学习的训练参数，是直接叠加在tokens上（加法操作）的，因此它的维度与tokens保持一致，这样才能做相加运算。最后我们得到作为Transformer编码器输入的tokens的维度为[1, num_patches+1, patch_dim]。

Transformer 编码器

Self-Attention

原论文链接：Attention Is All You Need

图3 原论文中的self-attention的解释框图

接下来将详细讨论其中的实现过程。为方便解释，假设Transformer编码器的输入维度为(2, 2),即有两个patch,每个patch都是维度为2(patch_dim)的一维向量。Attention公式中的Q、K、V的维度和输入向量的维度是一致的，它们由输入向量经过线性变换得到的，分别由三个可学习的维度为[patch_dim, patch_dim]的变换矩阵W_Q、W_K、W_V计算得到,变换矩阵对于所有向量是共享的，如图4所示。X中的每一行表示一个patch对应的一维向量，即[x₁₁,x₁₂]和[x₂₁,x₂₁]。

图4 Q、K、V的计算过程 Attention公式的理解：

• Q：即query,作为查询后面会和K去进行匹配，也就是QK^T。
• K：即key，Q的匹配对象，QK^T的输出经过Softmax处理后，每一行的值的和为1，匹配的过程可以理解成计算输入向量两两之间的相关性，相关性越大对应V中元素的权重也就越大。
• V：即value, 可以理解为从输入X中学习到的信息。

下面举一个例子进行计算的演示过程，方便理解。

Multi-Head Attention

理解了Self-Attention,Multi-Attention的理解就水到渠成了。它的理解可以类比分组卷积的操作。由图3所示，简单来说，将输入X分别经过变换矩阵W_Q、W_K、W_V计算得到Q、K、V之后，将它们分别均分成head等份。拿我们前面得到的tokens为例，其维度为[1, num_patches+1, patch_dim]，经过线性变换后Q、K、V的维度也为[1, num_patches+1, patch_dim]，此时将Q、K、V均分成head等份，假设head=2，这时Q、K、V将分别拆分成[Q₁, Q₂]、[K₁, K₂]、[V₁, V₂]，这时，Q₁、K₁、V₁的维度将变为[1, 1, num_patches+1, patch_dim],实际上经过均分后Q、K、V的维度可以理解为[1, head,num_patches+1, patch_dim//head],这里的 patch_dim//head = 768// 2=384。
紧接着分别将Q₁、K₁、V₁；Q₂、K₂、V₂进行Self-Attention的计算，计算之后将得到的两组结果进行拼接，接着将拼接后的结果通过W^O（可学习的参数）进行融合,也是一个线性操作的过程，可以理解为一个矩阵运算，如下面的公式所示：

主要的代码实现如下：

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 torch script 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)
        # 将每个head完成Self-attention之后的拼接结果通过线性操作来实现融合。
        x = self.proj(x)
        x = self.proj_drop(x)
        return x

Self-Attention与Multi-Head Attention 计算量对比

从原论文的描述中可以了解到两者的计算量其实相差不大。使用多头注意力机制主要是为了能够融合来自不同head部分学习到的信息。

MLP模块

MLP模块就是包含两个线性连接层和激活函数以及dropout的多层感知机结构，如图5所示。其中第一个全连接层会将输入维度扩大k倍，然后经第二个全连接层进行恢复到原来的维度。

图5 MLP模块结构图 至此，Transformer 编码器相关内容的介绍到此就结束了。

MLP head

经过Transforme编码器处理后的输出，将分类所用的class token 部分单独提取出来，其维度为[1, patch_dim]，每张图片对应一个一维向量，该向量直接通过一个全连接层来进行分类。

完整的模型框图

代码链接