paper:AN IMAGE IS WORTH 16X16 WORDS:TRANSFORMERS FOR IMAGE RECOGNITION AT SCALE

论文链接

• 🔗：AN IMAGE IS WORTH 16X16 WORDS:TRANSFORMERS FOR IMAGE RECOGNITION AT
• 🔗：李沐论文精读系列二：Vision Transformer、MAE、Swin-Transformer

说在前面

这篇文章挑战了自从2012年AlexNet提出以来卷积神经网络在计算机视觉里绝对统治的地位。结论是如果在足够多的数据上做预训练，也可以不需要卷积神经网路，直接使用标准的transformer也能够把视觉问题解决的很好。它打破了CV和NLP在模型上的壁垒，开启了CV的一个新时代，推进了多模态领域的发展。

paperswithcode可以查询现在某个领域或者说某个数据集表现最好的一些方法有哪些。图像分类在ImageNet数据集上排名靠前的全是基于Vision Transformer。

对于目标检测任务在COCO数据集上，排名靠前都都是基于Swin Transformer。Swin Transformer是ICCV 21的最佳论文，可以把它想象成一个多尺度的Vit（Vision Transformer）。

在其他领域（语义分割、实例分割、视频、医疗、遥感），基本上可以说Vision Transformer将整个视觉领域中所有的任务都刷了个遍。

1. 引言

1.1 Vision Transformer的一些有趣特性

作者介绍的另一篇论文：《Intriguing Properties of Vision Transformer》

• 图a表示的是遮挡，在这么严重的遮挡情况下，不管是卷积神经网络，人眼也很难观察出图中所示的是一只鸟
• 图b表示数据分布上有所偏移，这里对图片做了一次纹理去除的操作，所以图片看起来比较魔幻
• 图c表示在鸟头的位置加了一个对抗性的patch
• 图d表示将图片打散了之后做排列组合

上述例子中，卷积神经网络很难判断到底是一个什么物体，但是对于所有的这些例子Vision Transformer都能够处理的很好。

1.2 标题

一张图片等价于很多1616大小的单词。为什么是1616的单词？把图片分割成很多方格patch的形式，每一个方格的大小都是1616，那么这张图片就相当于是很多1616的patch组成的整体

1.3 摘要

在VIT之前，self-attention在CV领域的应用很有限，要么和卷积一起使用，要么就是把CNN里面的某些模块替换成self-attention，但是整体架构不变。

这篇文章证明了，在图片分类任务中，只使用纯的Vision Transformer结构直接作用于一系列图像块，也可以取的很好的效果（最佳模型在ImageNet1K上能够达到88.55%的准确率）。尤其是当在大规模的数据上面做预训练然后迁移到中小型数据集（ImageNet、CIFAR-100、VATB）上面使用的时候，Vision Transformer能够获得跟最好的卷积神经网络相媲美的结果。Transformer的另外一个好处：它只需要更少的训练资源，而且表现还特别好。

作者这里指的少的训练资源是指2500天TPUv3的天数。这里的少只是跟更耗卡的模型去做对比。

1.4 引言

Transformer在NLP领域的应用

基于self-attention的模型架构，特别是Transformer，在NLP领域几乎成了必选架构。现在比较主流的方式，就是先去一个大规模的数据集上去做预训练，然后再在一些特定领域的小数据集上面做微调。多亏了Transformer的高效性和可扩展性，现在已经可以训练超过1000亿参数的大模型（GPT3）。随着模型和数据集的增长，还没有看到性能饱和的现象。

• 很多时候不是一味地扩大数据集或者说扩大模型就能够获得更好的效果的，尤其是当扩大模型的时候很容易碰到过拟合的问题，但是对于transformer来说目前还没有观测到这个瓶颈
• 微软和英伟达联合推出了一个超级大的语言生成模型Megatron-Turing，它已经有5300亿参数了，还能在各个任务上继续大幅度提升性能，没有任何性能饱和的现象

将transformer运用到视觉领域的难处

Transformer在做自注意力的时候是两两互相的，这个计算复杂度是跟序列的长度呈平方倍的。目前一般在自然语言处理中，硬件能支持的序列长度一般也就是几百或者是上千（比如说BERT的序列长度也就是512）。

首先要解决的是如何把一个2D的图片变成一个1D的序列（或者说变成一个集合）。最直观的方式就是把每个像素点当成元素，将图片拉直放进transformer里，看起来比较简单，但是实现起来复杂度较高。

一般来说在视觉中训练分类任务的时候图片的输入大小大概是224224，如果将图片中的每一个像素点都直接当成元素来看待的话，序列长度就是224224=50176个像素点，这个大小就相当于是BERT序列长度的100倍。这还仅仅是分类任务，对于检测和分割，现在很多模型的输入都已经变成600600或者800800或者更大，计算复杂度更高，所以在视觉领域，卷积神经网络还是占主导地位的，比如AlexNet或者是ResNet。

将自注意力用到机器视觉的相关工作

受NLP启发，很多工作研究如何将自注意力用到机器视觉中。一些工作是说把卷积神经网络和自注意力混到一起用；另外一些工作就是整个将卷积神经网络换掉，全部用自注意力。这些方法其实都是在干一个事情：因为序列长度太长，所以导致没有办法将transformer用到视觉中，所以就想办法降低序列长度

Non-local Neural Networks（CVRP，2018）：将网络中间层输出的特征图作为transformer输入序列，降低序列的长度。比如ResNet50在最后一个Stage的特征图size=14×14，把它拉平，序列元素就只有196了，这就在一个可以接受的范围内了。

《Stand-Alone & Self-Attention in Vision Models》（NeurIPS，2019）：使用孤立注意力Stand-Alone和 Axial-Attention来处理。具体的说，不是输入整张图，而是在一个local window（局部的小窗口）中计算attention。窗口的大小可以控制，复杂度也就大大降低。（类似卷积的操作）

《Axial-DeepLab: Stand-Alone Axial-Attention for Panoptic Segmentation》（ECCV，2020a）：

• 孤立自注意力：不使用整张图，就用一个local window（局部的小窗口），通过控制这个窗口的大小，来让计算复杂度在可接受的范围之内。这就类似于卷积操作（卷积也是在一个局部的窗口中操作的）
• 轴自注意力：之所以视觉计算的复杂度高是因为序列长度N=H*W，是一个2D的矩阵，将图片的这个2D的矩阵想办法拆成2个1D的向量，所以先在高度的维度上做一次self-attention（自注意力），然后再在宽度的维度上再去做一次自注意力，相当于把一个在2D矩阵上进行的自注意力操作变成了两个1D的顺序的操作，这样大幅度降低了计算的复杂度

这些模型虽然理论上是非常高效的，但事实上这个自注意力操作都是一些比较特殊的自注意力操作，无法在现在的硬件上进行加速，所以就导致很难训练出一个大模型。因此在大规模的图像识别上，传统的残差网络还是效果最好的。

所以，自注意力早已经在计算机视觉里有所应用，而且已经有完全用自注意力去取代卷积操作的工作了。本文是被transformer在NLP领域的可扩展性所启发，直接应用一个标准的transformer作用于图片，尽量做少的修改。

vision transformer如何解决序列长度的问题

To do so, we split an image into patches and provide the sequence of linear embeddings of these patches as an input to a Transformer. Image patches are treated the same way as tokens (words) in an NLP application. We train the model on image classification in supervised fashion.

• vision transformer将一张图片打成了很多的patch，每一个patch是16*16
• 假如图片的大小是224224，则sequence lenth（序列长度）就是N=224224=50176，如果换成patch，一个patch相当于一个元素的话，有效的长宽就变成了224/16=14，所以最后的序列长度就变成了N=14*14=196，对于普通的transformer来说是可以接受的
• 将每一个patch当作一个元素，通过一个全连接层就会得到一个linear embedding，这些就会当作输入传给transformer。这时候一张图片就变成了一个一个的图片块了，可以将这些图片块当成是NLP中的单词，一个句子中有多少单词就相当于是一张图片中有多少个patch，这就是题目中所提到的一张图片等价于很多16*16的单词

有监督的训练

本文训练vision transformer使用的是有监督的训练。为什么要突出有监督？因为对于NLP来说，transformer基本上都是用无监督的方式训练的，要么是用language modeling，要么是用mask language modeling，都是用的无监督的训练方式。但是对于视觉来说，大部分的基线（baseline）网络还都是用的有监督的训练方式去训练的.

前人最相关的工作

本文把视觉当成自然语言处理的任务去做的，尤其是中间的模型就是使用的transformer encoder，跟BERT完全一样。这么简单的想法，之前其实也有人想到过去做，跟本文的工作最像的是一篇ICLR 2020的paper

• 这篇论文是从输入图片中抽取2*2的图片patch
• 为什么是22？因为这篇论文的作者只在CIFAR-10数据集上做了实验，而CIFAR-10这个数据集上的图片都是3232的，所以只需要抽取22的patch就足够了，1616的patch太大了
• 在抽取好patch之后，就在上面做self-attention

从技术上而言这就是Vision Transformer，但是本文的作者认为二者的区别在于，本文的工作证明了如果在大规模的数据集上做预训练的话，那么就能让一个标准的Transformer，不用在视觉上做任何的更改或者特殊的改动，取得比现在最好的卷积神经网络差不多或者还好的结果。

This model is very similar to ViT, but our work goes further to demonstrate that large scale pre-training makes vanilla transformers competitive with (or even better than) state-of-the-art CNNs.

这篇文章的主要目的就是说，Transformer在Vision领域能够扩展的有多好，就是在超级大数据集和超级大模型两方的加持下，transformer也能在视觉中起到很好的效果

ViT和CNN网络使用效果的比较

在中型大小的数据集上（比如说ImageNet）上训练的时候，如果不加比较强的约束，ViT的模型其实跟同等大小的残差网络相比要弱一点。

作者对此的解释是：transformer跟CNN相比，缺少了一些CNN所带有的归纳偏置（inductive bias，是指一种先验知识或者说是一种提前做好的假设）。

CNN的归纳偏置一般来说有两种：

• locality：CNN是以滑动窗口的形式一点一点地在图片上进行卷积的，所以假设图片上相邻的区域会有相邻的特征，靠得越近的东西相关性越强；
• translation equivariance（平移等变性）：写成公式就是f(g(x))=g(f(x))，不论是先做 g 这个函数，还是先做 f 这个函数，最后的结果是不变的；其中f代表卷积操作，g代表平移操作。因为在卷积神经网络中，卷积核就相当于是一个模板，不论图片中同样的物体移动到哪里，只要是同样的输入进来，然后遇到同样的卷积核，那么输出永远是一样的

一旦神经网络有了这两个归纳偏置之后，他就拥有了很多的先验信息，所以只需要相对较少的数据就可以学习一个相对比较好的模型。但是对于transformer来说，它没有这些先验信息，所以它对视觉的感知全部需要从这些数据中自己学习。

为了验证这个假设， 作者在更大的数据集（ImageNet 22k数据集， 14M个样本&JFT 300M数据集， 300M个样本）上做了预训练，然后发现在有足够的数据做预训练的情况下，Vit能够获得跟现在最好的残差神经网络相近或者说更好的结果

Our Vision Transformer (ViT) attains excellent results when pre-trained at sufficient scale and transferred to tasks with fewer datapoints. When pre-trained on the public ImageNet-21k dataset or the in-house JFT-300M dataset, ViT approaches or beats state of the art on multiple image recognition benchmarks. In particular, the best model reaches the accuracy of 88.55% on ImageNet, 90.72% on ImageNet-ReaL, 94.55% on CIFAR-100, and 77.63% on the VTAB suite of 19 tasks

上面VTAB也是作者团队所提出来的一个数据集，融合了19个数据集，主要是用来检测模型的稳健性，从侧面也反映出了VisionTransformer的稳健性也是相当不错的。

2.结论

这篇论文的工作是直接拿NLP领域中标准的Transformer来做计算机视觉的问题，跟之前用自注意力的那些工作的区别在于，除了在刚开始抽图像块的时候，还有位置编码用了一些图像特有的归纳偏置，除此之外就再也没有引入任何图像特有的归纳偏置了。这样的好处就是可以直接把图片当做NLP中的token，拿NLP中一个标准的Transformer就可以做图像分类了。

当这个简单而且扩展性很好的策略和大规模预训练结合起来的时候效果出奇的好：Vision Transformer在很多图像分类的benchmark上超过了之前最好的方法，而且训练起来还相对便宜

作者对未来的展望：

• Vit不只做分类，还有检测和分割
  • DETR：目标检测的一个力作，相当于是改变了整个目标检测之前的框架
  • 在Vit出现短短的一个半月之后，2020年12月出来了一个叫Vit-FRCNN的工作，将Vit用到检测上面了
  • 2020年12月有一篇SETR的paper将Vit用到分割里了
  • 3个月之后Swin Transformer横空出世，它将多尺度的设计融合到了Transformer中，更加适合做视觉的问题，真正证明了Transformer是能够当成一个视觉领域的通用骨干网络
• 探索一下自监督的预训练方案：因为在NLP领域，所有大的transformer全都是用自监督的方式训练的，Vit这篇paper也做了一些初始实验，证明了用这种自监督的训练方式也是可行的，但是跟有监督的训练比起来还是有不小的差距的
• 将Vision Transformer变得更大，有可能会带来更好的结果：过了半年，同样的作者团队又出了一篇paper叫做Scaling Vision Transformer，就是将Transformer变得很大，提出了一个Vit-G，将ImageNet图像分类的准确率提高到了90以上了

3.相关工作

transformer在NLP领域的应用

自从2017年transformer提出做机器翻译以后，基本上transformer就是很多NLP任务中表现最好的方法。现在大规模的transformer模型一般都是先在一个大规模的语料库上做预训练，然后再在目标任务上做一些细小的微调，这当中有两系列比较出名的工作：BERT和GPT。BERT是用一个denoising的自监督方式（其实就是完形填空，将一个句子中某些词划掉，再将这些词预测出来）；GPT用的是language modeling（已经有一个句子，然后去预测下一个词是什么，也就是next word prediction）做自监督。这两个人物其实都是人为定的，语料是固定的，句子也是完整的，只是人为划掉其中的某些部分或者把最后的词拿掉，然后去做完形填空或者是预测下一个词，所以这叫自监督的训练方式.

自注意力在视觉中的应用

视觉中如果想简单地在图片上使用自注意力，最简单的方式就是将每一个像素点当成是一个元素，让他们两两做自注意力就好了，但是这个是平方复杂度，所以很难应用到真实的图片输入尺寸上。像现在分类任务的224*224，一个transformer都很难处理，更不用提人眼看的比较清晰的图片了，一般是1k或者4k的画质，序列长度都是上百万，直接在像素层面使用transformer的话不太现实，所以如果想用transformer就一定得做一些近似

• 复杂度高是因为用了整张图，所以序列长度长，那么可以不用整张图，就用local neighborhood（一个小窗口）来做自注意力，那么序列长度就大大降低了，最后的计算复杂度也就降低了
• 使用Sparse Transformer，就是只对一些稀疏的点去做自注意力，所以只是一个全局注意力的近似
• 将自注意力用到大小不同的block上，或者说在极端的情况下使用轴注意力（先在横轴上做自注意力，然后再在纵轴上做自注意力），序列长度也是大大减小的

这些特制的自注意力结构其实在计算机视觉上的结果都不错，表现都是没问题的，但是它们需要很复杂的工程去加速算子，虽然在CPU或者GPU上跑得很快或者说让训练一个大模型成为可能

跟本文工作最相似的是一篇ICLR2020的论文，区别在于Vision Transformer使用了更大的patch，更大的数据集

在计算机视觉领域还有很多工作是把卷积神经网络和自注意力结合起来的，这类工作相当多，而且基本涵盖了视觉里的很多任务（检测、分类、视频、多模态等）

还有一个工作和本文的工作很相近，叫imageGPT:

• GPT是用在NLP中的，是一个生成性的模型。imageGPT也是一个生成性模型，也是用无监督的方式去训练的，和Vit相近的地方在于它也用了transformer
• image GPT最终所能达到的效果：如果将训练好的模型做微调或者就把它当成一个特征提取器，它在ImageNet上的最高的分类准确率也只能到72，Vit最终的结果已经有88.5了，远高于72
• 但是这个结果也是最近一篇paper叫做MAE爆火的原因。因为在BEiT和MAE这类工作之前，生成式网络在视觉领域很多任务上是没有办法跟判别式网络相比的，判别式网络往往要比生成式网络的结果高很多，但是MAE做到了，它在ImageNet-1k数据集上训练，用一个生成式的模型，比之前判别式的模型效果好很多，而且不光是在分类任务上，最近发现在目标检测上的迁移学习的效果也非常好

Vit其实还跟另外一系列工作是有关系的，用比ImageNet更大的数据集去做预训练，这种使用额外数据的方式，一般有助于达到特别好的效果

• 2017年介绍JFT 300数据集的paper研究了卷积神经网络的效果是怎么随着数据集的增大而提高的
• 一些论文是研究了在更大的数据集（比如说ImageNet-21k和JFT 300M）上做预训练的时候，迁移到ImageNet或者CIFAR-100上的效果如何

这篇论文也是聚焦于ImageNet-21k和JFT 300M，但是训练的并不是一个残差网络，而是训练transformer.

4.ViT模型

4.1 整体结构和前向传播

在模型的设计上尽可能按照最原始的transformer来做的，这样做的好处是transformer在NLP领域已经火了很久了，它有一些非常高效的实现，可以直接拿来使用

简单而言，模型由三个模块组成：

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

前向传播过程：

• Pacth embedding: 一张图片先分割成n个patchs，然后这些patchs变成序列，每个patch输入线性投射层，得到Pacth embedding。比如ViT-L/16表示每个patchs大小是16×16。
• position embedding：self-attention本身没有考虑输入的位置信息，无法对序列建模。而图片切成的patches也是有顺序的，打乱之后就不是原来的图片了。于是和transformer一样，引入position embedding。
• class token：在所有tokens前面加一个新的class token作为这些patchs全局输出，相当于transformer中的CLS（这里的加是concat拼接）。而且它也是有position embedding，位置信息永远是0
• Pacth embedding+position embedding+class token一起输入Transformer Encoder，得到输出。
• 因为所有的token都在跟其它token做交互信息，所以class embedding能够从别的embedding中学到有用的信息，class token的输出当做整个图片的特征，经过MLP Head得到分类结果（VIT只做分类任务）。最后用交叉熵函数进行模型的训练

模型中的Transformer encoder是一个标准的Transformer。整体上来看Vision Transformer的架构还是相当简洁的，它的特殊之处就在于如何把一个图片变成一系列的token

4.2 图片预处理

标准的Transformer模块要求输入的是token（向量）序列，即二维矩阵[num_token, token_dim]。对于图像数据而言，其数据为 [H, W, C]格式的三维矩阵，所以需要先通过一个Embedding层来对数据做变换.

首先将一张图片按给定大小分成一堆Patches。以ViT-B/16为例，将输入图片(224x224)按照16x16大小的Patch尺寸进行划分，划分后会得到196个Patches，每一个图像块的维度就是16* 16 *3=768

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

img_size=224, patch_size=16, in_c=3, embed_dim=768, norm_layer=None
nn.Conv2d(in_c, embed_dim, kernel_size=patch_size, stride=patch_size)

接着通过线性映射E将每个Patch映射到一维向量中。这个全连接层的维度是768*768，第二个768就是文章中的D。

全连接层维度是768*768的原因：每个Patch是一个1616*3的区域(16乘以16乘以3等于768)。

现在得到了patch embedding，它是一个196*768的矩阵，即现在有196个token，每个token向量的维度是768，到目前为止就已经成功地将一个vision的问题变成了一个NLP的问题了，输入就是一系列1d的token，而不再是一张2d的图片了

额外的cls token维度也是768，这样可以方便和后面图像的信息直接进行拼接。所以最后整体进入Transformer的序列的长度是197*768

Position Embedding是可以学习的，每一个向量代表一个位置信息（向量的维度是768），将这些位置信息加到所有的token中，序列还是197*768。

注：对于位置编码信息，本文用的是标准的可以学习的1d position embedding，它也是BERT使用的位置编码。作者也尝试了了别的编码形式，比如说2d-aware（它是一个能处理2d信息的位置编码），但是最后发现结果其实都差不多，没有什么区别

--------------------------图片预处理-example-BEGIN--------------------------
图片预处理的模块：

class Embed(nn.Module):
    def __init__(self, *, image_size, patch_size, dim, pool = 'cls', channels = 3, 
                 dim_head = 64, emb_dropout = 0.):
        super().__init__()
        assert image_size % patch_size == 0, 'Image dimensions must be divisible by the patch size.'  # 保证一定能够完整切块
        num_patches = (image_size // patch_size) ** 2  # 图像patch的个数
        patch_dim = channels * patch_size ** 2  # 线性变换时的输入大小，即下面的p1*p2*c
        assert pool in {'cls', 'mean'}, 'pool type must be either cls (cls token) or mean (mean pooling)'  # 池化方法必须为cls或者mean
 
        self.to_patch_embedding = nn.Sequential(
            Rearrange('b c (h p1) (w p2) -> b (h w) (p1 p2 c)', p1 = patch_size, p2 = patch_size), # 把b张c通道的图像分割成b*（h*w）张大小为p1*p2*c的图像块
            nn.Linear(patch_dim, dim),  # 对分割好的图像块进行线性处理，输入维度为每一个patch的所有像素个数，输出为dim（函数传入的参数）
        )
 
        self.pos_embedding = nn.Parameter(torch.randn(1, num_patches + 1, dim))  # 位置编码，获取一组正态分布的数据用于训练
        self.cls_token = nn.Parameter(torch.randn(1, 1, dim))  # 分类令牌，可训练
        self.dropout = nn.Dropout(emb_dropout)
 
 
    def forward(self, img):
        x = self.to_patch_embedding(img)  # 切块操作，shape (b, n, dim)，b为批量，n为切块数目，dim为线性操作时输入的神经元个数
        b, n, _ = x.shape  # shape (b, n, 1024)
 
        cls_tokens = self.cls_token.repeat([b, 1, 1])  # 将self.cls_token由（1, 1, dim）变为shape (b, 1, dim)
        x = torch.cat((cls_tokens, x), dim=1)  # 将分类令牌拼接到输入中，x的shape (b, n+1, 1024)
        x += self.pos_embedding[:, :(n + 1)]  # 加上位置编码，shape (b, n+1, 1024) 不知道[:, :(n + 1)]是干嘛用的，貌似去掉也行
        x = self.dropout(x)
        return x  

注意，patch embeding使用的不是卷积层，而是用了线性变换，这么做是有优势的。下面分别建立这样的两个网络。输出都是torch.Size([1, 64, 3072])

img = torch.randn([1,3,256,256])
net =  nn.Sequential(
            Rearrange('b c (h p1) (w p2) -> b (h w) (p1 p2 c)', p1 = 32, p2 = 32),
            nn.Linear(3*32*32, 3072),
        )
out_linear = net(img)
 
patch_embed = nn.Sequential(
    nn.Conv2d(3, 32*32*3, kernel_size=32, stride=32) 
)
out_conv = patch_embed(img).flatten(2).transpose(1, 2)

这两个网络的参数个数都是9440256（3072*3073）。卷积层的不方便之处是，embeding的最后一个维度大小实际上是限定死的，如果想要改变维度大小，还要再添加一个线性层。直接使用线性层的好处是，输出的最后一个的设定可以一步到位。

需要额外注意的是，线性层的输入要仔细做reshape，使得每一个patch中的像素在原图像中邻近（卷积做法没有这样的担心）。这里使用了Rearrange，原理是什么，用reshape函数应该怎么做，还没有搞懂

测试：

输入形状为[32,3,256,256]的矩阵，每个patch的大小是32，embeding之后最后一维的大小是1024

img = torch.randn([32,3,256,256])
v = Embed(
    image_size = 256,
    patch_size = 32, 
    dim = 1024,
    emb_dropout = 0.1
)
out = v(img)

out的形状是:

torch.Size([32, 65, 1024])

--------------------------图片预处理-example-END--------------------------

4.3 Transformer Encoder

Transformer Encoder其实就是重复堆叠Encoder Block L次。经过预处理，包括特殊的字符cls和位置编码信息，transformer输入的embedded patches就是一个197*768的tensor。

class Transformer(nn.Module):
    def __init__(self, dim, depth, heads, dim_head, mlp_dim, dropout = 0.):
        super().__init__()
        self.layers = nn.ModuleList([]) # Transformer包含多个编码器的叠加
        for _ in range(depth):
            # 编码器包含两大块：自注意力模块和前向传播模块
            self.layers.append(nn.ModuleList([
                PreNorm(dim, Attention(dim, heads = heads, dim_head = dim_head, dropout = dropout)),
                PreNorm(dim, FeedForward(dim, mlp_dim, dropout = dropout))
            ]))
    def forward(self, x):
        for attn, ff in self.layers:
            x = attn(x) + x  # 自注意力模块和前向传播模块都使用了残差的模式
            x = ff(x) + x
        return x

• Layer Norm层标准化：tensor先过一个layer norm，出来之后还是197*768。

class PreNorm(nn.Module):
    def __init__(self, dim, fn):
        super().__init__()
        self.norm = nn.LayerNorm(dim)  # 正则化
        self.fn = fn  # 具体的操作
    def forward(self, x, **kwargs):
        return self.fn(self.norm(x), **kwargs)

• Multi-Head Attention：假设使用的是ViT的base版本，即使用了12个头，那么k、q、v的维度变成了19764（768/12=64），进行12组k、q、v自注意力操作，最后再将12个头的输出拼接起来，输出还是197768

class Attention(nn.Module):
    def __init__(self, dim, heads = 8, dim_head = 64, dropout = 0.):
        super().__init__()
        inner_dim = dim_head *  heads
        project_out = not (heads == 1 and dim_head == dim) # 多头注意力或输入和输出维度不相同时为True
 
        self.heads = heads
        self.scale = dim_head ** -0.5  # 缩放操作
 
        self.attend = nn.Softmax(dim = -1)
        self.dropout = nn.Dropout(dropout)
 
        self.to_qkv = nn.Linear(dim, inner_dim * 3, bias = False) # 对Q、K、V三组向量线性操作
 
        # 线性全连接，如果多头注意力或输入和输出维度不相同，进行映射，变换维度
        self.to_out = nn.Sequential(
            nn.Linear(inner_dim, dim),
            nn.Dropout(dropout)
        ) if project_out else nn.Identity()
 
    def forward(self, x):
        qkv = self.to_qkv(x).chunk(3, dim = -1)  # 先对Q、K、V进行线性操作，然后chunk成三份
        q, k, v = map(lambda t: rearrange(t, 'b n (h d) -> b h n d', h = self.heads), qkv) # 整理维度，获得Q、K、V
 
        dots = torch.matmul(q, k.transpose(-1, -2)) * self.scale  # 计算相关性
 
        attn = self.attend(dots)
        attn = self.dropout(attn)
 
        out = torch.matmul(attn, v)  # # Softmax运算结果与Value向量相乘，得到最终结果
        out = rearrange(out, 'b h n d -> b n (h d)')  # 重新整理维度
        return self.to_out(out)  # 做线性的全连接操作或者空操作（空操作直接输出out）

torch.chunk(tensor, chunk_num, dim)函数的功能：与torch.cat()刚好相反，它是将tensor按dim（行或列）分割成chunk_num个tensor块，返回的是一个元组。

rearrange操作就是调整维度。比如:

• rearrange(out, ‘b h n d -> b n (h d)’)等于out.transpose(1, 2).reshape(B, N, C)
• q, k, v = map(lambda t: rearrange(t, ‘b n (h d) -> b h n d’, h = self.heads), qkv)可以写成下面

# qkv(): -> [batch_size, num_patches + 1, 3 * total_embed_dim]
# reshape: -> [batch_size, num_patches + 1, 3, num_heads, embed_dim_per_head]
# permute: -> [3, batch_size, num_heads, num_patches + 1, embed_dim_per_head]
qkv = self.qkv(x).reshape(B, N, 3, self.num_heads, total_embed_dim // self.num_heads).permute(2, 0, 3, 1, 4)
# [batch_size, num_heads, num_patches + 1, embed_dim_per_head]
q, k, v = qkv[0], qkv[1], qkv[2]

attention操作的整体流程：

• 首先对输入生成query, key和value，这里的“输入”有可能是整个网络的输入，也可能是某个hidden layer的output。在这里，生成的qkv是个长度为3的元组，每个元组的大小为(1, 65, 1024)
• 对qkv进行处理，重新指定维度，得到的q, k, v维度均为(1, 16, 65, 64)
• q和k做点乘，得到的dots维度为(1, 16, 65, 65)
• 对dots的最后一维做softmax，得到各个patch对其他patch的注意力得分
• 将attention和value做点乘
• 对各个维度重新排列，得到与输入相同维度的输出 (1, 65, 1024)
• 根据需要，做投射
  • Dropout/DropPath：在原论文的代码中是直接使用的Dropout层，在但rwightman实现的代码中使用的是DropPath（stochastic depth），可能后者会更好一点。
  • 再过一层layer norm，还是197*768
  • MLP Block，全连接+GELU激活函数+Dropout组成。把维度放大到4倍[197, 768] -> [197, 3072]，再还原回原节点个数[197, 3072] -> [197, 768]。

class FeedForward(nn.Module):
    def __init__(self, dim, hidden_dim, dropout = 0.):
        super().__init__()
        # 前向传播
        self.net = nn.Sequential(
            nn.Linear(dim, hidden_dim),
            nn.GELU(),
            nn.Dropout(dropout),
            nn.Linear(hidden_dim, dim),
            nn.Dropout(dropout)
        )
    def forward(self, x):
        return self.net(x)

进去Transformer block之前是197768，出来还是197768，这个序列的长度和每个token对应的维度大小都是一样的，所以就可以在一个Transformer block上不停地往上叠加Transformer block，最后有L层Transformer block的模型就构成了Transformer encoder

Transformer从头到尾都是使用D当作向量的长度的，都是768，同一个模型里这个维度是不变的。如果transformer变得更大了，D也可以相应的变得更大。

The Transformer uses constant latent vector size D through all of its layers, so we flatten the patches and map to D dimensions with a trainable linear projection. We refer to the output of this projection as the patch embeddings.

4.4 MLP Head和ViT-B/16模型结构图

对于分类，只需要提取出[class]token生成的对应结果就行，即[197, 768]中抽取出[class]token对应的[1, 768]，通过MLP Head得到最终的分类结果。MLP Head原论文中说在训练ImageNet21K时是由Linear+tanh激活函数+Linear组成，但是迁移到ImageNet1K上或者你自己的数据上时，只定义一个Linear即可。注意，在Transformer Encoder后其实还有一个Layer Norm，

class ViT(nn.Module):
    def __init__(self, image_size, patch_size, num_classes, 
                 dim, depth, heads, mlp_dim, pool = 'cls', 
                 channels = 3, dim_head = 64, dropout = 0., 
                 emb_dropout = 0.):
        super().__init__()
        self.embed = Embed(image_size, patch_size, dim, 
                 pool = 'cls', channels = 3, 
                 dim_head = 64, emb_dropout = 0.)
 
        self.transformer = Transformer(dim, depth, heads, dim_head, mlp_dim, dropout)  # Transformer模块
 
        self.pool = pool
        self.to_latent = nn.Identity()  # 占位操作
 
        self.mlp_head = nn.Sequential(
            nn.LayerNorm(dim),  # 正则化
            nn.Linear(dim, num_classes)  # 线性输出
        )
 
    def forward(self, img):
        x = self.embed(img)  # embeding操作
        x = self.transformer(x)  # transformer操作
 
        x = x.mean(dim = 1) if self.pool == 'mean' else x[:, 0]
 
        x = self.to_latent(x)
        return self.mlp_head(x)  # 线性输出

4.5 数学公式描述

（1）Xp​表示图像块的patch，一共有N个patch，E表示线性投影的全连接层，得到一些patch embedding。在它前面拼接一个class embedding（Xclass）。得到所有的tokens后，将位置编码信息Epos也加进去。

• Z0就是整个transformer的输入。

（2）-（3）循环

对于每个transformer block来说，里面都有两个操作：一个是多头自注意力，一个是MLP。在做这两个操作之前，都要先经过layer norm，每一层出来的结果都要再去用一个残差连接

• ZL’就是每一个多头自注意力出来的结果
• ZL就是每一个transformer block整体做完之后出来的结果

4）L层循环结束之后将ZL（最后一层的输出）的第一个位置上的ZL0，也就是class token所对应的输出当作整体图像的特征，去做最后的分类任务

完整代码

import torch
from torch import nn, einsum
import torch.nn.functional as F
 
from einops import rearrange, repeat
from einops.layers.torch import Rearrange
 
class Embed(nn.Module):
    def __init__(self, image_size, patch_size, dim, 
                 pool = 'cls', channels = 3, 
                 dim_head = 64, emb_dropout = 0.):
        super().__init__()
        assert image_size % patch_size == 0, 'Image dimensions must be divisible by the patch size.'  # 保证一定能够完整切块
        num_patches = (image_size // patch_size) ** 2  # 获取图像切块的个数
        patch_dim = channels * patch_size ** 2  # 线性变换时的输入大小，即每一个图像宽、高、通道的乘积
        assert pool in {'cls', 'mean'}, 'pool type must be either cls (cls token) or mean (mean pooling)'  # 池化方法必须为cls或者mean
 
        self.to_patch_embedding = nn.Sequential(
            Rearrange('b c (h p1) (w p2) -> b (h w) (p1 p2 c)', p1 = patch_size, p2 = patch_size),  # 把b张c通道的图像分割成b*（h*w）张大小为P1*p2*c的图像块
            nn.Linear(patch_dim, dim),  # 对分割好的图像块进行线性处理（全连接），输入维度为每一个小块的所有像素个数，输出为dim（函数传入的参数）
        )
 
        self.pos_embedding = nn.Parameter(torch.randn(1, num_patches + 1, dim))  # 位置编码，获取一组正态分布的数据用于训练
        self.cls_token = nn.Parameter(torch.randn(1, 1, dim))  # 分类令牌，可训练
        self.dropout = nn.Dropout(emb_dropout)
 
 
    def forward(self, img):
        x = self.to_patch_embedding(img)  # 切块操作，shape (b, n, dim)，b为批量，n为切块数目，dim为最终线性操作时输入的神经元个数
        b, n, _ = x.shape  # shape (b, n, 1024)
 
        cls_tokens = self.cls_token.repeat([b, 1, 1])  
        # 分类令牌，将self.cls_token（形状为1, 1, dim）赋值为shape (b, 1, dim)
        x = torch.cat((cls_tokens, x), dim=1)  # 将分类令牌拼接到输入中，x的shape (b, n+1, 1024)
        x += self.pos_embedding[:, :(n + 1)]  # 进行位置编码，shape (b, n+1, 1024)
        x = self.dropout(x)
        return x  
 
class PreNorm(nn.Module):
    def __init__(self, dim, fn):
        super().__init__()
        self.norm = nn.LayerNorm(dim)  # 正则化
        self.fn = fn  # 具体的操作
    def forward(self, x, **kwargs):
        return self.fn(self.norm(x), **kwargs)
 
class FeedForward(nn.Module):
    def __init__(self, dim, hidden_dim, dropout = 0.):
        super().__init__()
        # 前向传播
        self.net = nn.Sequential(
            nn.Linear(dim, hidden_dim),
            nn.GELU(),
            nn.Dropout(dropout),
            nn.Linear(hidden_dim, dim),
            nn.Dropout(dropout)
        )
    def forward(self, x):
        return self.net(x)
 
class Attention(nn.Module):
    def __init__(self, dim, heads = 8, dim_head = 64, dropout = 0.):
        super().__init__()
        inner_dim = dim_head *  heads
        project_out = not (heads == 1 and dim_head == dim) # 多头注意力或输入和输出维度不相同时为True
 
        self.heads = heads
        self.scale = dim_head ** -0.5  # 缩放操作
 
        self.attend = nn.Softmax(dim = -1)
        self.dropout = nn.Dropout(dropout)
 
        self.to_qkv = nn.Linear(dim, inner_dim * 3, bias = False) # 对Q、K、V三组向量线性操作
 
        # 线性全连接，如果多头注意力或输入和输出维度不相同，进行映射，变换维度
        self.to_out = nn.Sequential(
            nn.Linear(inner_dim, dim),
            nn.Dropout(dropout)
        ) if project_out else nn.Identity()
 
    def forward(self, x):
        qkv = self.to_qkv(x).chunk(3, dim = -1)  # 先对Q、K、V进行线性操作，然后chunk成三份
        q, k, v = map(lambda t: rearrange(t, 'b n (h d) -> b h n d', h = self.heads), qkv) # 整理维度，获得Q、K、V
 
        dots = torch.matmul(q, k.transpose(-1, -2)) * self.scale  # 计算相关性
 
        attn = self.attend(dots)
        attn = self.dropout(attn)
 
        out = torch.matmul(attn, v)  # # Softmax运算结果与Value向量相乘，得到最终结果
        out = rearrange(out, 'b h n d -> b n (h d)')  # 重新整理维度
        return self.to_out(out)  # 做线性的全连接操作或者空操作（空操作直接输出out）
 
class Transformer(nn.Module):
    def __init__(self, dim, depth, heads, dim_head, mlp_dim, dropout = 0.):
        super().__init__()
        self.layers = nn.ModuleList([])  # Transformer包含多个编码器的叠加
        for _ in range(depth):
            # 编码器包含两大块：自注意力模块和前向传播模块
            self.layers.append(nn.ModuleList([
                PreNorm(dim, Attention(dim, heads = heads, dim_head = dim_head, dropout = dropout)),  # 多头自注意力模块
                PreNorm(dim, FeedForward(dim, mlp_dim, dropout = dropout))  # 前向传播模块
            ]))
    def forward(self, x):
        for attn, ff in self.layers:
            # 自注意力模块和前向传播模块都使用了残差的模式
            x = attn(x) + x
            x = ff(x) + x
        return x
 
class ViT(nn.Module):
    def __init__(self, image_size, patch_size, num_classes, 
                 dim, depth, heads, mlp_dim, pool = 'cls', 
                 channels = 3, dim_head = 64, dropout = 0., 
                 emb_dropout = 0.):
        super().__init__()
        self.embed = Embed(image_size, patch_size, dim, 
                 pool = 'cls', channels = 3, 
                 dim_head = 64, emb_dropout = 0.)
 
        self.transformer = Transformer(dim, depth, heads, dim_head, mlp_dim, dropout)  # Transformer模块
 
        self.pool = pool
        self.to_latent = nn.Identity()  # 占位操作
 
        self.mlp_head = nn.Sequential(
            nn.LayerNorm(dim),  # 正则化
            nn.Linear(dim, num_classes)  # 线性输出
        )
 
    def forward(self, img):
        x = self.embed(img)  
        x = self.transformer(x)  # transformer操作
 
        x = x.mean(dim = 1) if self.pool == 'mean' else x[:, 0]
 
        x = self.to_latent(x)
        return self.mlp_head(x)  # 线性输出

• image_size：int 类型参数，图片大小。 如果矩形图像，为宽度和高度的最大值
• patch_size：int 类型参数，patch大小。image_size 必须能够被 patch_size整除。n must be greater than 16
• num_classes：int 类型参数，分类数目。
• dim：int 类型参数，embedding的维度。
• depth：int 类型参数，Transformer模块的个数。
• heads：int 类型参数，多头注意力中“头”的个数。
• mlp_dim：int 类型参数，多层感知机中隐藏层的神经元个数。
• channels：int 类型参数，输入图像的通道数，默认为3。
• dropout：float类型参数，Dropout几率，取值范围为[0, 1]，默认为 0.。
• emb_dropout：float类型参数，进行Embedding操作时Dropout几率，取值范围为[0, 1]，默认为0。
• pool：string类型参数，取值为 cls或者 mean 。