深度学习模型之CNN（二十一）Vision Transformer(vit)网络详解

前言

原论文：An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale

Transformer最初提出是针对NLP领域的，并且在NLP领域大获成功。这篇论文也是受到其启发，尝试将Transformer应用到CV领域。通过这篇文章的实验，给出的最佳模型在ImageNet1K上能够达到88.55%的准确率（先在Google自家的JFT数据集上进行了预训练，之后再到ImageNet1K上进行迁移学习），说明Transformer在CV领域确实是有效的，而且效果还挺惊人。

模型详解

在论文中作者主要拿ResNet、ViT（纯Transformer模型）以及Hybrid（传统CNN和Transformer混合模型）三个模型进行比较。本节课主要讲解ViT，也会简单了解Hybrid模型

Vision Transformer模型

模型由三个模块组成

• Linear Projection of Flattened Patches(Embedding层)
• Transformer Encoder(图右侧有给出更加详细的结构)
• MLP Head（最终用于分类的层结构）

• 如上图所示，首先输入一张图片，之后将图片分成一个个Paches，例如性能表中表头的ViT-L16，指的是每一个Patches是16x16大小的。
• 之后将每一个Patches输入到Embedding层，也就是Linear Projection of Flattened Patches。
• 通过Embedding层之后会得到一个个向量，通常称为token。
• 紧接着在一系列token的最前面加上新的token，专门用于分类的[class]token（这里增加一个[class]token是参考RERT网络）。那么这里的class所对应的token的维度和之前得到的token是一样的，都是一个向量，且向量长度相同。
• 之后还需要加上位置的信息，如上图左所示，Patch+Position Embedding（对应表中0，1，2…9），因此此刻的一系列token已经加上[class]token以及叠加上位置信息的token。
• 将一系列token输入到Transformer Encoder（对应上图右），就是将Encoder Block重复堆叠L次（上图右中信息[灰色框框]重复L次）
• 之后再将[class]token对应的输出（在Transformer Muti-Head Attention模块中，输入几个变量，那么就会输出几个输出的变量，都是一一对应的，但此刻仅用于分类，因此只提取针对[class]token所对应的输出）通过MLP Head得到最终分类结果。

Embedding层

对于标准的Transformer模块，要求输入的是token（向量）序列，即二维矩阵[num_token, token_dim]，如下图，token0-9对应的都是向量，以ViT-B/16为例，每个token向量长度为768（对应token_dim）。

对于图像数据而言，其数据格式为[H, W, C]是三维矩阵明显不是Transformer想要的。所以需要先通过一个Embedding层来对数据做个变换。如下图所示，首先将一张图片按给定大小分成一堆Patches。以ViT-B/16为例，将输入图片(224x224)按照16x16大小的Patch进行划分，划分后会得到$( 224 / 16 )^2=196$个Patches。接着通过线性映射将每个Patch映射到一维向量中，以ViT-B/16为例，每个Patche数据shape为[16, 16, 3]通过映射得到一个长度为768的向量（后面都直接称为token）。[16, 16, 3] -> [768]

在代码实现中，直接通过一个卷积层来实现。 以ViT-B/16为例，直接使用一个卷积核大小为16x16，步距为16，卷积核个数为768的卷积来实现。通过卷积[224, 224, 3] -> [14, 14, 768]，然后把H以及W两个维度展平即可[14, 14, 768] -> [196, 768]，此时正好变成了一个二维矩阵，正是Transformer想要的。

[224,224,3]->[14,14, 768] ->[196,768]

在输入Transformer Encoder之前注意需要加上[class]token以及Position Embedding。 在原论文中，作者说参考BERT，在刚刚得到的一堆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采用的是一个可训练的参数（1D Pos. Emb.），是直接叠加在tokens上的（add），所以shape要一样。以ViT-B/16为例，刚刚拼接[class]token后shape是[197, 768]，那么这里的Position Embedding的shape也是[197, 768]。

拼接[class]token: Cat([1, 768]，[196, 768]) ->[197, 768]

叠加Position Embedding: [197, 768]->[197, 768]

对于是否使用Position Embedding作者也有做一系列对比试验，在源码中默认使用的是1D Pos. Emb.，对比不使用Position Embedding准确率提升了大概3个点，和2D Pos. Emb.比起来没太大差别。

前文说到，会为每一个Patches增加一个位置编码，其实就是一个向量，我们可以针对每一个位置上的位置编码和其他位置上的位置编码求训练相似度，即能得到下图。如下图所示（关于训练得到的位置编码，在每个位置上和其他位置的余弦相似度），这里的Patches大小是32x32的，224/32 = 7，所以大小为7x7。

理解：会在每一个token上叠加一个位置编码，那么针对每一个patches的位置编码其实就是一个向量，那么可以针对每一个位置上的位置编码和其他位置上的位置编码去求一个余弦相似度，即能得到如下图。（就是每个小小格子和其他小小格子计算相似度，所有的大格子其实都是同一张图）

Transformer Encoder层

Transformer Encoder其实就是重复堆叠Encoder Block L次，下图是Encoder Block，主要由以下几部分组成：

• Layer Norm，这种Normalization方法主要是针对NLP领域提出的，这里是对每个token进行Norm处理
• Multi-Head Attention
• Dropout/DropPath，在原论文的代码中是直接使用的Dropout层，在但rwightman（大神）实现的代码中使用的是DropPath（stochastic depth），可能后者会更好一点。
• MLP Block，如图右侧所示，就是全连接+GELU激活函数+Dropout组成，需要注意的是第一个全连接层会把输入节点个数翻4倍[197, 768] -> [197, 3072]，第二个全连接层会还原回原节点个数[197, 3072] -> [197, 768]

MLP Head层

注意，在Transformer Encoder前有个Dropout层，后有一个LayerNorm（细节都在源码中）

训练ImageNet21K或者是更大数据集时是由Linear+tanh激活函数+Linear，但是迁移到ImageNet1K上或者自己的数据上时，只有一个Linear。因此可以将MLP Head简单理解为一个全连接层。如果需要最终输出一个类别的话，需要接上softmax函数。

Vision Transformer网络结构

以ViT-B/16为例，讲解网络过程。假设输入图片为一张RGB彩色图片（224x224x3）

• 首先经过Patch Embedding层（Conv2d+Flatten）

Conv2d：卷积层的卷积核大小为16x16，stride = 16，卷积核个数为768。通过卷积层之后，输入特征矩阵的shape由[224，224，3]->[14，14，768]；

Flatten：通过展平处理，将[14，14，768]->[196，768]。

• 紧接着进行Concat处理，[class]token的shape是[1，768]（可训练的参数），与[196，768]进行concat拼接（拼接操作），得到[197，768]。再加上（相加操作）Position Embedding（可训练的参数），因此shape依旧为[197，768]；
• 通过Dropout层；
• 通过Transformer Encoder，就是将Encoder Block重复L次（12次）；
• 通过Layer Norm得到的输出shape是[197，768]；
• 再提取[class]token所对应的输出，直接对[197，768]数据进行切片，只需要提取出[class]token所对应的输出即可，切片之后得到输出[1，768];
• 之后通过MLP Head得到最终的输出（如果ImageNet21K上进行预训练的话，Pre-Logits其实就是一个全连接层+tanh激活函数；但是在ImageNet1K上或者自己的数据上时，Pre-Logits可以不要它，即MLP Head只有一个全连接层）。

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

流程概述：

• 将[224，224，3]的RGB彩色图像输入进StdConv2d卷积层中，卷积核大小为7x7，stride = 2，卷积核个数为64（[224，224，3]->[112，112，64]）；
• 通过GroupNorm、ReLU、MaxPool，数据的shape变为[56，56，64]；
• 再依次通过Stage1、Stage2、Stage3，数据shape变为[14，14，1024]，刚好由[224，224，3]->[14，14，1024]下采样16倍，刚好同ViT网络中直接通过一个卷积核大小为16x16，步距为16的下采样率一样

注意：在原来的ResNet50网络中Stage3为重复6次，但因为把stage4中的3个Block移至stage3中，所以stage3中共重复堆叠6+3=9次。

• 之后通过Patch Embedding层，这里的Conv2d和刚刚的不一样，刚刚的是16x16，stride = 16，这里为1x1，stride = 1（目的是调节特征矩阵的channel）[14，14，1024]->[14，14，768]；
• 之后同前面ViT中讲的完全一样，就不在赘述。

ViT模型搭建参数

• Layers是Transformer Encoder中重复堆叠Encoder Block的次数；
• Hidden Size是通过Embedding层后每个token的dim (向量的长度)；
• MLP size是Transformer Encoder中MLP Block第一个全连接的节点个数(是Hidden Size的四倍)；
• Heads代表Transformer中Multi-Head Attention的heads数

Model	Patch Size	Layers	Hidden Size D	MLP size	Heads	Params
ViT-Base	16x16	12	768	3072	12	86M
ViT-Large	16x16	24	1024	4096	16	307M
ViT-Huge	14x14	32	1280	5120	16	632M

论文中的网络性能对比

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