金字塔视觉transformer PVT及其改进

PVT金字塔视觉transformer

(2021.2.24) 南京大学, 港大, 南理, IIAI, 商汤

PVT提出第一个用于不同像素级密集预测任务的视觉Transformer网络。

参考论文作者解读：https://zhuanlan.zhihu.com/p/353222035
论文地址：https://arxiv.org/abs/2102.12122
源码：https://github.com/whai362/PVT

研究问题

​ViT将输入图像tokens化后不断堆积相同的transformer encoders，但是如果应用在密集任务上，会遇到问题：一是分割和检测往往需要较大的分辨率输入，当输入图像增大时，ViT的计算量会急剧上升；二是ViT直接采用较大patchs进行token化，如采用16x16大小那么得到的粗粒度特征，对密集任务来说损失较大。

方法

1. 输入为H x W x 3的图像，划分为4x4大小的patch，就获得了$\frac{HW}{4^2}$ 个patch，每个patch大小为4×4×3，将展平的patch送入线性变换（卷积操作，卷积核和步长均为patch大小)，得到$\frac{HW}{4^2}\ast C_1$作为序列输入到transformer中。

2. 每个stage的输入都是一个维度$H_i\ast W_i\ast C_i$的3-D特征图，对于其它stage可以将tokens重新reshape成3-D特征图。

3. PVT共4个stage，这和ResNet类似，4个stage得到的特征图相比原图大小分别是1/4，1/8，1/16和1/32。由于不同的stage的tokens数量不一样，所以每个stage采用不同的position embeddings，在patch embed之后加上各自的position embedding，当输入图像大小变化时，position embeddings也可以通过插值来自适应。

PVT为了进一步减少计算量，将常规的multi-head attention (MHA)用spatial-reduction attention (SRA)来替换。SRA的核心是减少attention层的key和value对的数量，常规的MHA在attention层计算时key和value对的数量为sequence的长度，但是SRA将其降低为原来的$\frac{1}{R^2}$
SRA的具体结构如下所示：

PVT-v2

问题

​与 ViT 相同， PVTv1 将图像视为一系列非重叠的片元序列，这样对图像的非重叠切片与编码会在一定程度上损失图像中原有的连续性；其次，PVTv1 中的位置编码为固定尺寸，对于任意尺度的图像处理缺乏灵活性。当高分辨率的图像输入时，PVTv1 的计算复杂度就会飙升，大大降低了模型的性能。

改进

• 重叠片元 (patch) 嵌入编码。在 PVTv2 中使用了重叠片元嵌入来对图像进行编码。上图左侧展示了重叠片元操作的示意图，图中每个片元的窗口尺寸被放大、与相邻窗口互相重叠一半，同时对特征图进行0填充操作。而后使用卷积对填充后特征图进操作实现嵌入编码。具体来讲，给定 hxwxc 的输入，应用步长为S、尺寸为2S-1的卷积，填充大小为S-1，使用c’个卷积核最终获得 h/S x w/S x C’的输出结果。
  

• 卷积前传。为了解决图像大小灵活性的问题，在新版的 PVT 中固定位置编码被移除，并引入了填充零的位置编码机制，上图右侧显示了在前传网络和全连接层间插入的3x3零填充逐深度卷积。

• 线性空间缩减注意力机制。为了进一步优化PVT的计算开销，线性空间注意力 (Spatial Reduction Attention, SRA) 机制被引入到新的模型中来。与原始的SRA不同，线性空间注意力机制具有线性的计算复杂度和内存开销，对于输入为hxwxc的特征图，与SRA相比线性SRA的复杂度大大降低。