Vision Transformer 理论详解

Vison Transformer 介绍

Vison Transformer论文- An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale, 论文由Google 发表于2020 CVPR上。

VIT模型的效果

从图中可以看出VIT不仅在NLP领域应用，在CV领域的效果也非常好。

VIT模型架构

• 首先输入一张图片，然后将图片分成一个个patches.对于模型VIT-L/16,这里的16指的是每个patches的大小是16x16的。
• 然后将所有patches输入到Embedding层,也就是这里的“Linear Projection of Flattened Patches” ,通过Embedding层之后我们就可以得到一个个向量，通常称为Token。对应图上每个Patch对应得到一个Token。在Token最前面增加了一个带*的Token，专门用来分类的class token.
• 同时需要增加一个位置信息，在之前博客中我们讲到它的结构中是没有考虑到位置信息的，因此加上了一个Position Embedding，对应上图Embedding层得到的0,1,2,3,4,5,6,7,8,9标识的向量。
• 然后将得到的一系列Token包括class token以及position ,然后输入到Transformer Encoder中。Transformer Encoder的结构如上图所示，网络将Encoder Block重复堆叠L次。
• 紧接着将class token的所对应的输出，输入到MLP Head中，得到我们最终的分类结果。(我们在之前的博客讲到的muti-head attention模块，输入几个变量就能得到几个输出，这些都是一一对应的。由于这里我们只是分类，所以只提取class token经过Transformer Encoder对应的输出)。

根据VIT模型的网络结构，我们将模型分为3个部分：

1. Embedding 层(Linear Projection of Flattened Patches)
2. Transformer Encoder层
3. MLP Head(最终用于分类的层结构)

Embedding 层

• 对于标准的Transformer模块，要求输入的是token向量序列，即二维矩阵[num_token,token_dim]， 如下图，token0-9对应的都是向量
  
  在代码的实现中，直接通过一个卷积层来实现,以VIT-B/16为例，使用卷积核大小为16x16，stride为16，卷积核个数为768,这里的768对应的是token的维度。对输入为224 x 224 x 3的图片，经过卷积后得到14 x 14 x 768的特征矩阵，然后将宽高纬度信息展平，得到输出特征矩阵196 x 768，对应为196个纬度为768的token。这个过程就对应于上图的Embedding层。
• 紧接着加上一个class token，通过初始化一个可训练的参数，纬度为1x768。与我们通过将图片的patches，通过embedding层之后得到196x768的特征矩阵进行concate拼接，就得到了197x768的特征矩阵。得到的特征矩阵在输入Transformer Embedding之前需要叠加一个Position Embedding，它的纬度为197 x 768 。

Position Embedding

• 如果不去使用Position Embedding会有什么样的结果呢？ 论文进行了实验对比，如下
  
  对比发现没有使用位置编码，准确率为61.32，使用了1-D,2-D,Rel位置编码都能达到64个点，使用了位置编码可以提升大概3个点。但位置编码的不同形式对最终效果影响不大，所以源码中一般使用的是1 D位置编码。
• 我们会在每个token上加上一个位置编码，假设输入图片大小为224 x 224，patches的大小为32, 我们会得到7x7个token，在每个token上会加一个位置编码。我们对各个位置的位置编码求余弦相似度，可以得到如下的图形：
  
  可以看出位置编码与它所在行和列的余弦相似度比较大。

Transformer Encoder层

Transformer Encoder就是将我们的Encoder Block重复堆叠L次。对于Encoder Block的说明如下：

• 首先，将Embedding层的输出，经过Layer Norm（层归一化），紧接着通过Mutli-Head Attention,然后通过dropout层或者DropPath层（一般使用DropPath的效果会好点）得到的输出与捷径分支的输出进行Add相加操作。
• 紧接着将输出通过Layer Norm操作，然后通过MLP Block,紧接着通过Dropout或者DropParh的输出与捷径分支的输出进行Add相加操作。
• 其中MLP Block的结构如下图，通过全连接层，GELU激活函数，Dropout，全连接层以及最后Dropout层得到输出。注意第一个全连接层的节点个数3072，是输入节点个数768的4倍,通过第二个全连接层之后节点个数又变回了768。
  

MLP Head

• 上面通过Transformer Encoder后输出的shape和输入的shape是保持不变的，以ViT-B/16为例，输入的是[197, 768]输出的还是[197, 768]。注意，在Transformer Encoder前有个Dropout层，Transformer Encoder后其实还有一个Layer Norm没有画出来。后面通过手绘的ViT的模型可以看到详细结构。
• 这里我们只是需要分类的信息，所以我们只需要提取出[class]token生成的对应结果就行，即[197, 768]中抽取出[class]token对应的[1, 768]。接着我们通过MLP Head得到我们最终的分类结果。MLP Head原论文中说在训练ImageNet21K时是由Linear+tanh激活函数+Linear组成。但是迁移到ImageNet1K上或者你自己的数据上时，只用一个Linear即可。后面如果需要得到每个类别的概率的话，需要接上一个softmax激活函数。
  
  以上就是VIT模型网络结构的讲解。

---

Vision Transformer详细的网络结构

为了方便理解，博主：太阳花小绿豆手绘了更详细的图(以ViT-B/16为例)：

• 图中MLP Head如果是对于自己的数据集进行训练的话，Pre-Logits可以去掉，直接通过Linear全连接层输出结果。如果训练ImageNet21K时是，则Pre-Logits对应的是Linear+tanh激活函数

Hybrid模型详解

在论文4.1章节的Model Variants中有比较详细的讲到Hybrid混合模型，就是将传统CNN特征提取和Transformer进行结合。下图绘制的是以ResNet50作为特征提取器的混合模型，但这里的Resnet与之前讲的Resnet有些不同。首先这里的R50的卷积层采用的StdConv2d不是传统的Conv2d，然后将所有的BatchNorm层替换成GroupNorm层。在原Resnet50网络中，stage1重复堆叠3次，stage2重复堆叠4次，stage3重复堆叠6次，stage4重复堆叠3次，但在这里的R50中，把stage4中的3个Block移至stage3中，所以stage3中共重复堆叠9次。

通过R50 Backbone进行特征提取后，得到的特征矩阵shape是[14, 14, 1024]，接着再输入Patch Embedding层，注意Patch Embedding中卷积层Conv2d的kernel_size和stride都变成了1，只是用来调整channel。后面的部分和前面ViT中讲的完全一样，就不在赘述。

下表是论文用来对比ViT，Resnet（和刚刚讲的一样，使用的卷积层和Norm层都进行了修改）以及Hybrid模型的效果。通过对比发现，在训练epoch较少时Hybrid优于ViT，但当epoch增大后ViT优于Hybrid。

ViT模型搭建参数

在论文的Table1中有给出三个模型（Base/ Large/ Huge）的参数，在源码中除了有Patch Size为16x16的外还有32x32的。其中的Layers就是Transformer Encoder中重复堆叠Encoder Block的次数，Hidden Size就是对应通过Embedding层后每个token的dim（向量的长度），MLP size是Transformer Encoder中MLP Block第一个全连接的节点个数（是Hidden Size的四倍），Heads代表Transformer中Multi-Head Attention的heads数。

参考：Vision Transformer详解