Vision Transformer模型学习笔记

Vision Transformer 模型学习笔记

  • 模型构成
  • 理解Patch Embedding
  • Transformer 编码器
    • Self-Attention
    • Multi-Head Attention
    • Self-Attention与Multi-Head Attention 计算量对比
    • MLP模块
  • MLP head
  • 完整的模型框图
  • 代码链接

模型构成

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

  • 输入图片—>Patch Embedding
  • Transformer 编码器
  • 用于分类的MLP Head
    Vision Transformer模型学习笔记_第1张图片
图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

Vision Transformer模型学习笔记_第2张图片

图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
Vision Transformer模型学习笔记_第3张图片

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

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

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

下面举一个例子进行计算的演示过程,方便理解。
Vision Transformer模型学习笔记_第5张图片

Multi-Head Attention

理解了Self-Attention,Multi-Attention的理解就水到渠成了。它的理解可以类比分组卷积的操作。由图3所示,简单来说,将输入X分别经过变换矩阵WQ、WK、WV计算得到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将分别拆分成[Q1, Q2]、[K1, K2]、[V1, V2],这时,Q1、K1、V1的维度将变为[1, 1, num_patches+1, patch_dim],实际上经过均分后Q、K、V的维度可以理解为[1, head,num_patches+1, patch_dim//head],这里的 patch_dim//head = 768// 2=384。
紧接着分别将Q1、K1、V1;Q2、K2、V2进行Self-Attention的计算,计算之后将得到的两组结果进行拼接,接着将拼接后的结果通过WO(可学习的参数)进行融合,也是一个线性操作的过程,可以理解为一个矩阵运算,如下面的公式所示:
在这里插入图片描述
主要的代码实现如下:

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],每张图片对应一个一维向量,该向量直接通过一个全连接层来进行分类。

完整的模型框图

代码链接

你可能感兴趣的:(深度学习,深度学习,计算机视觉,pytorch)