Transformer学习记录（6）：Vision Transformer

背景

Transformer模型最初是使用在NLP中，但近几年Transformer模型在图像上的使用越来越频繁，最新的模型也出现了很多基于Transfomer的，而其中经典的是Vision Transformer(ViT)，它是用于图像分类的，这里就以ViT-B/16这个模型来学习Transformer模型是如何在图像领域使用的。

整体框架

ViT的网络结构如图所示，可以将ViT分为三个部分：

1. Embedding层：对原始的图像数据进行处理，转换为token（向量）序列，使其符合Transformer Encoder的输入要求
2. Transformer Encoders层：包含多个Encoder Block，对输入的token进行处理
3. MLP Head层：用于分类的层结构

各个部分的详解

Embedding层

对于标准的Transformer模块，要输入的是token（向量）序列，也就是二维矩阵[num_token，token_dim(向量的长度)]，以ViT-B/16为例，token的长度为768。

而对于图像数据而言，其格式为[H,W,C]，显然不是Transformer模块需要的，因此需要通过一个Embedding层来对数据格式进行转换。

操作过程如下图所示，以ViT-B/16为例，首先将输入图像（224x224）切分为196个16x16的Patch，这些Patch的格式为[16,16,3]。接下来将这些Patch进行线性映射，也就是转换为一维的向量：[16,16,3]-->[768]。

在代码实现中，直接使用一个卷积层来实现。以ViT-B/16为例，使用一个卷积核为16x16，步距为16，卷积核个数为768的卷积层来实现。通过卷积层会得到[224,224,3]-->[14,14,768]，然后在通过Flatten操作将[H,W]展平即可[14, 14, 768] -> [196, 768]，此时正好为二维矩阵，符合Transformer模块的需要。

在输入到Transformer模块前，需要加上[class]token(token所属的类别)与位置编码Position Embedding。

• [class]token：在刚刚得到的一堆tokens中插入一个专门用于分类的[class]token，这个[class]token是一个可训练的参数，数据格式和其他token一样都是一个向量，以ViT-B/16为例，就是一个长度为768的向量，与之前从图片中生成的tokens拼接在一起Cat([1, 768], [196, 768]) -> [197, 768]，增加了一个向量。
• Position Embedding：就是之前Transformer中讲到的Positional Encoding，这里的Position Embedding采用的是一个可训练的参数，这里是直接叠加（add）到token上。以ViT-B/16为例，刚刚拼接[class]token后shape是[197, 768]，那么这里的Position Embedding的shape也是[197, 768]。

Embedding 层的整体结构如图所示

Transformer Encoder层

Transformer Encoder其实就是重复堆叠Encoder Block L次，Encoder Block由以下部分构成：

• Layer Norm：对每个token进行标准化处理
• Multi-Head Attention：链接
• Dropout/DropPath：防止过拟合，在原论文的代码中是直接使用的Dropout层，但也有人使用DropPath
• MLP Block：如图右侧所示，由全连接层(Liner)+GELU激活函数+DropOut层构成，需要注意的是第一个全连接层会将节点个数翻4倍[197,768]-->[197,3072]，而第二个全连接层会将节点个数恢复原样[197,3072]-->[197,768]，这样经过MPL Block不会修改数据的结构。

通过Transformer Encoder后输出的shape和输入的shape是保持不变的，以ViT-B/16为例，输入的是[197, 768]输出的还是[197, 768]

而在代码实现中，会在Transformer模块前面加一个Dropout层，后面加一个Layer Norm层。

MLP Head

由于只需要分类的信息，因此只需要将[class]token生成的对应结果提取出来即可，即[197, 768]中抽取出[class]token对应的[1, 768]，最后在使用MLP Head得到最终的分类结果。

在论文中，训练ImageNet21K（大型数据集）时，MLP Head是由Linear+tanh激活函数+Linear组成。但是迁移到ImageNet1K上或者你自己的数据上时，只用一个Linear即可。

以ViT-B/16为例

不同种类的ViT

论文中提供了三种ViT模型的参数

• Layers:Transformer模块中Encoder模块的个数
• Hidden Size：经过Embedding层后每个token的长度
• MLP Size:MLP Block中第一个全连接层的输出节点个数，是Hidden Size的四倍。
• Heads：Multi-Head Attention的heads数。

代码实现

Embedding层

在代码实现时，要注意输入数据的格式，一般的格式为(B,C,H,W)，但后面需要的数据格式为(B,HW,C)，即二三维需要调换，因此需要进行转置操作。

class Patch_Embedding(nn.Module):
    def __init__(self, image_size=224, in_channel=3, embed_dim=768, patch_size=16, dropout=0.1):
        """
        嵌入层
        @param image_size: 图像大小
        @param in_channel: 输入通道数
        @param embed_dim: 输出向量的维度
        @param patch_size: patch的大小
        @param dropout: dropout操作
        """
        super(self,Patch_Embedding).__init__()
        self.patcher = nn.Sequential(
            nn.Conv2d(
                in_channels=in_channel,
                out_channels=embed_dim,
                kernel_size=patch_size,
                stride=patch_size
            ),
            nn.Flatten(2)  # 展平，(B,C,H,W)-->(B,C,H*W)
        )
        self.class_token = nn.Parameter(torch.randn(size=(1, 1, embed_dim)))  # 随机初始化class_token
        self.num_patch = (image_size // patch_size) * (image_size // patch_size)  # patch的个数
        self.position_embedding = nn.Parameter(torch.randn(size=(1,self.num_patch+1,embed_dim)))  # 随机初始化位置编码
        self.dropout = nn.Dropout(p=dropout)

    def forward(self,x):
        B,C,H,W = x.shape
        x = self.patcher(x)  # 卷积并展平
        x = x.permute(0,2,1)  # 交换 (B,C,HW)-->(B,HW,C)
        x = torch.cat((x,self.class_token),dim=1)  # 与class_token拼接
        x = x + self.position_embedding  # 与位置编码相加
        return x

Transformer Encoder层

有两种实现方法：1.使用pytorch中的API；2.手动搭建

① 使用pytorch中的API：
在pytorch中，有两个API：nn.TransformerEncoderLayer与nn.TransformerEnoder

其中TransformerEncoderLayer相当于图中的Encoder Block，而TransformerEncoder则包含所有Encoder Block。
因此使用时需要先构造出TransformerLayer对象，设置好参数，再使用TransformerLayer对象出构造TransformerEncoder。

参数：

• d-model(int)：输入数据的特征维度，在vit中也就是第三维的数据个数
• nhead(int)：多头自注意力机制中，head的个数
• dim_feedforward(int)：前馈网络中隐藏层的维度，默认为2048
• dropout(float)：dropout的比例，默认为0.1
• activation(str)：使用的激活函数，如‘“relu”与“gelu”，默认为“relu”
• layer_norm_eps (float) ：在layer normalization中使用的，默认为1e-5
• batch_first (bool)：若为True，则在输入与输出的数据中，batch在第一个维度，反之在第二个维度。默认为Flase
• norm_first (bool)：是否在执行起始的layer normalization。若为True则执行，默认为False
• bias (bool)：Liner与layer normalization是否使用bias。若为True则使用，默认为Flase。

参数：

• encoder_layer (TransformerEncoderLayer) ：传入TransformerEncoderLayer对象
• num_layers(int)：包含多少个Encoder Block
• norm (Optional[Module])： layer normalization使用的模块，默认为None

self.TransformerEnoderLayer = nn.TransformerEncoderLayer(  # Encoder Block
     d_model=embed_dim,
     nhead=num_head,
     dropout=dropout,
     activation=activation,
     batch_first = True,
     norm_first = True,
)
self.TransformerEncoder = nn.TransformerEncoder(self.TransformerEnoderLayer,num_layers)  # Encoder

MLP层

CNN与Transformer的混合网络：Hybrid

Hybrid将传统CNN的特征提取和Transformer进行结合。
下图是以ResNet50作为特征提取器的Hybrid的网络结构，与ViT区别在于开始阶段会使用ResNet50网络进行特征提取。但这里的ResNet50与普通的ResNet50不同。

首先这里的R50的卷积层采用的StdConv2d不是传统的Conv2d，然后将所有的BatchNorm层替换成GroupNorm层。
而且在原Resnet50网络中，stage1重复堆叠3次，stage2重复堆叠4次，stage3重复堆叠6次，stage4重复堆叠3次。但是在这里的R50中，stage4的3次被集成到stage3中，因此stage3堆叠了9次。

通过R50 Backbone进行特征提取后，得到的特征矩阵为[14,14,1024]，然后输入到Embedding层，要注意的是这个Embedding层中的卷积变为卷积核为1x1，步长为1，只是调整了channel为768。

后面的处理与ViT一样。