关于美团:NeurIPS-2021-|-Twins重新思考高效的视觉注意力模型设计

42次阅读

共计 8012 个字符,预计需要花费 21 分钟才能阅读完成。

Twins 是美团和阿德莱德大学单干提出的视觉注意力模型,相干论文已被 NeurIPS 2021 会议接管。本文次要讲述 Twins 解决的难点、设计和实现思路,以及在美团场景的摸索落地,心愿能对从事视觉算法研发的同学有所帮忙和启发。

概要

Twins [1] 是美团和阿德莱德大学单干提出的视觉注意力模型,相干论文已被 NeurIPS 2021 会议接管,代码也已在 GitHub 上进行开源。NeurIPS(Conference on Neural Information Processing Systems)是机器学习和计算神经科学相干的学术会议,也是人工智能方向的国内顶级会议。

Twins 提出了两类构造,别离是 Twins-PCPVT 和 Twins-SVT:

  • Twins-PCPVT 将金字塔 Transformer 模型 PVT [2] 中的固定地位编码(Positional Encoding)更改为团队在 CPVT [3] 中提出的条件式地位编码(Coditional Position Encoding, CPE),从而使得模型具备平移等变性(即输出图像产生平移后,输入同时相应发生变化),能够灵活处理来自不同空间尺度的特色,从而可能广泛应用于图像宰割、检测等变长输出的场景。
  • Twins-SVT 提出了空间可拆散自注意力机制(Spatially Separable Self-Attention,SSSA)来对图像特色的空间维度进行分组,别离计算各部分空间的自注意力,再利用全局自注意力机制对其进行交融。这种机制在计算上更高效,性能更优。

Twins 系列模型实现简略,部署敌对,在 ImageNet 分类、ADE20K 语义宰割、COCO 指标检测等多个经典视觉工作中均获得了业界当先的后果。

背景

2020 年 9 月,谷歌的视觉注意力模型(Vision Transformer,ViT)[4] 胜利将本来用于自然语言解决的 Transformer [5] 利用到视觉的分类工作中。ViT 将一幅输出图像切分为若干个图像块(Patch),并把一个图像块类比为一个文字(Word)作为 Transformer 编码器的输出(如图 1 所示),通过 L 层的编码器解决后应用一般的多层感知机(Multilayer Perceptron, MLP)映射到类别空间。ViT 的模型性能大幅超过了卷积神经网络,尔后迅速倒退成为了以后视觉畛域钻研的次要热点。

![图 1 视觉注意力模型(ViT)将用于自然语言解决工作的 Transformer 利用于视觉工作(起源:ViT [4])](https://p1.meituan.net/travel…)

Transformer 编码器中多头注意力(Multi-head attention)的根本计算方法由下式给定,其中 Q、K、V 别离为 Query(查问)、Key(键)、Value(值)的缩写,d 为编码维度,softmax 为归一化函数,注意力机制能够了解为对输出依照相关性加权的过程。

原生的视觉注意力模型做骨干网络并不能很好地适配指标检测、语义宰割等罕用的浓密预测工作。此外,相比于卷积神经网络,ViT 计算量通常要更大,推理速度变慢,不利于在理论业务中利用。因而设计更高效的视觉注意力模型,并更好地适配上游工作成为了当下钻研的重点。香港大学、商汤联结提出的金字塔视觉注意力模型 PVT [2] 借鉴了卷积神经网络中的图像金字塔范式来生成多尺度的特色,这种构造能够和用于浓密工作的现有后端间接联合,反对多种上游工作,如图 2(c)所示。但因为 PVT 应用了动态且定长的地位编码,通过插值形式来适应变长输出,不能针对性依据输出特色来编码,因而性能受到了限度。另外,PVT 沿用了 ViT 的全局自注意力机制,计算量仍然较大。

![图 2 PVT 将卷积神经网络(a)的金字塔范式迁徙到视觉注意力模型(b)失去(c),以适应分类、检测、宰割多种工作(起源:PVT [2])](https://p0.meituan.net/travel…)

微软亚研院提出的 Swin [6] 复用了 PVT 的金字塔构造。在计算自注意力时,应用了对特色进行窗口分组的办法(如图 3 所示),将注意力机制限定在一个个小的窗口(红色格子),而后通过对窗口进行错位使不同组的信息产生交互。这样能够防止计算全局自注意力而缩小计算量,其毛病是损失了全局的注意力,同时因为窗口错位产生的信息交互能力绝对较弱,肯定水平上影响了性能。

![图 3 Swin 计算每一个红色格子的部分自注意力,通过不同层间的窗口移位来使各部分注意力之间产生交互(起源:Swin [6])](https://p0.meituan.net/travel…)

视觉注意力模型设计的难点

简略总结一下,以后视觉注意力模型设计中须要解决的难点在于:

  • 高效率的计算:放大和卷积神经网络在运算效率上的差距,促成理论业务利用;
  • 灵便的注意力机制:即可能具备卷积的部分感触野和自注意力的全局感触野能力,兼二者之长;
  • 利于上游工作:反对检测、宰割等上游工作,尤其是输出尺度变动的场景。

Twins 模型设计

从这些难点问题登程,基于对以后视觉注意力模型的粗疏剖析,美团视觉智能部从新思考了自注意力机制的设计思路,提出了针对性的解决方案。首先将 PVT [2] 和 CPVT [4] 相结合,造成 Twins-PCPVT 来反对尺度变动场景的上游工作。再从自留神机制的效率和感触野角度登程,设计了兼容部分和全局感触野的新型自注意力,叫做 空间可拆散自注意力(Spatially Separable Self-Attention,SSSA),造成了 Twins-SVT。

Twins-PCPVT

Twins-PCPVT 通过将 PVT 中的地位编码(和 DeiT [7] 一样固定长度、可学习的地位编码)替换为 CPVT [4] 中的条件地位编码(Conditional Positional Encodings,CPE)。生成 CPE 的模块叫做地位编码器(Positional Encoding Generator,PEG),PEG 在 Twins 模型中的具体位置是在每个阶段的第 1 个 Transformer Encoder 之后,如下图 4 所示:

条件地位编码

下图 5 展现了团队在 CPVT [4] 中提出的条件地位编码器的编码过程。首先将 $N*d$ 的输出序列转为 $H*W*d$ 的输出特色,再用 $F$ 依据输出进行条件式的地位编码,而且输入尺寸和输出特色雷同,因而能够转为 $N*d$ 序列和输出特色进行逐元素的加法交融。

其中,编码函数 $F$ 能够由简略的深度可拆散卷积实现或者其余模块实现,PEG 局部的简化代码如下。其中输出 feat_token 为形态为 $B*N*d$ 的张量,$B$ 为 batch,$N$ 为 token 个数,$C$ 为编码维度(同图 5 中 $d$)。将 feat_token 转化为 $B*d*H*W$ 的张量 cnn_feat 后,通过深度可拆散卷积 (PEG) 运算,生成和输出 feat_token 雷同形态的张量,即条件式的地位编码。

class PEG(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)
        
    def forward(self, feat_token, H, W):
        B, N, C = feat_token.shape
        cnn_feat = feat_token.transpose(1, 2).view(B, C, H, W) 
        x = self.peg(cnn_feat) + cnn_feat
        x = x.flatten(2).transpose(1, 2)
        return x

因为条件地位编码 CPE 是依据输出生成,反对可变长输出,使得 Twins 可能灵活处理来自不同空间尺度的特色。另外 PEG 采纳卷积实现,因而 Twins 同时保留了其平移等变性,这个性质对于图像工作十分重要,如检测工作中指标产生偏移,检测框需随之偏移。试验表明 Twins-PCPVT 系列模型在分类和上游工作,尤其是在浓密工作上能够间接取得性能晋升。该架构阐明 PVT 在仅仅通过 CPVT 的条件地位编码加强后就能够取得很不错的性能,由此阐明 PVT 应用的地位编码限度了其性能施展。

Twins-SVT

Twins-SVT(如下图 6 所示)对全局注意力策略进行了优化改良。全局注意力策略的计算量会随着图像的分辨率成二次方增长,因而如何在不显著损失性能的状况下升高计算量也是一个钻研热点。Twins-SVT 提出新的交融了部分 - 全局注意力的机制,能够类比于卷积神经网络中的深度可拆散卷积(Depthwise Separable Convolution),并因而命名为空间可拆散自注意力(Spatially Separable Self-Attention,SSSA)。与深度可拆散卷积不同的是,Twins-SVT 提出的空间可拆散自注意力(如下图 7 所示)是对特色的空间维度进行分组,并计算各组内的自注意力,再从全局对分组注意力后果进行交融。

空间可拆散自注意力采纳部分 - 全局自注意力(LSA-GSA)互相交替的机制,分组计算的部分注意力能够高效地传导到全局。LSA 能够大幅升高计算成本,复杂度从输出的平方 $O(H^2W^2d)$ 降为线性的 $O(mnHWd)$。其中分组部分注意力 LSA 要害实现(初始化函数略)如下:

class LSA(nn.Module):
    def forward(self, x, H, W):
        B, N, C = x.shape
        h_group, w_group = H // self.ws, W // self.ws # 依据窗口大小计算长(H)和宽(W)维度的分组个数
        total_groups = h_group * w_group
        x = x.reshape(B, h_group, self.ws, w_group, self.ws, C).transpose(2, 3) # 将输出依据窗口进行分组 B* h_group * ws * w_group * ws * C
        qkv = 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,v 
        q, k, v = qkv[0], qkv[1], qkv[2]
        attn = (q @ k.transpose(-2, -1)) * self.scale # 计算各组的注意力
        attn = attn.softmax(dim=-1) # 注意力归一化
        attn = self.attn_drop(attn) # 注意力 Dropout 层
        attn = (attn @ v).transpose(2, 3).reshape(B, h_group, w_group, self.ws, self.ws, C) # 用各组内的部分自注意力给 v 进行加权
        x = attn.transpose(2, 3).reshape(B, N, C)
        x = self.proj(x) # MLP 层
        x = self.proj_drop(x) # Dropout 层
        return x

高效交融 LSA 注意力的 GSA 要害实现(初始化函数略)如下。相比于 ViT 原始的全局自注意力,GSA 的 K、V 是在放大特色的根底上计算的,但 Q 是全局的,因而注意力依然能够复原到全局。这种做法显著缩小了计算量。

class GSA(nn.Module):
    def forward(self, x, H, W):
        B, N, C = x.shape
        q = self.q(x).reshape(B, N, self.num_heads, C // self.num_heads).permute(0, 2, 1, 3) # 依据输出特色 x 计算查问张量 q
        x_ = 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_) # 层归一化 LayerNorm
        kv = self.kv(x_).reshape(B, -1, 2, self.num_heads, C // self.num_heads).permute(2, 0, 3, 1, 4) # 依据放大尺寸后的特色后 x_,计算 k, v 
        k, v = 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) # 依据全局自注意力对 v 加权
        x = self.proj(x)
        x = self.proj_drop(x)
        return x

从上述代码中能够看出,SVT 系列在实现上采纳现有支流深度学习框架中的已有操作,不须要额定的底层适配,因而部署起来比拟不便。

试验

ImageNet-1k 分类

Twins-PCPVT 和 Twins-SVT 在 ImageNet-1k 分类工作上,相比同等量级模型均获得 SOTA 后果,吞吐率占优。另外,Twins 反对 TensorRT 部署,Twins-SVT-S 模型应用 NVIDIA TensorRT 7.0 推理能够有 1.6 倍的减速,吞吐率能够从 PyTorch 实现的 1059(images/s)晋升到 1732。

ADE20K 宰割

在语义宰割工作 ADE20K 上,Twins 模型做主干网别离应用 FPN 和 Upernet 后端,相比 PVT 和 Swin 也达到了更好后果,见下表 2:

COCO 指标检测(RetinaNet 框架)

在经典的 COCO 指标检测工作中,应用 RetinaNet 框架,Twins 模型大幅优于 PVT。而且 Twins-PCPVT 系列证实 PVT 在通过 CPVT 的编码方式加强之后,能够媲美 Swin 同量级模型,见下表 3:

COCO 指标检测(Mask-RCNN 框架)

在 Mask-RCNN 框架下,Twins 模型在 COCO 上也有很好的性能劣势,且在更长时间训练(3x)时得以放弃,见下表 4:

在高精地图多因素语义宰割场景的利用

高精地图是主动驾驶中的要害组成部分,在美团无人配送、网约车等业务承当着十分重要的作用。路线场景要害因素的语义提取作为高精建图的前序流程,对建图的品质有间接的影响。多因素语义宰割是语义提取的重要一环,业界个别采纳经典的语义宰割算法来实现。

此处,咱们以 DeepLab 系列 [8] 为代表做介绍,宰割模型通常分为编码和解码两个阶段,应用卷积神经网络提取特色,并采纳空间金字塔迟化(Spatial Pyramid Pooling),以及不同尺度空洞卷积(Atrous Conv)操作(如下图 8a 所示)来减少全局感触野。这种设计一方面受限于卷积神经网络的特征提取能力,另一方面对全局关系的建模能力无限,导致在宰割工作上对细节关注不够,边缘往往不够清晰。

![图 8 经典语义宰割模型架构(DeepLabV3+ [8])](https://p0.meituan.net/travel…)

Twins 尽管大幅晋升了视觉注意力模型的效率和性能,但为了放弃和卷积神经网络的靠近的推理效率,咱们仍须要对模型的后端构造作进一步的优化。不同于论文中为了与其余办法做偏心比照而应用的较重的 FPN [9] 或 UperNet [10] 后端,咱们设计了如下图 9 所示的简略轻量的后端,并在业务数据集上的性能和推理速度之间获得了很好的均衡。这个后端是依据 Twins 的个性而设计,因为 Twins 兼顾了全局和部分两种注意力,因而后端无需采纳简单的设计来增大感触野,只通过各尺度特色的线性变动和缩放,就间接复原到雷同尺寸并进行拼接(Concat),简略维度变换后就能够输入宰割后果。

从下图 10 的宰割后果比照看,Twins 为主干网的模型能够提取失去更精密的图像边缘,如隔离带、路线标牌、路灯杆等要害路线因素和标注真值(Ground Truth)之间的差别更小。

总结

视觉注意力模型是以后视觉畛域的钻研重点,并且曾经在各类视觉工作上展现了相比经典卷积神经网络的优越性,但在效率上仍需精心优化,在成果上也须要持续晋升。摸索设计更高效率的注意力模型并促成前沿的视觉钻研转向工业落地,对美团业务也具备重要的意义。

此次,美团和阿德莱德大学单干设计的 Twins 系列模型架构,无效升高了计算成本,晋升了模型性能,更好地反对了如检测和宰割等浓密工作。此外,咱们将 Twins 利用在美团高精地图的因素语义宰割场景中,带来了更精密的宰割后果,晋升了高精地图的建图品质。后续,视觉团队将继续摸索高效的视觉注意力模型设计,并冀望在美团更宽泛的业务场景中失去实际和利用。

参考文献

  • [1] Twins: Revisiting the Design of Spatial Attention in Vision Transformers
  • [2] Pyramid Vision Transformer: A Versatile Backbone for Dense Prediction without Convolutions
  • [3] Conditional Positional Encodings for Vision Transformers
  • [4] An image is worth 16×16 words: Transformers for image recognition at scale
  • [5] Attention Is All You Need
  • [6] Swin Transformer: Hierarchical Vision Transformer using Shifted Windows
  • [7] Training data-efficient image transformers & distillation through attention
  • [8] Encoder-Decoder with Atrous Separable Convolution for Semantic Image Segmentation
  • [9] Panoptic Feature Pyramid Networks
  • [10] Unified Perceptual Parsing for Scene Understanding

作者简介

祥祥、田值、张勃、晓林在视觉智能部,海兵、华夏在自动车配送部。

团队简介及招聘信息

美团视觉智能部 AutoML 算法团队旨在通过 AutoML 及前沿的视觉技术赋能公司各项业务、减速算法落地,涵盖 AutoML、宰割、检测(2D、3D)、Self-training 等技术方向。欢送感兴趣的校招和社招同学发送简历至:chuxiangxiang@meituan.com,实习、正式均可。

美团自动车配送部高精地图团队是高精地图技术研发团队,咱们的职责是为美团主动驾驶提供高精度、高鲜度、大范畴的高精地图服务。高精地图是波及多种学科的综合技术,不仅须要依靠 SLAM、天文测绘、深度学习、多传感器定位等算法进行高精地图的构建,而且还须要利用大数据技术、高性能计算、高并发服务等进行大规模高精地图解决、存储和查问服务。高精地图团队长期招聘计算机视觉、SLAM、零碎开发等专家,感兴趣的同学能够将简历发送至:tech@meituan.com(邮件主题:美团高精地图)。

浏览美团技术团队更多技术文章合集

前端 | 算法 | 后端 | 数据 | 平安 | 运维 | iOS | Android | 测试

| 在公众号菜单栏对话框回复【2021 年货】、【2020 年货】、【2019 年货】、【2018 年货】、【2017 年货】等关键词,可查看美团技术团队历年技术文章合集。

| 本文系美团技术团队出品,著作权归属美团。欢送出于分享和交换等非商业目标转载或应用本文内容,敬请注明“内容转载自美团技术团队”。本文未经许可,不得进行商业性转载或者应用。任何商用行为,请发送邮件至 tech@meituan.com 申请受权。

正文完
 0