Vision Transformer原理分析

Vision Transformer原理分析

Vision Transformer是2021年Google团队提出的将Transformer应用在图像分类的模型，因为其模型简单、效果好、可扩展性强，成为CV领域的里程碑著作，也引爆了后续相关研究。

---

1、Motivation

在2020至2021年间，CV领域被CNN占据，NLP领域Transformer成为标配。许多工作致力于将NLP领域的思路推广到CV领域，大体可分为两类：第一、注意力机制与CNN结合。第二、在整体结构不变的情况下注意力机制替换CNN某些结构，比如将Self-Attention应用于语义分割领域，产生了诸如Non-local、DANet、CCNet等优秀的语义分割模型。

但是这些方法始终还是需要依赖于CNN来提取图像特征。因此，谷歌研究人员思考，可不可以完全不依赖CNN，直接用Transformer作为编码器提取图像特征，并使其成为如ResNet一样的CV领域提取图像特征的标配。

最直接的方法就是将NLP中的做法搬过来，NLP中提取文本特征的预训练模型的标配恰好就是谷歌发布的预训练模型BERT。因此，ViT就是在思考如何将BERT拓广到CV领域。对比可以发现，ViT实际上就是输入形式（需要将图像处理成NLP中的embedding形式）、预训练任务（去除了BERT中的NSP以及mask LM任务，改为图像分类任务）相较于BERT有所差别。

2、文章结论

ViT原论文中最核心的结论是，当拥有足够多的数据进行预训练的时候，ViT的表现就会超过CNN，突破Transformer缺少归纳偏置的限制，可以在下游任务中获得较好的迁移效果。但是当训练数据量不够大的时候，ViT的表现通常比同等大小的ResNet要差一些，因为Transformer和CNN相比缺少归纳偏置（inductive bias），即一种先验知识。

3、模型细节

Vision Transformer的网络架构图如下：

第1部分：模型输入

假设模型输入是一张(512, 512, 3)尺度的三通道图像，初始值=torch.Size([1, 3, 512, 512])。

第2部分：Linear Projection of Flatten Patches

将整张(512, 512, 3)大小的图像划分成(16, 16)个patches，每个patches的尺度即为(32, 32, 3)，此时torch.Size([1, 3, 512, 512])可以转化成torch.Size([1, 3, 16, 32, 16, 32])。

将(16, 16)合并成一个维度，代表所有的patches；(3, 32, 32)合并成一个维度，代表每个patches所含有的像素点数量，此时可以转换得到torch.Size([1, 256, 3072])。

经过一个线性层nn.Linear(3072, 1024)，转换成torch.Size([1, 256, 1024])。再经过gelu激活层、LayerNorm归一化层，得到torch.Size([1, 256, 1024])。

第3部分：Patches+Position Embedding

构建一个nn.Embedding(config.max_position_embeddings=500, config.hidden_size=1024)的字典，用于将一个索引数字转变成一个指定维度的向量，这里字典的值也是需要训练的参数。

对于(16, 16)=256个patches，生成tensor([0, 1, 2, 3, …, 254, 255])，按索引值从Embedding字典中取出对应的（编码得到的）1024维度的向量，即torch.Size([1, 256, 1024])。

将位置编码向量与前面得到的图像特征相加，得到torch.Size([1, 256, 1024])。再经过LayerNorm归一化层和Dropout层，得到torch.Size([1, 256, 1024])。

第4部分：Transformer Encoder — Block叠加

Transformer Encoder包含一系列的Block（即不断重复第5部分到第8部分），每个Block结构都相同，输入维度为torch.Size([1, 256, 1024])，输出维度为torch.Size([1, 256, 1024])。

第5部分：Transformer Encoder — Block — Multi-Head Attention

输入维度为torch.Size([1, 256, 1024])，定义三个线性层nn.Linear(1024, 1024)，得到三个Q、K、V矩阵，维度均为torch.Size([1, 256, 1024])。

转化成多头形式，Q、K、V矩阵维度均转换成torch.Size([1, 16, 256, 64])，采用16个头。

在多头形式下进行自注意力运算，得到数据维度为torch.Size([1, 16, 256, 64])，然后再恢复成torch.Size([1, 256, 1024])的形式。

第6部分：Transformer Encoder — Block — 第一次残差连接

将自注意力机制得到的torch.Size([1, 256, 1024])，与进行自注意力机制之前的数据，即torch.Size([1, 256, 1024])累加，相当于残差连接。再经过LayerNorm归一化层，得到torch.Size([1, 256, 1024])。

第7部分：Transformer Encoder — Block — MLP

输入维度为torch.Size([1, 256, 1024])，定义MLP层nn.Linear(1024, 1024)，转化得到torch.Size([1, 256, 1024])。再经过gelu激活层，得到torch.Size([1, 256, 1024])。

第8部分：Transformer Encoder — Block — 第二次残差连接

将MLP得到的torch.Size([1, 256, 1024])，与进行MLP之前的数据torch.Size([1, 256, 1024])累加，相当于残差连接。再经过LayerNorm归一化层，得到torch.Size([1, 256, 1024])。

第9部分：图像分类输出结果

经过一系列Transformer Block模块的编码，最终提取到的特征维度为torch.Size([1, 256, 1024])。
在patches维度上进行平均，相当于将每个patch提取到的特征融合，得到torch.Size([1, 1024])。
最后经过线性层nn.Linear(1014, 10)，得到输出结果torch.Size([1, 10])。

4、相关思考

（1）对nn.Embedding函数的理解，以及其与nn.Linear函数有什么区别？

解答博客：https://blog.csdn.net/qq_43391414/article/details/120783887

torch.nn.Embedding是用来将一个索引数字变成一个指定维度的向量的，比如数字1变成一个128维的向量，数字2变成另一个128维度的向量，不过这128维的向量并不是永恒不变的。然后这些128维的向量会参与模型训练并且得到更新，从而数字1会有一个更好的128维向量的表示。显然，这非常像全连接层，所以很多人说，Embedding层是全连接层的特例。

Embedding和Linear几乎是一样的，区别就在于：输入不同，一个是输入数字，后者是输入one-hot向量。可以这么说，假设有K个不同的1维数值，每个都需要映射成N维的向量：
Embedding(K, N) = One-hot(K) + Linear(K, N)

5、源码

如果需要源代码，可以去我的主页寻找项目链接，以上代码和实验结果都由本人亲自实验得到：
https://blog.csdn.net/Twilight737