金字塔Transformer，更适合稠密预测任务的Transformer骨干架构

首发：AIWalker
作者:HappyAIWalker

本文是南京大学&港中文&南理工&商汤在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。

image-20210225163938302

Method

image-20210225164117660

本文旨在将金字塔结构嵌入到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)。

image-20210225170248387

类似于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