本文在transformer的基础上对ViT进行讲解,transformer相关部分可以看我另一篇博客(transformer中对于QKV的个人理解-CSDN博客)。
一、网络结构概览
上图展示了Vision Transformer (ViT) 的基本架构,我按照运行顺序分为三个板块进行讲解。以下是 ViT 处理图像的基本步骤:
-
Embedding(Linear Projection of Flattened Patches):首先,输入图像被分割成固定大小的patch。每个patch被展平并经过一个线性层,将其转换为特征向量。此外,每个特征向量都会附加一个表示其在原始图像中相对位置的位置嵌入。
-
Transformer 编码器(Transformer Encoder):这些预处理过的特征向量作为输入馈送到一系列的 Transformer 层进行编码。每个 Transformer 层包含多头注意力机制(Multi-Head Attention)和前馈神经网络(Feed-Forward Network)。在每个 Transformer 层之间有层归一化(Layer Normalization),以保持特征分布的一致性。
-
MLP Head 和分类(MLP Head):最后,来自 Transformer 编码器的最后一层输出通过一个多层感知机(MLP)头部,通常是一个简单的全连接层,用于执行最终的任务,如图像分类。在这个阶段,模型将学习到的特征映射到类概率空间。
上图是另外一个大佬做的细致一点的网络结构图。
1. Embedding
我将Embedding分为4个步骤:
- 分割patch
首先输入为一张图片,vit中默认将图像缩放为224*224,通常利用16*16的patch进行划分得到(224*224) / (16*16)=14*14=196张子图,每张patch的大小为16*16,3通道。 - 转换获得token
将每个二维的补丁转换为可以输入到 Transformer 的一维序列,将14*14个patch展平成一个长向量,再通过线性变换得到196个token,token的长度为768。
然而,在vit中是利用768个卷积核大小为16*16,stride为16,padding为0的一个卷积层直接对224*224的输入图像进行卷积,从而得到14*14*768的输出,再对其展平得到[196,768]的token,相当于合并了第一步和第二步。 - 拼接上类别token
之后在得到的196个 token 的前面加上加上一个新的Class Token(即图中0号紫色框右边带*的那个框,这不是通过某个patch产生的。其作用类似于BERT中的Class Token。),得到[197,768]的数据。
在ViT中,Class Token(通常记为CLS)是一个可学习的参数。它在网络初始化时被随机初始化,类似于其他神经网络权重参数。在经过多层Transformer编码器后,Class Token会聚集来自所有图像块的信息,形成图像的全局表示。最终的Class Token表示(即最后一层Transformer编码器输出的Class Token)被输入到一个分类头(通常是一个全连接层)中,用于图像的分类任务。分类头的输出即为预测的类别概率分布。 - 加上位置编码
patch得到的图像是没有位置信息的,需要用position embedding将位置信息加到模型中去。如上图所示,编号有0-9的紫色框表示各个位置的position embedding
位置编码是可训练的编码,通过叠加加入到[197,768]的数据中,最终输出[197,768]的数据。
加了位置编码,性能有明显提升,但是不同编码器的方式对性能提升差不多。所以源码中使用的是1-D位置编码。
上图呈现了一个热力图,其中水平轴代表输入补丁的列数,垂直轴代表输入补丁的行数,颜色深浅表示相似程度。
通过观察热力图,我们可以看到随着补丁距离变远,它们的位置嵌入变得越来越不相似。在真实世界中,离得近的物体往往比离得远的物体具有更强的空间相关性。这种特性有助于 ViT 更好地理解图像内容及其结构。
2. Transformer 编码器
Transformer encoder层如图所示,将Transformer encoder重复堆叠 L 次,整个模型也就包括 L 个 Transformer。关于Layer Norm和多头注意力模块的具体解析可以看我主页其他的博客,这里不过多赘述。
MLP中先通过一个线性层将输入数据的通道数变为原来的4倍,之后通过GELU激活函数和Dropout,再通过一个线性层将4倍通道数变为原来的通道数。
3. MLP Head 和分类
MLP Head 层位于 Transformer 编码器之后,用于完成特定任务,如图像分类。该层通常是一个多层感知机(Multilayer Perceptron,简称 MLP),它接收来自 Transformer 编码器的输出,并对其进行进一步处理以生成最终的预测。在不同的场景下,MLP Head 的结构可能会有所不同。
在这张图片中,我们能看到两种情况下的 MLP Head 设计:
-
训练 ImageNet21K 时的 MLP Head:
当训练 ViT 时,特别是在大型数据集如 ImageNet21K 上,MLP Head 包含三个组成部分:线性层、tanh 激活函数以及另一个线性层。这样的设计允许 MLP Head 对编码器输出进行非线性的转换,以适应复杂的模式和特征。激活函数 tanh 提供了一种非线性变形,可以帮助模型学习更复杂的表示。 -
迁移至 ImageNet1K 或其他数据集时的 MLP Head:
当将 ViT 迁移到较小的数据集,如 ImageNet1K 或者自己的数据集时,通常只保留一个线性层。这是因为较小的数据集可能不需要那么多的复杂性,一个线性层就足以提供足够的泛化能力。减少层数也可以降低过拟合的风险。