TPAMI 2022｜金字塔池化的骨干网络，各大任务都涨点！南开&达摩院联合推出P2T

作者丨吴宇寰

编辑丨极市平台

导读

 

该工作主要受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

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

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