极智AI | 详解 ViT 算法实现
MobileViT模型简介
ECCV 2022丨力压苹果MobileViT,这个轻量级视觉模型新架构火了
ECCV 2022丨轻量级模型架构火了,力压苹果MobileViT(附代码和论文下载)
再读VIT,还有多少细节是你不知道的
CNN 里最重要的算子是卷积,卷积具有两个很重要的特性:translation equivariance 平移等价性 和 locality 局部性。
归纳偏置用大白话来说,就是一种假设,或者说一种先验知识。有了这种先验,我们就能知道哪一种方法更适合解决哪一类任务。所以归纳偏置是一种统称,不同的任务其归纳偏置下包含的具体内容不一样。
对图像任务来说,它的归纳偏置有以下两点:
在这两种先验假设下,CNN成为了图像任务最佳的方案之一。卷积核能最大程度保持空间局部性(保存相关物体的位置信息)和平移等边性,使得在训练过程中,最大限度学习和保留原始图片信息。
如果说ViT相比于卷积,在图像任务上没有显著优势,那大概率ViT对这两种先验的维护没有CNN做的好,具体来看:
图中箭头所指的两部分都属于同一栋建筑。在卷积中,我们可以用大小适当的卷积核将它们圈在一起。但是在ViT中,它们之间的位置却拉远了,如果我把patch再切分细一些,它们的距离就更远了。虽然attention可以学习到向量间的想关系,但是ViT在空间局部性的维护上,确实没有卷积做的好。而在平移等边性上,由于ViT需要对patch的位置进行学习,所以对于一个patch,当它位置变幻时,它的输出结果也是不一样的。所以,ViT的架构没有很好维护图像问题中的归纳偏置假设。
但是,这就意味着ViT没有翻盘的一天了吗?当然不是,不要忘了,**Transformer架构的模型都有一个广为人知的特性:大力出奇迹。**只要它见过的数据够多,它就能更好地学习像素块之间的关联性,当然也能抹去归纳偏置的问题。
预训练好的ViT模型是个有力的特征提取器,我们可以用它输出的特征,去做更多有趣的下游任务(downstream task)。例如拿它去做类型更丰富的分类,目标检测等事情。在做这些任务时,我们会喂给预训练模型一堆新的数据,同时尽量保证模型的主体架构不变(例如ViT整体参数不动,只在输出层后接一个新模型,再次训练时只对新模型做参数更新之类)。这种既利用了已有模型的特征提取能力,又能让模型更好适应不同任务的操作,称为微调(fine-tune)。
在fine-tune的时候,我们用的图像大小可能和预训练时的并不一致,比如:
224*224*3
大小的图片,fine-tune时为了效果更好,一般选择分辨率更高的图片,例如1024*1024*3
。那么多出来的patch,在fine-tune时要怎么给它们位置编码呢?如果统一都赋成0向量,然后在fine-tune的时候再去训练这些向量,看起来可以,但这样粗暴的赋值不仅增加了计算量,也浪费了已有的信息(例如,是否能从已有的位置编码粗略地初始化一些新的位置编码出来?)考虑到这一点,ViT在fine-tune时,对预训练阶段的位置编码做了2D插值处理。
如图绿色部分所示,在fine-tune阶段要处理的patch/token数 s f i n e t u n e s_{\mathrm{finetune}} sfinetune 可能比预训练阶段要处理的 s p r e t r a i n s_{\mathrm{pretrain}} spretrain 要多。图中红色部分演示了如何通过插值方法将 s p r e t r a i n s_{\mathrm{pretrain}} spretrain 扩展至 s f i n e t u n e s_{\mathrm{finetune}} sfinetune。其中interpolate部分就是2D插值,这部分是重点,我们直接看下代码中的操作:
new_pos_embedding_img = nn.functional.interpolate(
pos_embedding_img,
size=new_seq_length_1d,
mode=interpolation_mode,
align_corners=True,
)
可以发现这里用了pytorch内置的interpolate函数,mode表示具体的插值方法,在ViT中采用的是bicubic。align_corners=True
的意思是在固定原矩阵四角的情况下按mode进行插值,可以参考上图中,白色圆圈表示原始的矩阵,蓝色点表示做完插值后的矩阵。插值后矩阵的四角保持不变,中间则按设置的方法做插值。关于插值位置编码更详细的讲解,可以参考:【ViT 微调时关于position embedding如何插值(interpolate)的详解】
NLP Transformer | ViT |
---|---|
句子 | 图像 |
words | patchs |
word embedding | patch embedding |
ViT 全称 Vision Transformer,不同于传统的基于CNN的网络结果,是基于transformer结构的cv网络。
ViT模型主要应用于图像分类领域。因此,其模型结构相较于传统的Transformer有以下几个特点:
ViT模型利用Transformer模型在处理上下文语义信息的优势,将图像转换为一种“变种词向量”然后进行处理,而这种转换的意义在于,多个Patch之间本身具有空间联系,这类似于一种“空间语义”,从而获得了比较好的处理效果。
解决图像分类任务的利器——Vision Transformer
mindspore vit模型
如下图所示, BiT 代表 ResNet,ViT* 代表 ViT 系列,可以看出:在相对小一些的数据集上如 ImageNet,ViT 普遍比不过 ResNet,而在 ImageNet-21k 这种中型的数据集上 ViT 性能和 ResNet 旗鼓相当,慢慢开始超越了,当在 JFT-300M 这种大型一些的数据集上时,ViT 开始全面超越 ResNet。
ViT预训练了三种不同参数规模的模型,分别是ViT-Base
,ViT-Large
和ViT-Huge
。其规模可具体见上图。
在论文及实际使用中,我们常用ViT-size/patch_size
的形式来表示该模型是在“什么规模”及“多大的patch尺寸”上预训练出来的。例如ViT-H/14
就表示该模型是在Huge规模上,用patch尺寸为14的数据做预训练的。
总结起来,ViT的训练其实就在做一件事情:把图片打成Patch,送入Transformer Encoder,然后拿对应位置的向量,过一个简单的softmax多分类模型,去预测原始图片中描绘的物体类别即可。
你可能会想:“这个分类任务只用一个简单的softmax,真得能分准吗?”其实,这就是ViT的精华所在:**ViT的目的不是让这个softmax分类模型强大,而是让这个分类模型的输入强大。这个输入就是Transformer Encoder提炼出来的特征。**分类模型越简单,对特征的要求就越高。
**所以,为什么说Transformer开启了大一统模型的预训练大门呢?主要原因就在于它对特征的提炼能力——这样我们就可以拿这个特征去做更多有趣的任务了。**这也是ViT能成为后续多模态backbone的主要原因。
ViT是基于多个transformer encoder模块串联起来,由多个inception模块串联起来,基本结构由patch_embeding + n transformer layer + head(分类网络中就是FC)构成。
ViT模型输入:Patch Embedding
+Class Embedding
+Position Embedding
。
功能:通过Patch Embedding操作,得到一维向量。
在 NLP Transformer中,句子都是一维的,而图像数据是二维的,那怎么把二维的图像数据变成跟 NLP 一样一维的呢,有几种方法:
ViT 利用等分窗口图片块的思想,将图像分成块,每个小块称作Patch,每个Patch块看作NLP Transformer中的一个单词。
例如,假设原始图片尺寸大小为:224*224*3(H*W*C)
。
每个Patch的尺寸设为16(P=16),则每个Patch下图片的大小为:16*16*3
,Patch共有 (224/16) x (224/16) = 14 x 14=196个。
Patch Embedding将每一个Patch的矩阵拉伸成为一个1维向量,从而获得近似词向量堆叠的效果。
如上图所示,每个Patch对应着一个token,将每个Patch展平,则得到输入矩阵X,其大小为(196, 768)
,其中16*16*3=768
,也就是每个token是768维。通过这样的方式,我们成功将图像数据处理成自然语言的向量表达方式。
那么现在问题来了,对于图中每一个16*16*3
的小方块,我要怎么把它拉平成1*768
维度的向量呢?
比如说,我先把第一个channel拉成一个向量,然后再往后依次接上第二个channel、第三个channel拉平的向量。但这种办法下,同一个pixel本来是三个channel的值共同表达的,现在变成竖直的向量之后,这三个值的距离反而远了。基于这个原因,你可能会想一些别的拉平方式,但归根究底它们都有一个共同的问题:太规则化,太主观。
采用768个16*16*3
尺寸的卷积核,stride=16,padding=0。这样我们就能得到14*14*768
大小的特征图。如图所示,特征图中每一个1*1*768
大小的子特征图,都是由卷积核对第一块patch做处理而来,因此它就能表示第一块patch的token向量。
Patch Embedding之后,会经过 Class Embedding 和 Position Embedding 两个过程。
功能:通过Class Embedding操作,得到类别向量。
Class Embedding
简介Class Embedding
主要借鉴了BERT模型的用于文本分类时的思想,在每一个word vector之前增加一个类别值,通常是加在向量的第一位。例如,Patch Embedding
得到的196维的向量加上 Class Embedding
后,变成197维。Class Embedding
用于最后的类别输出,可参考BERT 的 class token
,整个过程示意如下图:
Class Embedding
是可以学习的参数,经过网络的不断训练,最终以输出向量的第一个维度的输出来决定最后的输出类别。由于输入是 16x16个Patch,所以输出进行分类时是取16x16个Class Embedding
进行分类。
Class Embedding
预测类别的方式Class Embedding
有两种预测类别的方式:
这两种方式都是可行的,更倾向于使用 class token
是因为想把原滋原味的 transformer 直接应用到 CV 领域。这两种预测类别的方式,试验效果如下:
其中,**蓝色是 class token **,橙色和绿色是全局平均池化,橙色的存在告诉你需要好好调参,结果的好坏和你调参的姿态关系很大。
功能:通过Position Embedding操作,得到位置向量。
为什么要加这个位置编码,加上以后会有什么效果?
图像切分重排后,失去了位置信息,并且Transformer的内部运算是空间信息无关的,所以需要把位置信息编码重新传进网络。
Position Embedding
将位置编码嵌入图像块,用于表达图像块在原图的位置信息。位置编码随位置变化,即位置差别越大,位置编码差别越大。
在NLP Transformer中,把一个word单词转换成vector向量,就是把一个单词映射到了一个高维空间的位置,意思相近的词会在高维空间内比较靠近。
Position Embedding
操作会创建一个197维的可训练的向量,加入到经过 Class Embedding
的向量中。Position Embedding的长度和 Class Embedding
一致,两个embedding直接相加。
产生位置信息的方式主要分两大类,一类是直接通过固定算法产生,一种是训练获得。但加位置信息的方式还是比较统一且粗暴的。具体从方法上可以位置编码分为几种:
虽然位置编码的方法挺多,但从实验来看,对网络最后的结果影响不大(No Pos 会相对低一点),数据如下:
Transformer Encoder是两个块的堆叠,然后再整体叠加 L 次。这两个块指的是:
在 CV 里用的比较多的是 BatchNorm,那在NLP里为啥不喜欢用 BN 呢?因为 NLP 里输入序列往往是动态的,即序列的长度不定,一个序列就是一个样本。而BN 计算的是样本间的归一化,这样做一定会导致值域波动很大;而LN是在样本内做,不用考虑类间差异,波动就相对小很多。简单理解,Batch Normalization是对每个通道的**所有样本(样本间)进行归一化,而Layer normalization是对每个样本(样本内)**的所有特征进行归一化。
# NLP Example
batch, sentence_length, embedding_dim = 20, 5, 10
embedding = torch.randn(batch, sentence_length, embedding_dim)
# 指定归一化的维度
layer_norm = nn.LayerNorm(embedding_dim)
# 进行归一化
layer_norm(embedding)
# Image Example
N, C, H, W = 20, 5, 10, 10
input = torch.randn(N, C, H, W)
# Normalize over the last three dimensions (i.e. the channel and spatial dimensions)
# as shown in the image below
layer_norm = nn.LayerNorm([C, H, W])
output = layer_norm(input)
在ViT中,虽然LN处理的是图片数据,但在进行LN之前,图片已经被切割成了Patch,而每个Patch表示的是一个词,因此是在用语义的逻辑在解决视觉问题,因此在ViT中,LN也是按语义的逻辑在用的。
Multi-Head Attention 多头注意力机制,来源于论文《Attention Is All You Need》,示意如下:
多头即将模型分为多个头,形成多个子空间,让模型去关注不同方面的信息,将 Scaled Dot-Product Attention 过程做 h 次,再把输出做 cat。这样做的目的是为了使网络能够综合利用多方面角度提取更加准确的表示,从而可以捕捉到更加丰富的特征,可以类比 CNN 中多个核分别提取特征的作用。
Patch embedding -> 加cls -> 加pos embedding -> 用blocks进行encoding -> layer normalization -> 输出图的embedding。
def forward_features(self, x):
# x由(B,C,H,W)->(B,N,E)
x = self.patch_embed(x)
# stole cls_tokens impl from Phil Wang, thanks
# cls_token由(1, 1, 768)->(B, 1, 768), B是batch_size
cls_token = self.cls_token.expand(x.shape[0], -1, -1)
# dist_token是None,DeiT models才会用到dist_token。
if self.dist_token is None:
# x由(B, N, E)->(B, 1+N, E)
x = torch.cat((cls_token, x), dim=1)
else:
# x由(B, N, E)->(B, 2+N, E)
x = torch.cat((cls_token, self.dist_token.expand(x.shape[0], -1, -1), x), dim=1)
# +pos_embed:(1, 1+N, E),再加一个dropout层
x = self.pos_drop(x + self.pos_embed)
x = self.blocks(x)
# nn.LayerNorm
x = self.norm(x)
if self.dist_token is None:
# 不是DeiT,输出就是x[:,0],(B, 1, 768),即cls_token
return self.pre_logits(x[:, 0])
else:
# 是DeiT,输出就是cls_token和dist_token
return x[:, 0], x[:, 1]
这里在Patch 那个维度加入了一个cls_token,可以这样理解这个存在,其他的embedding表达的都是不同的Patch的特征,而cls_token是要综合所有Patch的信息,产生一个新的embedding,来表达整个图的信息。而dist_token则是属于DeiT网络的结构。
forward features -> 最终输出。
def forward(self, x):
#(B,C,H,W)-> (B, 1, 768)
# (B,C,H,W) -> (B, 1, 768), (B, 1, 768)
x = self.forward_features(x)
if self.head_dist is not None:
# 如果num_classes>0, (B, 1, 768)->(B, 1, num_classes)
# 否则不变
x, x_dist = self.head(x[0]), self.head_dist(x[1])
if self.training and not torch.jit.is_scripting():
return x, x_dist
else:
# during inference,
# return the average of both classifier predictions
return (x + x_dist) / 2
else:
# 如果num_classes>0, (B, 1, 768)->(B, 1, num_classes)
# 否则不变
x = self.head(x)
return x
MLP 全称 multi-layer perceptron,里面使用非线性激活函数去做分类的预测。
Encoder在ViT中的实现细节如下面代码所示(layer normalization -> multi-head attention -> drop path -> layer normalization -> mlp -> drop path),换了个名字,叫block了:
class Block(nn.Module):
def __init__(self, dim, num_heads, mlp_ratio=4., qkv_bias=False, drop=0., attn_drop=0.,
drop_path=0., act_layer=nn.GELU, norm_layer=nn.LayerNorm):
super().__init__()
# 将每个样本的每个通道的特征向量做归一化
# 也就是说每个特征向量是独立做归一化的
# 我们这里虽然是图片数据,但图片被切割成了Patch,用的是语义的逻辑
self.norm1 = norm_layer(dim)
self.attn = Attention(dim, num_heads=num_heads, qkv_bias=qkv_bias, attn_drop=attn_drop, proj_drop=drop)
# NOTE: drop path for stochastic depth, we shall see if this is better than dropout here
self.drop_path = DropPath(drop_path) if drop_path > 0. else nn.Identity()
self.norm2 = norm_layer(dim)
mlp_hidden_dim = int(dim * mlp_ratio)
# 全连接,激励,drop,全连接,drop,若out_features没填,那么输出维度不变。
self.mlp = Mlp(in_features=dim, hidden_features=mlp_hidden_dim, act_layer=act_layer, drop=drop)
def forward(self, x):
# 最后一维归一化,multi-head attention, drop_path
# (B, N, C) -> (B, N, C)
x = x + self.drop_path(self.attn(self.norm1(x)))
# (B, N, C) -> (B, N, C)
x = x + self.drop_path(self.mlp(self.norm2(x)))
return x
在ViT中这样的block会有好几层,形成blocks:
# stochastic depth decay rule
dpr = [x.item() for x in torch.linspace(0, drop_path_rate, depth)]
self.blocks = nn.Sequential(*[
Block(
dim=embed_dim, num_heads=num_heads, mlp_ratio=mlp_ratio, qkv_bias=qkv_bias, drop=drop_rate,
attn_drop=attn_drop_rate, drop_path=dpr[i], norm_layer=norm_layer, act_layer=act_layer)
for i in range(depth)])
如果drop_path_rate大于0,每一层block的drop_path的会线性增加。depth是一个blocks里block的数量。也可以理解为blocks这个网络块的深度。
极智AI | 详解 ViT 算法实现
你可能想问,为什么一定要先分patch,再从patch转token呢?
第一个原因,是为了减少模型计算量。
在Transformer中,假设输入的序列长度为N,那么经过attention时,计算复杂度就为,因为注意力机制下,每个token都要和包括自己在内的所有token做一次attention score计算。
在ViT中,,当patch尺寸P越小时,N越大,此时模型的计算量也就越大。因此,我们需要找到一个合适的P值,来减少计算压力。
第二个原因,是图像数据带有较多的冗余信息。
和语言数据中蕴含的丰富语义不同,像素本身含有大量的冗余信息。比如,相邻的两个像素格子间的取值往往是相似的。因此我们并不需要特别精准的计算粒度(比如把P设为1)。这个特性也是之后MAE之类的像素级预测模型能够成功的原因之一。
实验图刻画了ViT的16个multi-head attention学到的像素距离信息。横轴表示 网络的深度, 纵轴表示“平均注意力距离”, 我们设第 i i i 个和第 j j j 个像素的平均注意力 距离为 d i j d_{ij} dij, 真实像素距离为 d i j ′ d_{ij}^{\prime} dij′, 这两个像素所在patch某一个head上的attention score为 a i j a_{ij} aij, 则有: d i j = a i j ∗ d i j ′ d_{ij}=a_{ij}*d_{ij}^{\prime} dij=aij∗dij′。当 d i j d_{ij} dij越大时, 说明ViT的attention机制能让它关 注到距离较远的两个像素, 类似于CNN中的“扩大感受野”。
图中每一列上,都有16个彩色原点,它们分别表示16个head观测到的平均像素距离。由图可知,在浅层网络中,ViT还只能关注到距离较近的像素点,随着网络加深,ViT逐渐学会去更远的像素点中寻找相关信息了。这个过程就和用在CNN中用卷积逐层去扩大感受野非常相似。
下图的左侧表示原始的输入图片,右侧表示ViT最后一层看到的图片信息,可以清楚看见,ViT在最后一层已经学到了将注意力放到关键的物体上了,这是非常有趣的结论:
图像的空间局部性(locality),即有相关性的物体(例如太阳和天空)经常一起出现。CNN采用卷积框取特征的方式,极大程度上维护了这种特性。其实,ViT也有维护这种特性的方法,上面所说的attention是一种,位置编码也是一种。
上图是ViT-L/32
模型下的位置编码信息,图中每一个方框表示一个patch,图中共有7_7个patch。而每个方框内,也有一个7_7的矩阵,这个矩阵中的每一个值,表示当前patch的position embedding和其余对应位置的position embedding的余弦相似度。颜色越黄,表示越相似,也即patch和对应位置间的patch密切相关。
注意到每个方框中,最黄的点总是当前patch所在位置,这个不难理解,因为自己和自己肯定是最相似的。除此以外颜色较黄的部分都是当前patch所属的行和列,以及以当前patch为中心往外扩散的一小圈。这就说明ViT通过位置编码,已经学到了一定的空间局部性。
在工业界,人们的标注数据量和算力都是有限的,因此CNN可能还是首要选择。但是,ViT的出现,不仅是用模型效果来考量这么简单,今天再来看这个模型,发现它的意义在于: