Pyramid Vision Transformer, PVT（ICCV 2021）原理与代码解读

勿忘初心做自己 · 2024-6-21 16:11:16

paper：Pyramid Vision Transformer: A Versatile Backbone for Dense Prediction without Convolutions
official implementation：GitHub - whai362/PVT: Official implementation of PVT series
存在的问题

现有的 Vision Transformer (ViT) 主要设计用于图像分类任务，难以直接用于像素级密集猜测任务，如目标检测和分割。这是由于存在以下问题

低分辨率输出：传统的Vision Transformer（ViT）在处理密集猜测任务（如目标检测和语义分割）时，输出分辨率较低，难以获得高质量的像素级别猜测。
高计算和内存开销：ViT在处理大尺寸输入图像时，计算和内存开销较高，限定了其在实际应用中的效率。

本文的创新点

为了办理上述问题，作者提出了 Pyramid Vision Transformer (PVT)， PVT结合了卷积神经网络的金字塔结构和Transformer的全局感受野，旨在克服传统Transformer在处理密集猜测任务时遇到的分辨率低、计算和内存开销大的问题。它可以作为 CNN 骨干网络的替换品，用于多种下游任务，包括图像级猜测和像素级密集猜测。具体包括：

金字塔结构：PVT引入了金字塔结构，可以生成多尺度的特性图，这对于密集猜测任务是有益的。
空间缩减注意力层（SRA）：为了处理高分辨率特性图并减少计算/内存本钱，作者设计了 SRA 层来替换传统的多头注意力 (MHA) 层。
纯Transformer骨干：PVT 是一个没有卷积的纯 Transformer 骨干网络，可以用于各种像素级密集猜测任务，并与 DETR 结合构建了一个完全无需卷积的目标检测系统。

实际结果

PVT 在多个下游任务上进行了广泛的实验验证，包括图像分类、目标检测、实例和语义分割等，并与流行的 ResNets 和 ResNeXts 进行了比较。
实验结果表明，在参数数目相称的情况下，PVT 在 COCO 数据集上使用 RetinaNet 作为检测器时，PVT-Small 模子达到了 40.4 的 AP（平均精度），超过了 ResNet50+RetinaNet（36.3 AP）4.1 个百分点。
PVT-Large 模子达到了 42.6 的 AP，比 ResNeXt101-64x4d 高出 1.6 个百分点，同时参数数目减少了 30%。
这些结果表明 PVT 可以作为 CNN 骨干网络的一个有效的替换，用于像素级猜测，并推动将来的研究。

方法先容

Overall Architecture

PVT的整体结构如图3所示

和CNN backbone雷同，PVT也有四个stage来生成差别尺度的特性图。全部stage都有一个相似的架构，包括一个patch embedding层和 \(L_i\) 个Transformer encoder层。
在第一个stage，给定大小为 \(H\times W\times 3\) 的输入图片，我们起首将其划分为 \(\frac{HW}{4^2}\) 个patch，每个大小为4x4x3。然后将展平的patch送入一个线性映射层得到大小为 \(\frac{HW}{4^2}\times C_1\) 的输出。然后将输出和位置编码一起进入有 \(L_1\) 层的Transformer encoder，得到的输出reshape成大小为 \(\frac{H}{4}\times \frac{W}{4}\times C_1\) 的特性图 \(F_1\)。同样的方式，以前一个stage的输出特性图作为输入，我们得到特性图 \(F_2,F_3,F_4\)，相对于原始输入图片的步长分别为8，16,32。用了特性图金字塔 \(\{F_1,F_2,F_3,F_4\}\)，我们的方法可以很容易地应用于大多数下游任务，包括图像分类、目标检测和语义分割。
Feature Pyramid for Transformer

和CNN backbone用差别stride的卷积来得到差别尺度特性图差别，PVT使用一个渐进式shrinking策略，通过patch embedding层来控制特性图的尺度。
我们用 \(P_i\) 来表现第 \(i\) 个stage的patch size，在stage \(i\) 的开始，我们起首将输入特性图 \(F_{i-1}\in \mathbb{R}^{H_{i-1}\times W_{i-1}\times C_{i-1}}\) 匀称地划分成 \(\frac{H_{i-1}W_{i-1}}{P_i^2}\) 个patch，然后将每个patch展平并映射得到一个 \(C_i\) 维的embedding。在线性映射后，embedded patch的大小为 \(\frac{H_{i-1}}{P_i}\times \frac{W_{i-1}}{P_i}\times C_i\)，其中宽高比输入小了 \(P_i\) 倍。
如许，我们就可以在每个stage灵活地调整特性图的尺度，从而将Transformer构建成金字塔结构。
Transforme Encoder

由于PVT必要处理高分辨率（stride-4）的特性图，我们提出了一种spatial-reduction attention（SRA）来替换encoder中传统的multi-head attention（MHA）。
和MHA雷同，SRA的输入包括一个query \(Q\)，一个key \(K\)，一个value \(V\)。差别的是SRA在attention operation之前减小了 \(K\) 和 \(V\) 的大小，如图4所示，这大大减少了计算和内存的开销。

stage \(i\) 的SRA如下

其中 \(Concat(\cdot)\) 是拼接操作。\(W^{Q}_j\in \mathbb{R}^{C_i\times d_{head}},W^{K}_j\in \mathbb{R}^{C_i\times d_{head}},W^{V}_j\in \mathbb{R}^{C_i\times d_{head}},W^O\in \mathbb{R}^{C_i\times C_i}\) 是线性映射参数。\(N_i\) 是stage \(i\) 中attention层的head数目，所以每个head的维度（即\(d_{head}\)）即是 \(\frac{C_i}{N_i}\)。\(SR(\cdot)\) 是降低输入序列（即 \(K\) 或 \(V\)）空间维度的操作，如下：

其中 \(\mathbf{x}\in\mathbb{R}^{(H_iW_i)\times C_i}\) 表现一个输入序列，\(R_i\) 表现stage \(i\) 中attention层的reduction ratio。\(Reshape(\mathbf{x},R_i)\) 是将输入序列 \(\mathbf{x}\) reshape成大小为 \(\frac{H_iW_i}{R^2_i}\times (R^2_iC_i)\) 的序列的操作。\(W_S\in \mathbb{R}^{(R^2_iC_i)\times C_i}\) 是一个linear projection，它将输入序列的维度降低到 \(C_i\)。\(Norm(\cdot)\) 是layer normalization。和原始的Transformer一样，attention operation按下式计算

通过上述公式我们可以发现，MSA的计算/内存开销是MHA的 \(\frac{1}{R^2}\)，因此MSA可以在有限的资源下处理更大的输入特性图或序列。
代码解析

见PVT v2的代码解析 PVT v2 原理与代码解析-CSDN博客
实验结果

模子涉及到的一些超参总结如下：