背景
Transformer模型最初是使用在NLP中,但近几年Transformer模型在图像上的使用越来越频繁,最新的模型也出现了很多基于Transfomer的,而其中经典的是Vision Transformer(ViT),它是用于图像分类的,这里就以ViT-B/16这个模型来学习Transformer模型是如何在图像领域使用的。
整体框架
ViT的网络结构如图所示,可以将ViT分为三个部分:
- Embedding层:对原始的图像数据进行处理,转换为token(向量)序列,使其符合Transformer Encoder的输入要求
- Transformer Encoders层:包含多个Encoder Block,对输入的token进行处理
- MLP Head层:用于分类的层结构
各个部分的详解
Embedding层
对于标准的Transformer模块,要输入的是token(向量)序列,也就是二维矩阵[num_token,token_dim(向量的长度)],以ViT-B/16为例,token的长度为768。
而对于图像数据而言,其格式为[H,W,C],显然不是Transformer模块需要的,因此需要通过一个Embedding层来对数据格式进行转换。
操作过程如下图所示,以ViT-B/16为例,首先将输入图像(224x224)
切分为196个16x16
的Patch,这些Patch的格式为[16,16,3]
。接下来将这些Patch进行线性映射,也就是转换为一维的向量:[16,16,3]-->[768]
。
在代码实现中,直接使用一个卷积层来实现。以ViT-B/16为例,使用一个卷积核为16x16,步距为16,卷积核个数为768的卷积层来实现。通过卷积层会得到[224,224,3]-->[14,14,768]
,然后在通过Flatten操作将[H,W]展平即可[14, 14, 768] -> [196, 768]
,此时正好为二维矩阵,符合Transformer模块的需要。
在输入到Transformer模块前,需要加上[class]token(token所属的类别)与位置编码Position Embedding。
- [class]token:在刚刚得到的一堆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采用的是一个可训练的参数,这里是直接叠加(add)到token上。以ViT-B/16为例,刚刚拼接[class]token后shape是
[197, 768]
,那么这里的Position Embedding的shape也是[197, 768]
。
Transformer Encoder层
Transformer Encoder其实就是重复堆叠Encoder Block L次,Encoder Block由以下部分构成:
- Layer Norm:对每个token进行标准化处理
- Multi-Head Attention:链接
- Dropout/DropPath:防止过拟合,在原论文的代码中是直接使用的Dropout层,但也有人使用DropPath
- MLP Block:如图右侧所示,由全连接层(Liner)+GELU激活函数+DropOut层构成,需要注意的是第一个全连接层会将节点个数翻4倍
[197,768]-->[197,3072]
,而第二个全连接层会将节点个数恢复原样[197,3072]-->[197,768]
,这样经过MPL Block不会修改数据的结构。
通过Transformer Encoder后输出的shape和输入的shape是保持不变的,以ViT-B/16为例,输入的是[197, 768]输出的还是[197, 768]
而在代码实现中,会在Transformer模块前面加一个Dropout层,后面加一个Layer Norm层。
MLP Head
由于只需要分类的信息,因此只需要将[class]token生成的对应结果提取出来即可,即[197, 768]
中抽取出[class]token对应的[1, 768]
,最后在使用MLP Head得到最终的分类结果。
在论文中,训练ImageNet21K(大型数据集)时,MLP Head是由Linear+tanh激活函数+Linear组成。但是迁移到ImageNet1K上或者你自己的数据上时,只用一个Linear即可。
以ViT-B/16为例
不同种类的ViT
论文中提供了三种ViT模型的参数
- Layers:Transformer模块中Encoder模块的个数
- Hidden Size:经过Embedding层后每个token的长度
- MLP Size:MLP Block中第一个全连接层的输出节点个数,是Hidden Size的四倍。
- Heads:Multi-Head Attention的heads数。
CNN与Transformer的混合网络:Hybrid
Hybrid将传统CNN的特征提取和Transformer进行结合。
下图是以ResNet50作为特征提取器的Hybrid的网络结构,与ViT区别在于开始阶段会使用ResNet50网络进行特征提取。但这里的ResNet50与普通的ResNet50不同。
首先这里的R50的卷积层采用的StdConv2d不是传统的Conv2d,然后将所有的BatchNorm层替换成GroupNorm层。
而且在原Resnet50网络中,stage1重复堆叠3次,stage2重复堆叠4次,stage3重复堆叠6次,stage4重复堆叠3次。但是在这里的R50中,stage4的3次被集成到stage3中,因此stage3堆叠了9次。
通过R50 Backbone进行特征提取后,得到的特征矩阵为[14,14,1024]
,然后输入到Embedding层,要注意的是这个Embedding层中的卷积变为卷积核为1x1,步长为1,只是调整了channel为768。
后面的处理与ViT一样。