【vision transformer复现】vit整体架构

本文内容主要为vision transformer整体架构如何实现（包括pytorch 及 paddle版代码）来自笔者学习paddle vit课程笔记，经过整理得来，如有错误，请您不吝指出。

本文将解决的关键问题包括：

多头注意力为什么需要缩放，缩放因子为什么是 <math xmlns="http://www.w3.org/1998/Math/MathML"> d k \sqrt{d_k} </math>dk
为什么使用pre norm更多而不是post norm

vit整体架构如图

vit 整体架构

按照数据流向顺序依次为：

patch embedding
encoders (包含5个相同的encoder layer)
- multihead attention 多头注意力
  - qkv 线性变换
  - q, k, v 分块
  - S = q * k 矩阵点乘 matmul
  - <math xmlns="http://www.w3.org/1998/Math/MathML"> A t t e n t i o n = s o f t m a x ( S d k ) Attention = softmax(\frac{S}{\sqrt{d_k}}) </math>Attention=softmax(dk S) 经过缩放后计算注意力分数
    - 为什么要缩放？方差有放大效应，不稳定，容易导致梯度爆炸。需要将方差稳定回1.0
  - attention * v == softmax( q * k * scaler ) * v
- add & norm
  - pre norm better than post norm： pre norm的残差分支对层依赖性不高，不易导致梯度爆炸
- ffn 前馈神经网络其实就是mlp
  - fc1 -> act -> dropout ->
  - fc2 -> dropout
- add & norm
classifier head （线性层 + softmax)
- 线性层
- softmax

patch embedding

将图像分块，笔者理解为功能近似于nlp模型的tokenizer

使用conv2d卷积层来做分块比较容易，这时的卷积核大小等于步长 kenel_size = patch_size , stride = patch_size

添加cls_token 以及position_embedding ，两者均为可学习的参数，使用create_parameter()函数

cls_token是用来分类的token，初始化为常数0， position_embedding 的个数与n_patches相当，又因为concat了一个cls_token在前面，所以position_embedding的个数为 n_patches + 1,我们使用标准差为0.2的截断正态分布初始化它。

deit中还会添加distill_token，蒸馏的时候一般以截断正态分布和偏差值为0初始化蒸馏权重。

paddle版

python 复制代码

class PatchEmbedding(nn.Layer):
    def __init__(self, image_size=224, patch_size=16, in_channels=3, embed_dim=768, dropout=0.):
        super().__init__()
        self.n_patches = (image_size // patch_size) * (image_size // patch_size)
        # 如果输入的图像是长方形的，需要修改n_patches的计算方式。它应该是图像高度除以patch大小得到的patch数量乘以图像宽度除以patch大小得到的patch数量。
        self.patch_embedding = nn.Conv2D(in_channels=in_channels,
                                         out_channels=embed_dim,
                                         kernel_size=patch_size,
                                         stride=patch_size) #补丁嵌入：通过卷积运算 （Conv2D） 将图像的每个补丁转换为嵌入维度 （embed_dim）。如果将图像划分为网格（例如，7x7），您将获得 49 个补丁（标记），因此 n_patches = 49。
        # 这一步将3通道通过卷积（卷积核等同stride说明不提取特征）改变为768的嵌入维度

        self.dropout = nn.Dropout(dropout)
        # Adding class token
        # cls_token 是一种特殊标记，它位于嵌入式补丁序列之前。其目的是在变压器前向传递期间通过自注意力机制聚合来自整个图像的信息
        self.cls_token = self.create_parameter( # cls_token 是一个可学习的参数，它在模型的初始化阶段被创建，并且在训练过程中被优化。这个参数是模型的一部分，但它并不是一个层（layer），而是一个单独的参数。
            shape = [1, 1, embed_dim],  # shape 应该是与标准的token具有相同维度的。这里，class token的 shape 被设定为 [1, 1, embed_dim]，这表示有一个class token，它的特征维度与视觉token的嵌入维度（embed_dim）一致。因为ViT中的每一个patch都被映射成了一个特定维度的特征向量，class token也需要与它们保持一致。
            # 为什么shape的第一个维度是1呢？
            # 这是因为在初始化时，我们不知道实际的batch size是多少。
            # 所以我们用1来代替，之后在模型的前向传播过程中将其复制到每个样本上,
            # 当实际的batch size确定后，cls_token会通过expand方法复制到每个样本，以符合输入数据的batch size。
            
            # default_initializer 是设置这个参数初始值的方法。
            # 通常，可以使用 nn.initializer.Constant(value=0) 来将class token的初始值设置为零。
             dtype = "float32",
            default_initializer = nn.initializer.Constant(0.))
        # self.add_parameter("cls_token", self.cls_token)

        # Adding position embeddings
        self.pos_embedding = self.create_parameter(
            shape = [1, self.n_patches + 1, embed_dim], 
            # 这里的1是什么意思？
            # 这里的1表示这个embedding是共享的，即所有的图片将使用相同的位置嵌入。
            # 在模型的前向传播时，pos_embedding会自动扩展到每个batch的数据上，这个过程称为广播（broadcasting）。
            # 广播允许较小维度的Tensor在算术操作中与较大维度的Tensor兼容。
            dtype = "float32",
            default_initializer = nn.initializer.TruncatedNormal(std=.02)) # nn.initializer.TruncatedNormal(std=.02) 用截断正态分布初始化位置嵌入，这有助于模型的训练稳定性。有助于模型学习到不同位置之间的区别。
        # self.add_parameter("pos_embedding", self.pos_embedding)
    def forward(self, x):
        # [n, c, h, w]
        # TODO: forward
        # B,N,_ = x.shape 

        x = self.patch_embedding(x)
        x = x.flatten(2)
        x = x.transpose([0, 2, 1])  # Change to [B, N, D] # transpose([0, 2, 1]) 调整维度顺序
        # classtoken  and pos_embed
        class_token = self.cls_token.expand([x.shape[0], -1, -1]) # for batch [B, -1, -1]
        # expand([x.shape[0], -1, -1])表示我们想要在第一个维度（batch size维度）上复制它以匹配实际的batch size，
        # 而保持其他两个维度不变。
        # 在这个调用中，-1表示不改变那个维度的大小。
        x = paddle.concat([class_token, x], axis=1)

        print("patch_embed x shape",x.shape )
        x = x + self.pos_embedding  # Add position embeddings
        x = self.dropout(x)

        return x

pytorch版

python 复制代码

class PatchEmbedding(nn.Module):
    def __init__(self, image_size=224, patch_size=16, in_channels=3, embed_dim=768, dropout=0.):
        super().__init__()
        self.n_patches = (image_size // patch_size) * (image_size // patch_size)
        self.patch_embedding = nn.Conv2d(in_channels=in_channels,
                                         out_channels=embed_dim,
                                         kernel_size=patch_size,
                                         stride=patch_size)
        self.dropout = nn.Dropout(dropout)
        self.cls_token = nn.Parameter(torch.zeros(1, 1, embed_dim))
        self.pos_embedding = nn.Parameter(torch.zeros(1, self.n_patches + 1, embed_dim))

    def forward(self, x):
        B = x.shape[0]
        x = self.patch_embedding(x)
        x = x.flatten(2)
        x = x.transpose(1, 2)  # [B, N, D]
        cls_tokens = self.cls_token.expand(B, -1, -1)
        x = torch.cat((cls_tokens, x), dim=1)
        x = x + self.pos_embedding
        x = self.dropout(x)
        return x

encoder_layer

以下详细描述encoder的模型结构

pytorch版

python 复制代码

class EncoderLayer(nn.Module):
    def __init__(self, embed_dim, num_heads, mlp_ratio=4.0, qkv_bias=True, dropout=0., attention_dropout=0.):
        super().__init__()
        self.attn_norm = nn.LayerNorm(embed_dim)
        self.attn = Attention(embed_dim, num_heads, qkv_bias, dropout, attention_dropout)
        self.mlp_norm = nn.LayerNorm(embed_dim)
        self.mlp = Mlp(embed_dim, mlp_ratio, dropout)

    def forward(self, x):
        h = x
        x = self.attn_norm(x)
        x = self.attn(x) + h

        h = x
        x = self.mlp_norm(x)
        x = self.mlp(x) + h
        return x

multi_head attention

多头其实就是把 d_model 拆成多个头并行处理，加快速度

但是这里是vit，其实是可以使用batch norm的，但是它已经使用了patch embedding，所以把它当成序列可能会更好？

norm

pre norm better than post norm

是因为这样训练起来更稳定

注意力分数计算不至于太大，从而除以 <math xmlns="http://www.w3.org/1998/Math/MathML"> d k \sqrt{d_k} </math>dk 就容易进入（0，1）的区间，这样softmax在这个区间里近似线性，不像如果x太大的时候softmax虽然还是比relu的折角平滑，但是它的梯度变化的很快，这样就人为造成很多崎岖不平的探索区域，这完全是没有必要的。

mlp

简要来说就是两层mlp，中间使用gelu激活函数，并且添加dropout层

参考文献

百度七天目标检测实战营（vit复现）