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

AIWalker

发布于 2021-03-04 14:47:03

8850

发布于 2021-03-04 14:47:03

文章被收录于专栏：AIWalkerAIWalker

标题&作者团队

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

在第一个阶段，给定尺寸为

H \times W \times 3

的输入图像，我们按照如下流程进行处理：

首先，将其划分为

\frac{HW}{4^2}

的块(这里是为了与ResNet对标，最大输出特征的尺寸为原始分辨率的1/4)，每个块的大小为

4\times 4\times 3

；

然后，将展开后的块送入到线性投影曾得到尺寸为

\frac{HW}{4^2} \times C_1

的嵌入块；

其次，将前述嵌入块与位置嵌入信息送入到Transformer的Encoder，其输出将为reshap为

\frac{H}{4} \times \frac{W}{4} \times C_1

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

F_2, F_3, F_4

。基于特征金字塔

F_1, F_2, F_3, F_4

，所提方案可以轻易与大部分下游任务(如图像分类、目标检测、语义分割)进行集成。

Feature Pyramid for Transformer

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

我们假设第i阶段的块尺寸为

P_i

，在每个阶段的开始，我们将输入特征

F_{i-1} \in R^{H_{i-1}\times W_{i-1} \times C_{i-1}}

均匀的拆分为

\frac{H_{i-1}W_{i-1}}{P_i^2}

个块，然后每个块展开并投影到

C_i

维度的嵌入信息，经过线性投影后，嵌入块的尺寸可以视作

\frac{H_{i-1}}{P_i} \times \frac{W_{i-1}}{P_i} \times C_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阶段，它具有

L_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的处理过程可以描述如下：

SRA(Q,K,V) = Concat(head_0, \cdots, head_{N_i})W^O \\ head_j = Attention(QW_j^Q, SR(K)W_j^K, SR(V)W_j^V)

也就是说每个head的维度等于

\frac{C_i}{N_i}

，

SR(\cdot)

为空间尺度下采样操作，定义如下：

SR(x) = Norm(Reshape(x, R_i)W^S)

其他的就跟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