【Computer Vision论文笔记】PVT && PVTv2

最新推荐文章于 2024-06-07 19:38:02 发布

ZouJiaYuY

最新推荐文章于 2024-06-07 19:38:02 发布

阅读量440

点赞数

文章标签：人工智能计算机视觉论文阅读

本文链接：https://blog.csdn.net/zjy_cn/article/details/133579129

版权

本文介绍了PyramidVisionTransformer(PVT)和其改进版PVTv2，这是一种基于Transformer的金字塔结构网络，用于密集预测任务。PVT通过SpatialReductionAttention减少计算量，而PVTv2引入了LinearSpatialReductionAttention和改进的嵌入方法。尽管实验显示在某些任务上优于CNN，但PVTv2的创新点并非显著提升精度。

摘要由CSDN通过智能技术生成

论文1：Pyramid Vision Transformer: A Versatile Backbone for Dense Prediction without Convolutions

论文2：PVT v2: Improved Baselines with Pyramid Vision Transformer

代码：https://github.com/whai362/PVT

前言

Pyramid Vision Transformer(PVT) 是致力于密集预测的主干网络的研究，是一个由Transformer组成的金字塔结构的网络，能够在多种视觉任务中使用。金字塔结构是CNN中一种非常经典的设计，作者在此将其应用到Transformer中，如下图：

图（a）所示的 CNN 结构，feature maps会随着网络深度的增加，大小逐渐缩小。图（b）是ViT的维度变化情况，输入与输出大小保持不变。结合两者，就有了图（c）所示的一Tansformer为基础的金字塔结构。

PVTv2 是在 PVT 的基础上做了些小改进。

模型介绍

PVT

首先看 PVT，它的整体结构如下：

采用的是四阶层设计，每个阶层包含 L 个 Transformer Encoder，feature maps的大小逐阶减小、通道数逐阶增加（progressive shrinking strategy，这是ResNet提出的设计结构）。

论文的主要亮点是 Transformer Encoder 中的 Spatial Reduction Attention (SRA)。

在标准 Transformer 中的 Multi-Head Attention(MHA)，它存在一个很大的弊端，就是参数量大，需要较高的算力支持，尤其是在面对高分辨率的图像，计算开销巨大。论文作者从注意力机制入手，在进行 attention 前减小了空间尺度的大小，并且取得良好的实验效果。

Spatial Reduction会在K、V矩阵进行进入 MHA 前缩小其大小，如下：

由此可见，SRA 的计算成本为 MHA 的 $\frac{1}{R_{i}^{2}}$ ，这样一来，就能更加高效进行训练。

SRA 的实现相当简单，它引入了一个缩减率参数 R，代码如下：

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 # 缩减率，sr_ratios=[8, 4, 2, 1]，对应不同阶层的缩减大小
        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

这种粗暴地缩小输入的空间尺度，相较于 MHA，还是会影响精度的，只能说是在精度和计算量之间进行取舍。

PVTv2

v2版本的 PVT 只是在几个小地方稍加修改，主要是三个方面。

第一个是 attention，提出了Linear Spatial Reduction Attention (LInear SRA)，它使用平均池化在注意力操作之前将空间维度（即：h×w）减少到固定大小（即：P×P），因此，Linear SRA 像卷积层一样具有线性计算和内存成本：

其计算量对比：

再看Linear SRA的代码：

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, linear=False):
        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.linear = linear
        self.sr_ratio = sr_ratio
        if not linear:
            if sr_ratio > 1:
                self.sr = nn.Conv2d(dim, dim, kernel_size=sr_ratio, stride=sr_ratio)
                self.norm = nn.LayerNorm(dim)
        else:
            self.pool = nn.AdaptiveAvgPool2d(7)
            self.sr = nn.Conv2d(dim, dim, kernel_size=1, stride=1)
            self.norm = nn.LayerNorm(dim)
            self.act = nn.GELU()
        self.apply(self._init_weights)

    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 not self.linear:
            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)
        else:
            x_ = x.permute(0, 2, 1).reshape(B, C, H, W)
            x_ = self.sr(self.pool(x_)).reshape(B, C, -1).permute(0, 2, 1)
            x_ = self.norm(x_)
            x_ = self.act(x_)
            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