当前位置: 首页 > 科技观察

双胞胎:重新思考高效的视觉注意模型设计

时间:2023-03-19 22:16:35 科技观察

作者|香香田志张博等。Twins是美团点评联合阿德莱德大学提出的视觉注意力模型,相关论文已被NeurIPS2021会议录用。本文主要介绍Twins所解决的难点、设计与实现思路,以及在美团场景的探索与落地,希望对从事视觉算法研发的同学有所帮助和启发。Twins[1]是美团点评与阿德莱德大学联合提出的视觉注意力模型。相关论文已被NeurIPS2021会议录用,代码也已在GitHub上开源。NeurIPS(ConferenceonNeuralInformationProcessingSystems)是机器学习和计算神经科学相关的学术会议,也是人工智能方向的国际顶级会议。Twins提出了两类结构,即Twins-PCPVT和Twins-SVT:Twins-PCPVT将金字塔Transformer模型PVT[2]中的固定位置编码(PositionalEncoding)改为CPVT[3]中团队提出的条件公式CoditionalPositionEncoding(CPE),使模型具有平移等变性(即输入图像平移后,输出同时发生变化),可以灵活处理来自不同空间尺度的特征,因此可以广泛应用在图像分割中,检测等变长度输入的场景。Twins-SVT提出空间可分离自注意力机制(SpatiallySeparableSelf-Attention,SSSA)对图像特征的空间维度进行分组,分别计算每个局部空间的自注意力,然后利用全局自注意力机制进行分开他们。进行融合。这种机制在计算上更有效,并提供更好的性能。Twins系列模型实现简单、部署友好,在ImageNet分类、ADE20K语义分割、COCO目标检测等多项经典视觉任务中取得了业界领先的成绩。背景2020年9月,谷歌的视觉注意力模型(VisionTransformer,ViT)[4]成功地将原本用于自然语言处理的Transformer[5]应用到视觉分类任务中。ViT将一张输入图像分成若干个图像块(Patch),并将一个图像块与一个词(Word)进行比较,作为Transformer编码器的输入(如图1所示),经过L层编码器的处理,然后使用一个普通的多层感知器(MultilayerPerceptron,MLP)映射到类别空间。ViT的模型性能已经大大超越了卷积神经网络,从此迅速发展成为当前视觉领域研究的一大热点。图1视觉注意力模型(ViT)将自然语言处理任务的Transformer应用于视觉任务(来源:ViT[4])Transformerencoder中Multi-headattention的基本计算方法由下式给出,其中Q,K,和V分别是Query(查询)、Key(键)、Value(值)的缩写,d是编码维度,softmax是归一化函数。注意力机制可以理解为输入根据相关性加权的过程。作为骨干网络的原生视觉注意力模型不太适合常用的密集预测任务,如对象检测和语义分割。此外,与卷积神经网络相比,ViT通常计算量较大,推理速度较慢,不利于实际业务中的应用。因此,设计更高效的视觉注意力模型,更好地适应下游任务成为当前研究的重点。香港大学与商汤科技联合提出的PyramidVisualAttentionModelPVT[2]借鉴了卷积神经网络中的图像金字塔范式,生成多尺度特征。这种结构可以直接与现有后端一起用于密集任务。组合后,支持多个下游任务,如图2(c)所示。但由于PVT采用静态定长位置编码,采用插值来适应变长输入,无法根据输入特征进行编码,性能受到限制。另外,PVT继续使用ViT的globalself-attention机制,计算量还是很大。图2PVT将卷积神经网络(a)的金字塔范式迁移到视觉注意力模型(b)得到(c)适应分类、检测、分割等各种任务(来源:PVT[2])微软亚洲研究所提出的Swin[6]重用了PVT的金字塔结构。在计算self-attention时,采用windowgroupingfeatures的方法(如图3所示),将attention机制限制在小窗口(红色格子)内,然后不同组的Information产生交互。这样可以避免计算globalself-attention,减少计算量。缺点是失去了全局注意力。同时,窗口错位带来的信息交互能力相对较弱,在一定程度上影响了性能。图3Swin计算每个红色网格的localself-attention,并通过在不同层之间移动窗口来实现localattention之间的交互(来源:Swin[6])。简单总结一下视觉注意力模型的设计难点,目前视觉注意力模型设计需要解决的难点有:高效计算:缩小与卷积神经网络在计算效率上的差距,以及促进实际业务应用;灵活的attention机制:即localreceptivefield和self-attention的globalreceptivefieldcapability是两者的强项;对下游任务有利:支持检测、分割等下游任务,尤其是在输入尺度变化的场景下。Twins模型设计正是从这些难题出发,在对当前视觉注意力模型进行详细分析的基础上,美团视觉智能部重新思考了self-attention机制的设计思路,并提出了有针对性的解决方案。首先,PVT[2]和CPVT[4]结合形成Twins-PCPVT,以支持规模变化场景中的下游任务。从self-attention机制的效率和感受野的角度出发,设计了一种兼容局部和全局感受野的新型self-attention,称为SpatiallySeparableSelf-Attention(SSSA),形成Twins-SVT。Twins-PCPVTTwins-PCPVT用CPVT[4]中的条件位置编码(ConditionalPositionalEncodings,CPE)代替了PVT中的位置编码(与DeiT[7]相同的固定长度、可学习的位置编码)。生成CPE的模块称为PositionalEncodingGenerator(PEG)。PEG在Twins模型中的具体位置在每个stage的第一个TransformerEncoder之后,如下图4所示:图4Twins-PCPVT-S模型结构,使用CPVT提出的位置编码器(PEG)条件位置编码图下图5展示了CPVT[4]团队提出的条件位置编码器的编码过程。首先将输入序列转化为输入特征,然后根据输入进行条件位置编码,输出大小与输入特征相同,因此可以转化为序列和输入特征进行element-wise加法融合。图5ConditionalPositionEncoder(PEG)其中编码功能可以通过简单的depthwiseseparableconvolution或其他模块来实现,PEG部分的简化代码如下。输入的feat_token是形状为的张量,是batch,是token的数量,是编码维度(与图5相同)。将feat_token转换为tensorcnn_feat后,进行深度可分离卷积(depthwiseseparableconvolution,PEG)运算,生成与输入feat_token形状相同的tensor,即条件位置编码。classPEG(nn.Module):def__init__(self,in_chans,embed_dim):super(PEG,self).__init__()self.peg=nn.Conv2d(in_chans,embed_dim,3,1,1,bias=True,groups=embed_dim)defforward(self,feat_token,H,W):B,N,C=feat_token.shapecnn_feat=feat_token.transpose(1,2).视图(B,C,H,W)x=self.peg(cnn_feat)+cnn_featx=x.flatten(2).transpose(1,2)returnx由于条件位置编码CPE是根据输入生成的,支持变长输入,所以Twins可以灵活处理来自不同空间尺度的特征。另外,PEG是通过卷积实现的,所以Twins也保留了它的平移等可变性。此属性对于图像任务非常重要。如果目标在检测任务中发生了移动,则检测框也需要相应地移动。实验表明,Twins-PCPVT系列模型可以直接提高分类和下游任务的性能,尤其是在密集任务中。该架构表明,PVT仅通过CPVT的条件位置编码进行增强后即可获得非常好的性能,这表明PVT使用的位置编码限制了其性能。Twins-SVTTwins-SVT(如下图6所示)优化了全局注意力策略。全局注意力策略的计算量会随着图像分辨率的增加呈二次方增长,因此如何在不显着损失性能的情况下降低计算量也是研究热点。Twins-SVT提出了一种结合local-globalattention的新机制,可以类比于卷积神经网络中的深度可分离卷积(DepthwiseSeparableConvolution),因此命名为空间可分离自注意力(SpatiallySeparableSelf-Attention).-注意,SSSA)。与depth-separableconvolution不同,Twins-SVT提出的spatiallyseparableself-attention(如下图7)是对特征的空间维度进行分组,计算每组内的self-attention,然后从分组注意力结果的全局融合。图6Twins-SVT-S模型结构,右边是两个相邻TransformerEncoder的组合图7Twins提出的Spatialseparableself-attentionmechanism(SSSA)Spatialseparableself-attention采用local-globalself-attention(LSA-GSA)alternate机制,可以将group计算出的localattention高效传递到global。LSA可以大大降低计算成本,复杂度从输入的平方降低到线性。其中,grouplocalattentionLSA(初始化函数省略)的关键实现如下:classLSA(nn.Module):defforward(self,x,H,W):B,N,C=x.shapeh_group,w_group=H//自己。ws,W//self.ws#根据窗口大小计算长(H)和宽(W)维度的组数total_groups=h_group*w_groupx=x.reshape(B,h_group,self.ws,w_group,self.ws,C).transpose(2,3)#根据窗口对输入进行分组B*h_group*ws*w_group*ws*Cqkv=self.qkv(x).reshape(B,total_groups,-1,3,self.num_heads,C//self.num_heads).permute(3,0,1,4,2,5)#计算q,k,vq,k,v=qkv[0],qkv[1],qkv[2]attn=(q@k.transpose(-2,-1))*self.scale#计算每组attentionattn=attn.softmax(dim=-1)#Attentionnormalizationattn=self.attn_drop(attn)#AttentionDropoutlayerattn=(attn@v).transpose(2,3).reshape(B,h_group,w_group,self.ws,self.ws,C)#在每组Attention中使用局部selfweightsvx=attn.transpose(2,3).reshape(B,N,C)x=self.proj(x)#MLP层x=self.proj_drop(x)#Dropout层returnx高效融合GSA的关键实现对于LSAattention(初始化函数省略)如下。与ViT原来的globalself-attention相比,GSA的K和V是在减少特征的基础上计算的,但是Q是全局的,所以attention还是可以恢复到全局的。这种方法大大减少了计算量。classGSA(nn.Module):defforward(self,x,H,W):B,N,C=x.shapeq=self.q(x).reshape(B,N,self.num_heads,C//self.num_heads).permute(0,2,1,3)#根据输入特征x计算查询张量qx_=x.permute(0,2,1).reshape(B,C,H,W)x_=self.sr(x_).reshape(B,C,-1).permute(0,2,1)#减小输入特征的大小得到x_x_=self.norm(x_)#层归一化LayerNormkv=self.kv(x_).reshape(B,-1,2,self.num_heads,C//self.num_heads).permute(2,0,3,1,4)#根据x_缩小后计算k,vk,vfeatures=kv[0],kv[1]attn=(q@k.transpose(-2,-1))*self.scale#计算全局自注意力attn=attn.softmax(dim=-1)attn=self.attn_drop(attn)x=(attn@v).transpose(1,2).reshape(B,N,C)#根据globalself-attentionx=self.proj(x)x=self加权v。proj_drop(x)returnx从上面的代码可以看出,SVT系列采用了现有主流深度学习框架中已有的操作,不需要额外的底层适配,因此更方便dep洛伊。实验ImageNet-1k分类Twins-PCPVT和Twins-SVT在ImageNet-1k分类任务上相比同量级模型都取得了SOTA的成绩,吞吐率更胜一筹。此外,Twins支持TensorRT部署。Twins-SVT-S模型使用NVIDIATensorRT7.0推理可以加速1.6倍,吞吐率可以从PyTorch达到的1059(images/s)提高到1732。表1ImageNet-1k分类ADE20K分割在语义分割任务ADE20K上,以Twins模型为骨干网络分别使用FPN和Uppernetbackend,取得了比PVT和Swin更好的效果,见下表2:表2ADE20KsegmentationCOCOTargetdetection(RetinaNetframework)在经典的COCO目标检测任务中,使用RetinaNet框架,Twins模型比PVT好很多。而且,Twins-PCPVT系列证明了PVT经过CPVT编码方式增强后,可以和同量级的Swin模型进行对比,如下表3所示:表3COCO目标检测(RetinaNet框架)COCO目标Mask-RCNN框架下的检测(Mask-RCNN框架),Twins模型在COCO上也有很好的性能优势,在更长时间的训练(3x)中也能保持,见下表4:表4COCO目标检测(Mask-RCNN框架)在高精度地图多元素语义分割场景应用高精度地图是自动驾驶的关键组成部分,在美团无人配送、网约车等业务中发挥着非常重要的作用。道路场景关键要素语义提取作为高精度制图的前置过程,直接影响制图质量。多元素语义分割是语义抽取的重要环节,业界普遍采用经典的语义分割算法来实现。这里,我们以DeepLab系列[8]为代表进行介绍。分割模型通常分为编码和解码两个阶段,使用卷积神经网络提取特征,并使用空间金字塔池化(SpatialPyramidPooling),以及不同尺度的AtrousConv操作(如下图8a所示)来增加全局感受野。这种设计一方面受限于卷积神经网络的特征提取能力,另一方面全局关系的建模能力有限,导致在分割任务中对细节的关注不够,边缘往往不够清晰。图8经典语义分割模型架构(DeepLabV3+[8])虽然Twins在视觉注意力模型的效率和性能上有了很大的提升,但为了保持与卷积神经网络相似的推理效率,我们仍然需要修改back-模型的最终结构。进一步优化。与论文中使用的较重的FPN[9]或UpperNet[10]后端不同,为了与其他方法进行公平比较,我们设计了一个简单轻量级的后端,如下图9所示,并在业务数据集上进行了测试,取得了很好的平衡在性能和推理速度之间取得了平衡。这个后端是根据Twins的特点设计的。由于Twins兼顾了全局和局部注意力,后端不需要使用复杂的设计来增加感受野。可以直接还原为相同大小并进行拼接(Concat),分割结果经过简单的维度变换即可输出。图9LightweightTwins分割后端设计从下图10的分割结果对比来看,以Twins为骨干网络的模型可以提取更精细的图像边缘,比如隔离带、路标、路灯杆等关键道路elements和labelsgroundtruth之间的差异较小。图10道路多元素语义抽取结果对比总结视觉注意力模型是当前视觉领域研究的重点,在各种视觉任务中已经证明了其优于经典卷积神经网络的优势,但在视觉注意力方面仍需谨慎关注效率方面的优化,在效果方面也需要继续提升。探索设计更高效的注意力模型,将视觉前沿研究推进到产业落地,对美团业务也具有重要意义。此次,美团点评与阿德莱德大学联合设计的Twins系列模型架构有效降低了计算成本,提升了模型性能,更好地支持了检测、分割等密集任务。此外,我们在美团高精地图的特征语义分割场景中应用了Twins,带来了更精细的分割结果,提升了高精地图的构建质量。未来,视觉团队将继续探索高效的视觉注意力模型设计,期待在美团更广泛的业务场景中得到实践和应用。