【深度学习】（ICCV-2021）PVT-金字塔 Vision Transformer及PVT_V2

0. 详情

名称：Pyramid Vision Transformer: A Versatile Backbone for Dense Prediction without Convolutions
论文： 原文
代码： 官方代码

笔记参考：
1. 语义分割中的Transformer（第三篇）：PVT — 用于密集预测任务的金字塔 Vision Transformer
2 作者本人的解释！！-大白话Pyramid Vision Transformer
3 金字塔Transformer，更适合稠密预测任务的Transformer骨干架构
4. 对pvt的思考！
5 简洁版-PVT(Pyramid Vision Transformer)算法整理
6. 翻译版

1. 简述

• 之前的所总结的ViT backbone，本身并没有针对视觉中诸如分割、检测等密集预测型的任务的特定，设计合适结构。
  后续SERT等论文也只是简单的将VIT作为Encoder，将其提取到的单尺度特征通过一些简单的Decoder的处理，验证了transformer在语义分割任务上的效果。
  但是，我们知道，在语义分割任务上，多尺度的特征是非常重要的，因此在PVT中提出了一种能够提取多尺度特征的vision transformer backbone。
  

2.主要工作

2.1 ViT遗留的问题

ViT的结构是下面这样的，它和原版Transformer一样是柱状结构的 。这就意味着，
1）缺乏多尺度特征
ViT 输出的特征图和输入大小基本保持一致。这导致ViT作为Encoder时，只能输出单尺度的特征。
它全程只能输出16-stride或者32-stride的feature map；
2）计算开销剧增
一旦输入图像的分辨率稍微大点，占用显存就会很高甚至显存溢出。
分割和检测相对于分类任务而言，往往需要较大的分辨率图片输入。
因此，一方面，我们需要相对于分类任务而言划分更多个patch才能得到相同粒度的特征。如果仍然保持同样的patch数量，那么特征的粒度将会变粗，从而导致性能下降
另一方面，我们知道，Transformer的计算开销与token化后的patch数量正相关， patch数量越大，计算开销越大。所以，如果我们增大patch数量，可能就会让我们本就不富裕的计算资源雪上加霜。
以上是ViT应用于密集预测任务上的第一个缺陷。

解决
针对方案简单粗暴：

• 输出分辨率不够，那么就加大；
• patch token序列太长，导致attention矩阵的计算量太大，那么就针对性的缩减总体序列长度，或者是仅仅缩减k和v的长度(如下图)。

2.2 引入金字塔结构

计算机视觉中CNN backbone经过多年的发展，沉淀了一些通用的设计模式。
最为典型的就是金字塔结构。
简单的概括就是：

1）feature map的分辨率随着网络加深，逐渐减小；

2）feature map的channel数随着网络加深，逐渐增大。
几乎所有的密集预测（dense prediction）算法都是围绕着特征金字塔设计的

这个结构怎么才能引入到Transformer里面呢？
最终还是发现：简单地堆叠多个独立的Transformer encoder效果是最好的。
然后我们就得到了PVT，如下图所示。在每个Stage中通过Patch Embedding来逐渐降低输入的分辨率。

其中，除了金字塔结构以外。为了可以以更小的代价处理高分辨率（4-stride或8-stride）的feature map，我们对Multi-Head Attention也做了一些调整。

为了在保证feature map分辨率和全局感受野的同时降低计算量，我们把key（K）和value（V）的长和宽分别缩小到以前的1/R_i。通过这种方法，我们就可以以一个较小的代价处理4-stride，和8-stride的feature map了。

3.PVT的设计方案

模型总体上由4个stage组成用于生成不同尺度的特征，每个stage包含Patch Embedding和若干个Transformer模块（相对于原本的transformer有所改动）组成。

• Patch Embedding：目的在于将信息分块，降低单张图的图片大小，但会增加数据的深度
• Transformer Encoder：目的在于计算图片的attention value，由于深度变大了，计算复杂度会变大，所以在这里作者使用了Special Reduction来减小计算复杂度

在第一个阶段，给定尺寸为 H X WX3 的输入图像，我们按照如下流程进行处理：

首先，将其划分为HW/4的平方的块(这里是为了与ResNet对标，最大输出特征的尺寸为原始分辨率的1/4)，每个块的大小为；
然后，将展开后的块送入到线性投影曾得到尺寸为HW/4的平方 xC1 的嵌入块；
其次，将前述嵌入块与位置嵌入信息送入到Transformer的Encoder，其输出将为reshap为.H/4 X W/4 X C1

采用类似的方式，我们以前一阶段的输出作为输入即可得到特征F2 F3 F4。

H * W * 3 -> stage1 block -> H/4 * W/4 * C1 -> stage2 block -> H/8 * W/8 * C2 -> stage3 block -> H/16 * W/16 * C3 -> stage3 block -> H/32 * W/32 * C4

3.1 Patch embedding

Patch Embedding部分与ViT中对与图片的分块操作是一样的，即：

通过这种方式我们就可以灵活的调整每个阶段的特征尺寸，使其可以针对Transformer构建特征金字塔。

在每个stage开始，首先像ViT一样对输入图像进行token化，即进行patch embedding，patch大小除第1个stage的是4x4外，其余均采用2x2大小。这个思想有些类似于池化或者带步长的卷积操作，减小图像的分辨率，使得模型能够提取到更为抽象的信息。这意味着每个stage（第一个stage除外）最终得到的特征图维度是减半的，tokens数量对应减少4倍。

每个patch随后会送入一层Linear中，调整通道数量，然后再reshape以将patch token化。

这使得PVT总体上与resnet看起来类似，4个stage得到的特征图相比原图大小分别是1/4，1/8，1/16 和 1/32。这也意味着PVT可以产生不同尺度的特征。

Note：由于不同的stage的tokens数量不一样，所以每个stage采用不同的position embeddings，在patch embed之后加上各自的position embedding，当输入图像大小变化时，position embeddings也可以通过插值来自适应。

代码

1、首先输入的data的shape是（bs,channal,H,W)，为了方便直接用batchsize是1的图片做例子，因此输入是(1，3，224，224)
code对应：

model = pvt_small(**cfg)
data = torch.randn((1, 3, 224, 224))
output = model(data)

2、输入数据首先经过stage 1 block的Patch emb操作，这个操作首先把224224的图像分成44的一个个小patch，这个实现是用卷积实现的，用44的卷积和对224224的图像进行卷积，步长为4即可。
code对应：

self.proj = nn.Conv2d(in_chans=3, embed_dim=64, kernel_size=4, stride=4)
# 其中64是网络第一层输出特征的维度对应图中的C1
print(x.shape)# torch.Size([1, 3, 224, 224])
x = self.proj(x)
print(x.shape)# torch.Size([1, 64, 56, 56])

这样就可以用5656的矩阵每一个点表示原来44的patch

3、对16456*56的矩阵在进行第二个维度展平
code对应：

print(x.shape) # torch.Size([1, 64, 56, 56])
x = x.flatten(2)
print(x.shape) # torch.Size([1, 64, 3136])

这时候就可以用3136这个一维的向量来表示224*224的图像了

4、为了方便计算调换下第二第三两个维度，然后对数据进行layer norm。
code对应：

print(x.shape) # torch.Size([1, 64, 3136])
x = x.transpose(1, 2)
print(x.shape) # torch.Size([1, 3136, 64])
x = self.norm(x)
print(x.shape) # torch.Size([1, 3136, 64])

以上就完成了Patch emb的操作，完整代码对应：

def forward(self, x):
    B, C, H, W = x.shape # 1，3，224，224
    x = self.proj(x) # 卷积操作，输出1，64，56，56
    x = x.flatten(2) # 展平操作，输出1，64，3136
    x = x.transpose(1, 2) # 交换维度，输出 1，3136，64
    x = self.norm(x) # layer normal，输出 1，3136，64
    H, W = H // 4, W // 4 # 最终的高宽变成56，56
    return x, (H, W)

3.2position embedding

1、这部分和Vit的位置编码基本是一样的，创建一个可学习的参数，大小和patch emb出来的tensor的大小一致就是(1,3136,64)，这是个可学习的参数。
code对应：

pos_embed = nn.Parameter(torch.zeros(1, 3136, 64))

2、位置编码的使用也是和Vit一样，**直接和输出的x进行矩阵加，**因此shape不变化。
code对应：

print(x.shape) # torch.Size([1, 3136, 64])
x = x + pos_embed
print(x.shape) # torch.Size([1, 3136, 64])

3、相加完后，作者加了个dropout进行正则化。
code对应：

pos_drop = nn.Dropout(p=drop_rate)
x = pos_drop(x)

以上就完成了position embedding的操作，完整代码对应：

x = x + pos_embed
x = pos_drop(x)

3.3Encoder

第i个stage的encoder部分由depth[i]个block构成，对于pvt_tiny到pvt_large来说主要就是depth的参数的不同：

例如对于pvt_tiny来说，每个encoder都是由两个block构成，每个block的结构如下图所示：

对于第一个encoder的第一个block的输入就是我们前面分析的经过position embedding后拿到的tensor，因此他的输入的大小是(1,3136,64)，与此同时图像经过Patch emb后变成了56*56的大小。

1、首先从上图可以看出先对输入拷贝一份，给残差结构用。然后输入的x先经过一层layer norm层，此时维度不变，然后经过作者修改的Multi head attention层（SRA，后面再讲）与之前拷贝的输入叠加。
code对应：

print(x.shape) # (1,3136,64)
x = x + self.drop_path(self.attn(self.norm1(x), H, W))
print(x.shape) # (1,3136,64)

2、经过SRA层的特征拷贝一份留给残差结构，然后将输入经过layer norm层，维度不变，再送入feed forward层（后面再讲），之后与之前拷贝的输入叠加。
code对应：

print(x.shape) # (1,3136,64)
x = x + self.drop_path(self.mlp(self.norm2(x)))
print(x.shape) # (1,3136,64)

因此可以发现经过一个block，tensor的shape是不发生变化的。完整的代码对应：

def forward(self, x, H, W):
      x = x + self.drop_path(self.attn(self.norm1(x), H, W)) # SRA层
      x = x + self.drop_path(self.mlp(self.norm2(x))) # feed forward层
      return x

3、这样经过depth[i]个block之后拿到的tensor的大小仍然是(1,3136,64)，只需要将它的shape还原成图像的形状就可以输入给下一个stage了。而还原shape，直接调用reshape函数即可，这时候的特征就还原成(bs,channal,H,W)了，数值为(1,64,56,56)
code对应：

print(x.shape) # 1,3136,64
x = x.reshape(B, H, W, -1)
print(x.shape) # 1,56,56,64
x = x.permute(0, 3, 1, 2).contiguous()
print(x.shape) # 1,64,56,56

这时候stage2输入的tensor就是(1,64,56,56)，就完成了数据输出第一个stage的完整分析。
最后只要在不同的encoder中堆叠不同个数的block就可以构建出pvt_tiny、pvt_small、pvt_medium、pvt_large了。
完整图示如下：

所以！经过stage1，输入为（1，3，224，224）的tensor变成了（1，64，56，56）的tensor，这个tensor可以再次输入给下一个stage重复上述的计算就完成了PVT的设计。

3.2 Spatial-reduction attention（SRA）

在Patch embedding之后，需要将token化后的patch输入到若干个transformer 模块中进行处理。

不同的stage的tokens数量不同，越靠前的stage的patchs数量越多，self-attention的计算量与sequence的长度N的平方成正比，如果PVT和ViT一样，所有的transformer encoders均采用相同的参数，那么计算量肯定是无法承受的。

1.PVT 为了减少计算量，不同的stages采用的网络参数是不同的。
PVT不同系列的网络参数设置如下所示，这里P为patch的size，C为特征维度大小，N为MHA（multi-head attention）的heads数量，E为FFN的扩展系数，transformer中默认为4。

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

2.为了进一步减少计算量，作者将multi-head attention (MHA)用所提出的spatial-reduction attention (SRA)来替换。

SRA的核心是减少attention层的key和value对的数量，常规的MHA在attention层计算时key和value对的数量为sequence的长度，但是SRA将其降低为原来的1/Rd 平方。

SRA的处理过程可以描述如下：

也就是说每个head的维度等于Ci/Ni，SR( )为空间尺度下采样操作，定义如下：


(
总而言之，所提方案涉及的超参数包含如下：

P：阶段i的块尺寸；
C：阶段i的通道数；
L：阶段i的encoder层数；
R：阶段i中SRA的下采样比例；
N：阶段i的head数量；
E：阶段i中MLP的扩展比例。
)
具体来说：
在实现上，

• 首先将维度为（HW，C）的K，V通过 reshape变换到维度为（H，W， C）的3-D特征图，
• 然后均分大小为R * R的patchs，每个patchs通过线性变换将得到维度为（HW / RR，C）的patch embeddings（这里实现上其实和patch emb操作类似，等价于一个卷积操作），
• 最后应用一个layer norm层，这样就可以大大降低K和V的数量。其核心代码也是这么实现的：

self.sr = nn.Conv2d(dim, dim, kernel_size=sr_ratio, stride=sr_ratio)
self.norm = nn.LayerNorm(dim)```

每个stage，经过若干个SRA模块的处理后，将得到的特征，再次reshape成3D特征图的形式输入到下一个Stage中。

1、首先如果参数self.sr_ratio为1的话，那么pvt的attetion就和vit的attetion一模一样了：

2、因此分析不一样的地方

self.sr = nn.Conv2d(dim, dim, kernel_size=sr_ratio, stride=sr_ratio)
x_ = x.permute(0, 2, 1).reshape(B, C, H, W)
x_ = self.sr(x_).reshape(B, C, -1).permute(0, 2, 1)
x_ = self.norm(x_)
kv = self.kv(x_).reshape(B, -1, 2, self.num_heads, C // self.num_heads).permute(2, 0, 3, 1, 4)

2.1：首先输入进来的x的shape是(1，3136，64)，
2.2：先permute置换维度得到(1，64，3136)，
2.3：reshape得到(1，64，56，56)
2.4：self.sr(x_)是一个卷积操作，卷积的步长和大小都是sr_ratio，这个数值这里是8因此相当于将56*56的大小长宽缩小到8分之一，也就是面积缩小到64分之一，因此输出的shape是(1,64,7,7)
2.5：reshape(B, C, -1)得到(1,64,49)
2.6：permute(0, 2, 1)得到(1,49,64)
2.7：经过layer norm后shape不变
2.8：kv = self.kv(x_).reshape(B, -1, 2, self.num_heads, C // self.num_heads).permute(2, 0, 3, 1, 4)这一行就是和vit一样用x生成k和v，不同的是这里的x通过卷积的方式降低了x的大小。
这一行shape的变化是这样的:(1,49,64)->(1,49,128)->(1,