论文速递：金字塔Transformer，更适合稠密预测任务的Transformer骨干架构

本文是南京大学&港中文&南理工&商汤在Transformer方面的又一次探索，该文主要针对稠密预测任务的骨干的Transformer架构设计进行了一次创新性的探索，将特征金字塔结构与Transformer进行了一次融合，使其可以更好的输出多尺度特征，进而更方便与其他下游任务相结合。本文为Transformer在计算机视觉领域的应用提供了一个非常好的开头，期待有更优秀的Transformer在CV领域的应用于探索。

Abstract

以CNN为骨干的方案在计算机视觉的各个领域均取得极大的成功，本文研究了一种简单的无卷积骨干网络用于诸多稠密预测任务(包含检测、分割等)。近来提出的Transformer(ViT)主要是针对图像分类任务而设计，我们提出稠密预测任务提出了一种PVT方案(Pyramid Vision Transformer)，它克服了将ViT向稠密预测任务迁移的困难。

(1) 不同于ViT的低分辨率输出、高计算复杂度与内存消耗，PVT不仅可以得到更高分辨率的输出(这对于稠密预测任务尤为重要)，同时按照金字塔形式渐进式收缩特征分辨率；

(2) PVT继承了CNN与Transformer的优势，通过简单的替换CNN骨干使其成为不同视觉任务的统一骨干结构；

(3) 通过充分的实验验证了PVT的优越性，在目标检测、语义分割、实例分割等任务上PVT均取得了优异性能。比如RetinaNet+PVT在COCO数据集上取得了40.4AP指标，超越了RetinaNet+ResNet50的36.3AP。

Method

本文旨在将金字塔结构嵌入到Transformer结构用于生成多尺度特征，并最终用于稠密预测任务。上图给出了本文所提PVT的架构示意图，类似与CNN骨干结构，PVT同样包含四个阶段用于生成不同尺度的特征，所有阶段具有相类似的结构：Patch Embedding+Transformer Encoder。

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

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

采用类似的方式，我们以前一阶段的输出作为输入即可得到特征。基于特征金字塔，所提方案可以轻易与大部分下游任务(如图像分类、目标检测、语义分割)进行集成。

Feature Pyramid for Transformer

不同于CNN采用stride卷积获得多尺度特征，PVT通过块嵌入曾按照progressive shrinking策略控制特征的尺度。

我们假设第i阶段的块尺寸为，在每个阶段的开始，我们将输入特征均匀的拆分为个块，然后每个块展开并投影到维度的嵌入信息，经过线性投影后，嵌入块的尺寸可以视作。通过这种方式我们就可以灵活的调整每个阶段的特征尺寸，使其可以针对Transformer构建特征金字塔。该过程的code描述如下：

class PatchEmbed(nn.Module):    """ Image to Patch Embedding    """
    def __init__(self, img_size=224, patch_size=16, in_chans=3, embed_dim=768):        super().__init__()        img_size = to_2tuple(img_size)        patch_size = to_2tuple(patch_size)
        self.img_size = img_size        self.patch_size = patch_size        assert img_size[0] % patch_size[0] == 0 and img_size[1] % patch_size[1] == 0, \            f"img_size {img_size} should be divided by patch_size {patch_size}."        self.H, self.W = img_size[0] // patch_size[0], img_size[1] // patch_size[1]        self.num_patches = self.H * self.W        self.proj = nn.Conv2d(in_chans, embed_dim, kernel_size=patch_size, stride=patch_size)        self.norm = nn.LayerNorm(embed_dim)
    def forward(self, x):        B, C, H, W = x.shape
        x = self.proj(x).flatten(2).transpose(1, 2)        x = self.norm(x)        H, W = H // self.patch_size[0], W // self.patch_size[1]
        return x, (H, W)

从实现上来看，其实这里还是通过stride=2的卷积进行的实现，汗一个。

Transformer Encoder

对于Transformer encoder的第i阶段，它具有个encoder层，每个encoder层由注意力层与MLP构成。由于所提方法需要处理高分辨率特征，所以我们提出了一种SRA(spatial-reduction attention)用于替换传统的MHA(multi-head-attention)。

类似于MHA，SRA同样从输入端接收到了Q、K、V作为输入，并输出精炼后特征。SRA与MHA的区别在于：SRA会降低K和V的空间尺度，见上图。SRA的处理过程可以描述如下：

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

其他的就跟MHA没什么区别了，这里提供一下SRA的code实现，关键区别就是sr_ratio相关的那一部分。

class Attention(nn.Module):    def __init__(self, dim, num_heads=8, qkv_bias=False, qk_scale=None, attn_drop=0., proj_drop=0., sr_ratio=1):        super().__init__()        assert dim % num_heads == 0, f"dim {dim} should be divided by num_heads {num_heads}."
        self.dim = dim        self.num_heads = num_heads        head_dim = dim // num_heads        self.scale = qk_scale or head_dim ** -0.5
        self.q = nn.Linear(dim, dim, bias=qkv_bias)        self.kv = nn.Linear(dim, dim * 2, bias=qkv_bias)        self.attn_drop = nn.Dropout(attn_drop)        self.proj = nn.Linear(dim, dim)        self.proj_drop = nn.Dropout(proj_drop)
        self.sr_ratio = sr_ratio        if sr_ratio > 1:            self.sr = nn.Conv2d(dim, dim, kernel_size=sr_ratio, stride=sr_ratio)            self.norm = nn.LayerNorm(dim)
    def forward(self, x, H, W):        B, N, C = x.shape        q = self.q(x).reshape(B, N, self.num_heads, C // self.num_heads).permute(0, 2, 1, 3)
        if self.sr_ratio > 1:            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)        else:            kv = self.kv(x).reshape(B, -1, 2, self.num_heads, C // self.num_heads).permute(2, 0, 3, 1, 4)        k, v = kv[0], kv[1]
        attn = (q @ k.transpose(-2, -1)) * self.scale        attn = attn.softmax(dim=-1)        attn = self.attn_drop(attn)
        x = (attn @ v).transpose(1, 2).reshape(B, N, C)        x = self.proj(x)        x = self.proj_drop(x)
        return x
Model Details

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

• ：阶段i的块尺寸；
• ：阶段i的通道数；
• ：阶段i的encoder层数；
• ：阶段i中SRA的下采样比例；
• ：阶段i的head数量；
• ：阶段i中MLP的扩展比例。

延续ResNet的设计准则，我们在浅层采用较小的输出通道数；在中间阶段聚焦主要的计算量。为更好的讨论与分析，我们设计了不同大小的PVT模型，定义见下表。

Experiments

为说明所提方案的有效性，我们在图像分类、目标检测、语义分割、实例分割等任务上进行了对比分析。

首先，我们来看一下ImageNet数据集上的性能对比，结果见上表。从中可以看到：

• 相比CNN，在同等参数量与计算约束下，PVT-Small达到20.2%的误差率，优于ResNet50的21.5%；
• 相比其他Transformer(如ViT、DeiT)，所提PVT以更少的计算量取得了相当的性能。

然后，我们再来看一下目标检测方面的性能对比，结果见上表。从中可以看到：在参数量相当的情况下，PVT显著性的优于CNN方案，这说明PVT是一个CNN之外的一个好的选择。类似的现象在Mask RCNN体系下仍可发现，见下表。

其次，我们看一下所提方案在语义分割中的性能对比，见上表。可以看到：不同参数配置下，PVT均可取得优于ResNet与ResNeXt的性能。这侧面说明：相比CNN，受益于全局注意力机制，PVT可以提取更好的特征用于语义分割。

全文到此结束，更多消融实验与分析建议各位同学查看原文。