新闻  |   论坛  |   博客  |   在线研讨会
TPAMI 2022|金字塔池化的骨干网络,各大任务都涨点!南开&达摩院联合推出P2T
计算机视觉工坊 | 2022-09-26 06:53:40    阅读:878   发布文章

作者丨吴宇寰

编辑丨极市平台

导读

 

该工作主要受PVT和MViT的启发,第一次将金字塔池化融入到视觉transformer的骨干网络中,在减小输入序列长度的同时去捕捉多层次的丰富特征表达。P2T在图像分类、语义分割、目标检测、实例分割等多个领域中,并取得了比现有CNN/Transformer骨干更优异的性能。

图片

作者列表:吴宇寰,刘云,占新,程明明

作者单位:南开大学,阿里达摩院

代码位置:https://github.com/yuhuan-wu/P2T

论文地址:https://arxiv.org/abs/2106.12011

IEEE地址:https://ieeexplore.ieee.org/document/9870559

中文版本:https://mmcheng.net/wp-content/uploads/2022/08/22PAMI_P2T_CN.pdf

1. 引言Transformer技术在计算机视觉领域的兴起

近年来,NLP中的transformer技术在计算机视觉领域掀起新的热潮。Transformer可以充分地捕捉长距离的全局特征关系,正好满足了计算机视觉大感受野的需求。视觉transformer [1] 也在近年来的各大领域取得了巨大的成功,证明了纯transformer架构在巨大的数据规模下训练也能在大规模的ImageNet分类数据集上与CNN方法取得相同的分类效果。DeiT [2] 则随后证明了使用知识蒸馏可以只在ImgeNet-1K分类数据集上就能与CNN取得类似的效果。

然而与NLP领域特征长度有限不同,计算机视觉的图像具有密集特性,其拉直(flattened)后的特征长度要远远大于NLP任务中的序列长度。比如如WMT2014数据集 [3] 的平均序列长度为25,ImageNet-1K数据集 [4] 的一般使用图像分辨率为224x224的大小,即使使用常规的倍降采样,拉直后的图像序列长度仍有3136。而Transformer因其需要捕捉充分的全局注意力关系,其多头自我注意力(multi-head self-attention, MHSA)的计算复杂度是与输入特征呈平方关系。CNN网络因为主要捕捉局部信息,其计算复杂度只与图像的特征长度呈线性相关。

因此,在计算机视觉领域中应用transformer技术的关键困难在于如何使计算注意力关系的过程更加高效。视觉transformer [1]作为第一个解决方案,首先提出使用大的降采样尺度(如16、32倍)来对输入特征进行降采样,并设计了一个纯transformer的网络骨干(backbone),并在图像分类任务上取得了巨大的成功。而随后PVT [5]、Swin [6]、MViT [7] 等金字塔结构的网络骨干使用类似CNN的逐渐下采样的手段,来支持各类下游视觉任务。

目前transformer网络骨干的思路

如上文所述,transformer应用到计算机视觉领域的关键问题在于如何使计算MHSA的过程更加高效。PVT和MViT对特征图采用了单个池化操作来减少MHSA的计算量和显存占用,相应地是进行输入特征图元素与不同区域之间的关系建模,而不是特征图元素之间的关系建模。Swin transformer则是另一种思路:它先将特征划分为小的区块,并在每个小区块内直接计算MHSA,即进行小区块内元素之间的关系建模;随后,它将区块划分进行偏移,同时计算偏移后的小区块的MHSA。这样,类似CNN网络一样,随着网络的堆叠,Swin transformer网络的感受野会逐步增强。然而,视觉transformer的最根本特性源于其对特征全局的关系建模,这也是为什么我们要尝试将transformer引入到过度由CNN统治的计算机视觉领域中。

本文的工作

我们的工作主要受PVT和MViT的启发,即使用特征降采样的形式来减少MHSA的计算量。PVT和MViT使用单个尺度的降采样后的特征来计算MHSA,这种特征表达的模式有限。我们想到,如果使用金字塔池化(pyramid pooling)来生成特征,既能够大大缩小特征序列长度,又能引入更丰富的多层次特征。

金字塔池化在计算机视觉中有着悠久的历史 [8-9]。它利用多个不同步长的池化操作来提取丰富不同感受野的池化特征,其在目标检测 [10] 和语义分割中 [11] 也被证明十分有效。然而,它们都依赖于一个CNN骨干网络,所以它们也只是被专门设计用于某个特定的任务。目前并没有探索金字塔池化在骨干网络设计的工作。而一个新型的骨干网络可以应用于多种下游任务中。我们第一次将金字塔池化融入到视觉transformer的骨干网络中,在减小输入序列长度的同时去捕捉多层次的丰富特征表达。

利用金字塔池化,我们设计了基于的金字塔池化的MHSA(pooling-based MHSA, P-MHSA)。金字塔池化作为计算P-MHSA中降低序列长度的基本操作。使用P-MHSA替换掉MHSA后,即为P2T基础模块。通过堆叠不同数量的P2T基础模块,我们设计了P2T-Tiny/Small/Base/Large骨干网络,分别对应ResNet-18/50/101/152。P2T系列的多层次特征捕捉能力也要强于之前的PVT与MViT方法。我们也在图像分类、语义分割、目标检测、实例分割等多个领域中,并取得了比现有CNN/Transformer骨干更优异的性能,如ResNet, ResNeXt, Res2Net, PVT, Swin, Twins, 及PVTv2。

图片

图 1 P2T与其他方法在ADE20K数据集上的对比图(参数量、mIoU精度)。

2. P2T 方法

P2T的总体结构如图 2所示。每个阶段开始时有一个patch embedding模块,用于对输入的特征进行下采样。第一个Patch Embedding模块通过步长为4的7x7 卷积将输入特征下采样为原来的1/4,其他Patch Embedding则是将输入特征下采样为原来的2倍。每个Patch Embedding模块之后则堆叠了不同数量的P2T基础模块。其中,P2T基本模块的计算可由下式表示:

图片图片

图 2  P2T的总体架构。它有4个阶段,分别输出1/4、1/8、1/16和1/32大小的特征图。其中1/32大小的特征用于图像分类,所有的特征用于诸多下游任务。

我们通过调整P2T基础模块的数量,就可以得到不同规模的P2T模型,如图 3所示。

图片

图 3 不同设置下的P2T骨干网络。

基于金子塔池化的MHSA (P-MHSA)图片

图 4 P2T基本模块架构图。(a) P2T模块粗略图;(b) 基于金字塔池化的多头自我注意力(P-MHSA)架构。

计算流程。 对于给定的特征,先将用不同步长的池化操作下采样到不同的大小:

图片

可以得到不同大小的特征图。因为这些特征图是经池化后的特征,对每个特征作相对的位置编码(RPE)几乎不产生计算量:

图片

其中DWConv代表深度3x3卷积,即分组数等于通道数的3x3卷积。我们将这些编码后的特征进行拉直、特征拼接后,再分别生成计算自我注意力(self-attention)需要的query (Q)、key(K)、value(V):

图片

获得 、、 后, 就可以计算自我注意力  :


其中  为  的通道数, 用作尺度缩放。因为  和  的特征长度相对  要小得多, P-MHSA 的计算量也要相对 MHSA 小得多。

计算复杂度。 P-MHSA中主要由金字塔池化、自我注意力和一些线性层构成。设N和C分别为输入特征的长度和通道数。金字塔池化所占用的计算量仅为O(NC)。同时,也可以计算出P-MHSA的计算复杂度为:

图片

其中M为经金字塔池化采样拼接后的特征长度。对于默认的{12,16,20,24}的池化参数,M约为,与PVT大致相同。

前馈网络(Feed-forward Network, FFN)

在transformer中,前馈网络(FFN)是一种基础的特征增强模块。然而传统FFN仅用数个全连接层在特征的通道维度进行增强,其并没有考虑到图像是2维特征。通过将2D局部特征学习引入到FFN中,可以进一步提高场景理解所需要的特征表达能力。相对于普通的transformer使用的FFN,我们额外添加了一个深度3x3卷积来额外学习2D局部表征,即与MobileNetV2中的反向瓶颈模块(IRB)一致。假设  是已转化为 2 维的 ,,对原FFN引入2D局部特征学习:

图片3. 实验部分

我们在图像分类、语义分割、目标检测、实例分割上与其他方法进行了效果、计算量、运行速度上的对比。通过下列的实验可以发现我们的方法要显著优于现有的方法,如Swin、Twins、MViT、PVTv2等。

图像分类图片

图 5 图像分类结果。输入大小统一为224x224。FPS是在RTX 2070下测试的结果。

语义分割图片

图 6 在语义分割上的结果。结果基于Semantic FPN。输入大小为512x512。所有实验基于batchsize 16的情况下训练8w次。

目标检测与实例分割图片

图 7 目标检测与实例分割上的结果。图中所有的方法都基于RetinaNet和Mask R-CNN。所有实验基于在coco上训练12 epochs后的结果(无多尺度训练)。

消融实验图片

图 8 金字塔池化的设置。可以发现,在相同下采样系数(D Ratio)的前提下,金字塔池化要显著优于单尺度池化。

图片

图 9 相对于单尺度池化,该图显示的是在不同阶段下应用金字塔池化自注意力的效果。可以发现,随着金字塔池化在更多的阶段下的应用,效果也变的越好。

图片

图 10 进一步应用相对位置编码(RPE)、IRB和部分重叠的patch embedding(OPE)所带来的效果。RPE可以进一步大幅提升分类和效果,同时RPE的计算量仅占网络总计算量的0.1%(0.005 GFLOPS)。IRB则进一步显著地提升了分类和分割效果。OPE虽然进一步提升的分类相对前两者不多,但能显著提升分割效果。

此外,P2T采用了更轻量的Hardswish激活,GELU与其相比,训练内存占用额外多出了52%。我们也发现,采用Hardswish激活的P2T和采用GELU的P2T性能一致。所以,我们使用Hardswish激活,使得P2T可以在更小显存上的显卡上训练。

4. 总结

我们提出了基于金字塔池化的transformer基本单元P2T。相对于单尺度池化,金字塔池化既大幅减少了transformer中MHSA的计算量,又提供了强大的多层次的特征表达。我们利用P2T基本单元建立了一系列的P2T骨干:P2T-Tiny/Small/Base/Large。在图像分类、语义分割、目标检测、实例分割等任务上,基于P2T的方法也显著要好于其他骨干网络,如ResNet, ResNeXts, Res2Nets, 以及PVT、Swin、Twins、PVTv2。

相关文献

[1] An image is worth 16x16 words: Transformers for image recognition at scale, ICLR 2021

[2] Training data-efficient image transformers & distillation through attention, ICLR 2021

[3] Findings of the 2014 workshop on statistical machine translation, The Workshop on Statistical Machine Translation, 2014.

[4] ImageNet large scale visual recognition challenge, IJCV 2015

[5] Pyramid Vision Transformer: A versatile backbone for dense prediction without convolutions, ICCV 2021

[6] Swin Transformer: Hierarchical vision transformer using shifted windows, ICCV 2021

[7] Twins: Revisiting the design of spatial attention in vision transformers, NeurIPS 2021

[8] The pyramid match kernel: Discrim- inative classification with sets of image features, ICCV 2005.

[9] Beyond bags of features: Spatial pyramid matching for recognizing natural scene cate- gories, CVPR 2006.

[10] Spatial Pyramid Pooling in Deep Convolutional Networks for Visual Recognition, TPAMI 2015

[11] Pyramid Scene Parsing Network, CVPR 2017

本文仅做学术分享,如有侵权,请联系删文。


*博客内容为网友个人发布,仅代表博主个人观点,如有侵权请联系工作人员删除。

参与讨论
登录后参与讨论
推荐文章
最近访客