PVT(Pyramid Vision Transformer)

date

May 5, 2022

Last edited time

Mar 27, 2023 09:01 AM

status

Published

slug

PVT(Pyramid_Vision_Transformer)

PVT

Introduce

ViT存在的问题

CNN 结构常用的是一种金字塔架构，如上图所示，CNN 网络一般可以划分为不同的 stage，在每个 stage 开始时，特征图的长和宽均减半，而特征维度（channel）扩宽 2 倍。这主要有两个方面的考虑，一是采用 stride=2 的卷积或者池化层对特征降维可以增大感受野，另外也可以减少计算量，但同时空间上的损失用 channel 维度的增加来弥补。

但是 ViT 本身就是全局感受野，所以 ViT 就比较简单直接了，直接将输入图像 tokens 化后就不断堆积相同的 transformer encoders，这应用在图像分类上是没有太大的问题。但是如果应用在密集任务上，会遇到问题：

分割和检测往往需要较大的分辨率输入，当输入图像增大时，ViT 的计算量会急剧上升；

ViT 直接采用较大 patchs 进行 token 化，如采用 16x16 大小那么得到的粗粒度特征，对密集任务来说损失较大。

这正是 PVT 想要解决的问题，PVT 采用和 CNN 类似的架构，将网络分成不同的 stages，每个 stage 相比之前的 stage 特征图的维度是减半的，这意味着 tokens 数量减少 4 倍。

主要贡献

将细粒度图像块（即每个块 4×4 像素）作为输入来学习高像素，克服了传统 Transformer 的困难：大分辨率下的预测任务，这对于密集预测任务至关重要；
随着网络的深入，引入渐进式收缩金字塔来减少 Transformer 的序列长度，显着降低计算成本。
采用空间还原注意（SRA）层，进一步降低学习高分辨率特征时的资源消耗量。

Method

PVT网络结构

每个 stage 的输入都是一个维度的特征图，对于第 1 个 stage，输入就是 RGB 图像，对于其它 stage 可以将 tokens 重新 reshape 成特征图。在每个 stage 开始，首先像 ViT 一样对输入图像进行 token 化，即进行 patch embedding，patch 大小均采用大小（第 1 个 stage 的 patch 大小是），这意味着该 stage 最终得到的特征图维度是减半的，tokens 数量对应减少倍。PVT 共个 stage，这和 ResNet 类似，个 stage 得到的特征图相比原图大小分别是，，和。由于不同的 stage 的 tokens 数量不一样，所以每个 stage 采用不同的 position embeddings，在 patch embed 之后加上各自的 position embedding，当输入图像大小变化时，position embeddings 也可以通过插值来自适应。

不同的 stage 的 tokens 数量不同，越靠前的 stage 的 patchs 数量越多，我们知道 self-attention 的计算量与 sequence 的长度的平方成正比，如果 PVT 和 ViT 一样，所有的 transformer encoders 均采用相同的参数，那么计算量肯定是无法承受的。PVT 为了减少计算量，不同的 stages 采用的网络参数是不同的。PVT 不同系列的网络参数设置如下所示，这里为 patch_size，为特征维度大小，为 MHA（multi-head attention）的 heads 数量，为 FFN 的扩展系数，transformer 中默认为。

可以见到随着 stage，特征的维度是逐渐增加的，比如 stage1 的特征维度只有，而 stage4 的特征维度为，这种设置和常规的 CNN 网络设置是类似的，所以前面 stage 的 patchs 数量虽然大，但是特征维度小，所以计算量也不是太大。不同体量的 PVT 其差异主要体现在各个 stage 的 transformer encoder 的数量差异。

SRA

PVT 为了进一步减少计算量，将常规的 multi-head attention (MHA) 用 spatial-reduction attention (SRA) 来替换。SRA 的核心是减少 attention 层的 key 和 value 对的数量，常规的 MHA 在 attention 层计算时 key 和 value 对的数量为 sequence 的长度，但是 SRA 将其降低为原来的。

SRA 的具体结构如下所示：

在实现上，首先将维度为的 patch embeddings 通过 reshape 变换到维度为的 3-D 特征图，然后均分大小为的 patchs，每个 patchs 通过线性变换将得到维度为的 patch embeddings（这里实现上其实和 patch emb 操作类似，等价于一个卷积操作），最后应用一个 layer norm 层，这样就可以大大降低 K 和 V 的数量。具体实现代码如下：

从 PVT 的网络设置上，前面的 stage 的取较大的值，比如 stage1 的，说明这里直接将和的数量直接减为原来的，这个就大大降低计算量了。

PVT 具体到图像分类任务上，和 ViT 一样也通过引入一个 class token 来实现最后的分类，不过 PVT 是在最后的一个 stage 才引入：

Experience

具体到分类任务上，PVT 在 ImageNet 上的 Top-1 Acc 其实是和 ViT 差不多的。其实 PVT 最重要的应用是作为 dense 任务如分割和检测的 backbone，一方面 PVT 通过一些巧妙的设计使得对于分辨率较大的输入图像，其模型计算量不像 ViT 那么大，论文中比较了 ViT-Small/16 ，ViT-Small，PVT-Small 和 ResNet50 四种网络在不同的输入 scale 下的 GFLOPs，可以看到 PVT 相比 ViT 要好不少，当输入 scale=640 时，PVT-Small 和 ResNet50 的计算量是类似的，但是如果到更大的 scale，PVT 的增长速度就远超过 ResNet50 了。

PVT 的另外一个相比 ViT 的优势就是其可以输出不同 scale 的特征图，这对于分割和检测都是非常重要的。因为目前大部分的分割和检测模型都是采用 FPN 结构，而 PVT 这个特性可以使其作为替代 CNN 的 backbone 而无缝对接分割和检测的 heads。论文中做了大量的关于检测，语义分割以及实例分割的实验，可以看到 PVT 在 dense 任务的优势。比如，在更少的推理时间内，基于 PVT-Small 的 RetinaNet 比基于 R50 的 RetinaNet 在 COCO 上的 AP 值更高（38.7 vs. 36.3），虽然继续增加 scale 可以提升效果，但是就需要额外的推理时间：

PVTv2

PVTv2 通过添加三种设计来改进原始 Pyramid Vision Transformer（缩写为 PVTv1），包括 (1) 重叠块嵌入、(2) 卷积前馈网络和 (3) 线性复杂度注意层。通过这些修改，我们的 PVTv2 在分类、检测和分割等三个任务上显着改进了 PVTv1。此外，PVTv2 实现了与 Swin Transformer 等近期作品相当或更好的性能。

Method

Limitations in PVTv1

PVTv1的主要局限性有以下三个方面:

与 ViT 类似，在处理高分辨率输入时(如短边为800像素)，PVTv1的计算复杂度相对较大。

PVTv1 将一幅图像视为一组不重叠的patch序列，在一定程度上丧失了图像的局部连续性;

PVTv1中的位置编码是固定大小的，对于处理任意大小的图像是不灵活的。这些问题限制了PVTv1在视觉任务中的性能。

Linear Spatial Reduction Attention

用LinearSRA替代SRA。这里需要说明的一个问题，作者在PVTv1中说自己没用到卷积，但是在压缩K、V的时候使用的是Conv2D。在PVTv2中使用平均池化替代Conv2D。

Overlapping Patch Embedding

其次，为了对局部连续性信息进行建模，我们利用 Overlapping Patch Embedding 的方式对图像进行标记。如上图所示，我们放大 Patch 窗口，使相邻窗口重叠一半的面积，并用零填充特征图以保持分辨率。在这项工作中，我们使用带有零填充的卷积来实现 Overlapping Patch Embedding。具体来说，给定一个大小为的输入，我们将其送入一个跨度为、核大小为、填充大小为、核数为的卷积。输出的大小为。

Convolutional FeedForward

在前馈网络中删除固定大小位置编码,并引入零填充位置编码。我们添加一个padding为1的的深度卷积。放置在第一个(FC)层和GELU之间。