Transformer学习记录（6）：Vision Transformer

news2024/12/22 21:01:32

背景

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

整体框架

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

Embedding层：对原始的图像数据进行处理，转换为token（向量）序列，使其符合Transformer Encoder的输入要求
Transformer Encoders层：包含多个Encoder Block，对输入的token进行处理
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]。

在这里插入图片描述

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数。

在这里插入图片描述

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一样。

在这里插入图片描述