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

原文链接：https://www.techbeat.net/article-info?id=4057
作者：吴宇寰

现有的视觉Transformer中的两大问题：

1. 传统的Multi-Head Self-Attention (MHSA) 需要大量的计算、空间资源。

2. 视觉Transformer在图像分类中被过度地开发和调整，却忽视了图像分类（单一场景，与MLP较为相似）和各类下游场景理解任务（复杂场景，丰富结构和内容信息）的区别。

为了解决以上两大问题，研究人员提出了基于金字塔池化的视觉Transformer，第一次将金字塔池化思想引入到视觉Transformer中，从而减少使用传统MHSA带来的过高计算量和存储空间（问题1）。此外，研究人员注意到pyramid pooling因其强大的抽象上下文能力在各类下游视觉任务上的表现都十分出色，且其空间不变性的自然属性适合解决结构信息的丢失问题（问题2）。

论文链接：
https://arxiv.org/abs/2106.12011
中文版本：
https://mmcheng.net/wp-content/uploads/2022/08/22PAMI_P2T_CN.pdf
代码链接：
https://github.com/yuhuan-wu/P2T

一、引言

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精度）。

二、P2T方法

P2T的总体结构如图2所示。每个阶段开始时有一个patch embedding模块，用于对输入的特征进行下采样。第一个Patch Embedding模块通过步长为4的7x7 卷积将输入特征下采样为原来的1/4，其他Patch Embedding则是将输入特征下采样为原来的2倍。每个Patch Embedding模块之后则堆叠了不同数量的P2T基础模块。其中，P2T基本模块的计算可由下式表示：
$\begin{array}{rl}{\mathbf{X}}_{att}& =\text{LayerNorm}(\mathbf{X}+\mathrm{P}-\text{MHSA}(\mathbf{X}))\\ {\mathbf{X}}_{out}& =\text{LayerNorm}\left({\mathbf{X}}_{att}+\text{FFN}\left({\mathbf{X}}_{att}\right)\right)\end{array}$

图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）架构。

计算流程。 对于给定的特征 $\mathbf{X}$ ，先将 $\mathbf{X}$ 用不同步长的池化操作下采样到不同的大小：

$\begin{array}{rl}{\mathbf{P}}_1& ={\text{AvgPool}}_1(\mathbf{X})\\ {\mathbf{P}}_2& ={\text{AvgPool}}_2(\mathbf{X})\\ & \cdots \\ {\mathbf{P}}_n& ={\text{AvgPool}}_n(\mathbf{X})\end{array}$

可以得到不同大小的特征图 { P 1 , P 2 , … , P n } \{\mathbf{P}_1,\mathbf{P}_2,…,\mathbf{P}_n\} {P1​,P2​,…,Pn​} 。因为这些特征图是经池化后的特征，对每个特征作相对的位置编码（RPE）几乎不产生计算量：

${\mathbf{P}}_i^{enc}=\text{DWConv}\left({\mathbf{P}}_i\right)+{\mathbf{P}}_i,i=1,2,\cdots ,n$

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

$\mathbf{P}=\text{LayerNorm}\left(\text{Concat}\left({\mathbf{P}}_1^{enc},{\mathbf{P}}_2^{enc},\dots ,{\mathbf{P}}_n^{enc}\right)\right),$

$(\mathbf{Q},\mathbf{K},\mathbf{V})=\left({\mathbf{XW}}^q,\mathbf{P}{\mathbf{W}}^k,{\mathbf{PW}}^v\right),$

获得 $\mathbf{Q}$ 、 $\mathbf{K}$ 、 $\mathbf{V}$ 后, 就可以计算自我注意力 $\mathbf{A}$ :

$\mathbf{A}=\text{Softmax}\left(\frac{\mathbf{Q}\times {\mathbf{K}}^{\mathbf{T}}}{\sqrt{d_K}}\right)\times \mathbf{V},$

其中 $d_k$ 为 $\mathbf{K}$ 的通道数, 用作尺度缩放。因为 $\mathbf{K}$ 和 $\mathbf{V}$ 的特征长度相对 $\mathbf{Q}$ 要小得多, P-MHSA 的计算量也要相对 MHSA 小得多。

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

$O(\mathrm{P}-\mathrm{MHSA})=(N+2M)C^2+2NMC,$

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

前馈网络（Feed-forward Network, FFN）

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

${\mathbf{X}}_{\mathrm{IRB}}^1=\text{Act}\left({\mathbf{X}}_{att}^I{\mathbf{W}}_{\mathrm{IRB}}^1\right),$

${\mathbf{X}}_{\mathrm{IRB}}^{\text{out }}=\text{Act}\left(\text{DWConv}\left({\mathbf{X}}_{\mathrm{IRB}}^1\right)\right){\mathbf{W}}_{\mathrm{IRB}}^2,$

三、实验部分

我们在图像分类、语义分割、目标检测、实例分割上与其他方法进行了效果、计算量、运行速度上的对比。通过下列的实验可以发现我们的方法要显著优于现有的方法，如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可以在更小显存上的显卡上训练。

四、总结
我们提出了基于金字塔池化的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

Illustration by Dani Grapevine from icons8

-The End-

关于我“门”
▼

将门是一家以专注于发掘、加速及投资技术驱动型创业公司的新型创投机构，旗下涵盖将门创新服务、将门-TechBeat技术社区以及将门创投基金。

将门成立于2015年底，创始团队由微软创投在中国的创始团队原班人马构建而成，曾为微软优选和深度孵化了126家创新的技术型创业公司。

如果您是技术领域的初创企业，不仅想获得投资，还希望获得一系列持续性、有价值的投后服务，欢迎发送或者推荐项目给我“门”：

bp@thejiangmen.com