如何构建一个超快速语义分割Transformer方案？丢掉多头 ...

潮流服饰 · 发表于 2022-12-2 15:08:43

作者丨happy
编辑丨极市平台

论文链接：https://arxiv.org/abs/2210.07124 [NeurIPS2022]
代码地址：https://github.com/PaddlePaddle/PaddleSeg
尽管Transformer方案在语义分割领域取得了非常惊人的性能，但在实时性方面，纯CNN方案仍占据主流地位。本文提出了一种用于实时语义分割的高效对偶分辨率Transformer方案RTFormer，它具有比CNN方案更佳的性能-效率均衡。
为达成GPU设备上的高推理效率，所提RTFormer采用了线性复杂度的GPU友好注意力模块，同时消除了多头机制。此外，作者发现：跨注意力机制对于全局上下文信息聚合非常有效。多个主流基准数据集(Cityscapes, CamVid, COCOStuff, ADE20K)上的实验结果验证了所提RTFormer的有效性。下图给出了CAMVid数据集上不同方案的性能与推理速度对比，很明显：RTFormer具有最佳的性能-速度均衡。

出发点

ViT技术在CV领域证实其有效性后，相关技术迅速在各个领域取得了一系列的成果。比如语义分割领域的DPT、SegFormer、HRFormer、Segmentor等均取得了非常优异的成绩。但是，相比CNN方案，Transformer方案因自注意力机制问题存在高计算量、高显存占用问题，导致其推理效率明显不如CNN方案。
作者认为：注意力机制在推理效率方面的瓶颈主要源自以下两个维度：

现有注意力机制的计算属性对于GPU设备不够友好，如二次复杂度、多头机制；
仅在高分辨率特征图实施注意力可能并非最有效捕获长距离上下文关系的方案，这是因为高分辨率特征的单个特征向量的感受野非常有限。

基于上述所提到的两个局限性，本文提出了一种GPU友好的注意力模块与跨分辨率注意力模块，并由此构建了RTFormer。
本文方案

接下来，我们首先对本文所提GPU友好注意力RTFormer模块进行介绍，然后结合如何基于RTFormer模块构建RTFormer分割架构。
RTFormer block

上图给出了本文所提RTFormer模块示意图，它是一种对偶分辨率模块，它包含两种类型注意力模块。在低分辨率分支，作者采用了GPU友好的注意力模块以捕获高层全局上下文信息，而在高分辨率分支，作者则引入了跨分辨率注意力机制对高层全局上下文信息进行传播扩散，也就是将两个分辨率的特征通过注意力模块进行聚合。

GPU-Friendly Attention 相比其他类型的注意力，作者发现：External Attention(EA，见上图左)因其线性复杂度可成为GPU设备的潜在选择。基于此，作者提出了GPU-Friendly Attention(GFA，见上图右)。我们首先对EA进行简要介绍，假设X \in R^{N \times d}表示输入特征，那么EA可以描述如下：
EA(X, K, V) = DN(X \cdot K^T) \cdot V \\
其中K， V \in R^{M \times d}表示可学习参数，M表示参数维度，DN表示Double Normalization。多头版EA可描述如下：
MHEA(X) = Concat(h_1, h_2, \cdots, h_H) \\ h_i = EA(X_i, K^{&#39;}, V^{&#39;}), i \in [1, H] \\
尽管EA采用共享K与V，可以一定程度加速计算，但Splitted机制仍存在。为避免多头机制在GPU设备上的延迟问题，作者提出了一种简单而有效的GPU友好注意力，它在EA的基础上演变而来，表示如下：
GFA(X, K_g, V_g) = GDN(X \cdot K_g^T) \cdot V_g \\
其中K_g, V_g \in R^{M_g \times d}, M_g = M \times H，GDN表示分组DN。对比GFA与MHEA可以发现GFA有两个主要改进：

它使得矩阵乘操作成为一体且非常适合于GPU设备；
它在某种程度上了保持了多头机制的优势。

Cross-resolution Attention  已有研究证实：多分辨率融合对于稠密预测任务非常有效。对于多分辨率架构设计，我们在不同分支独立执行GAF处理，然后再进行信息交互。作者认为：直接在高分辨率特征执行注意力对于全局上下文学习不够高效。为更有效的获得全局上下文信息，作者提出了跨分辨率注意力，它可以充分利用从低分辨率分支学到的高层语义信息。该过程可描述如下：
CA(X_h, K_c, V_c) = Softmax(\frac{X_h \cdot K_c^T}{\sqrt{d_h}}) \cdot V_c \\ K_c, V_c = \phi(X_c) \\ X_c = \theta(X_l)  \\
需要注意的是：考虑到GPU设备的快速推理因素，这里同样消除了多头机制。
Feed Forward Network  SegFormer与HRFormer中的FFN由两个MLP层与3\times 3深度卷积构成。但是，这种FFN对于GPU设备不够高效。为平衡性能与效率，作者在FFN中采用了两个3\times 3卷积且不进行通道维度扩展。
RTFormer

上图给出了RTFormer架构示意图，它由Backbone与SegHead两部分构成：

Backbone：为提取高分辨率特征所需的局部信息，作者将卷积与所提RTFormer模块进行了组合以构建RTFormer。具体来说，RTFormer的Stem、Stage1以及Stage2部分由卷积与残差模块构成，从Stage3开始采用所提对偶分辨率模块以促进高分辨率分支与低分辨率分支特征的信息交互，在后三个阶段，高分辨率特征的stride保持为8不变，而低分辨率的stride则分别为16、32、32。值得说明的是，Stage分与Stage5由本文所提RTFormer模块构成以促进高效全局上下文建模，而Stage3则仍由残差模块构建。
SegHead：在该部分，作者引入了DAPPM以进一步提取多尺度特征，此时输出特征的stride=8。最后，输出特征将被送入像素分类头(该分类头由3\times 3卷积与1\times 1卷积构成)以进行稠密语义标签。

上表给出了RTFormer-Slim与RTFormer-Base的架构配置信息，很明显，RTFormer的骨干部分由5个stage构成，其中3-5stage由对偶分辨率特征构成，在分割头方面，RTFormer引入了DPPM模块进一步提取多尺度特征。
本文实验

上表与图给出了Cityscapes与CamVid数据集上的性能对比，从中可以看到：

在Cityscapes数据集上，在所有实时分割方案中，RTFormer取得了最佳的速度-精度均衡；
RTFormer-Slim取得了76.3%mIoU指标且推理速度高达110.0FPS，优于STDC2-Seg75与DDRNet-23-Slim。
RTFormer-Base取得了79.3%mIoU指标且推理速度高达39.1FPS，取得了新的SOTA结果。
在CamVid数据集上，所提方案仅需ImageNet预训练即取得了82.5%mIoU指标且推理速度达94.FPS，优于采用额外Cityscapes预训练的STDC2-Seg；
RTFormer-Slim仅需4.8M参数即取得了81.4mIoU指标且推理速度高达190.7FPS，优于STDC2-Seg与DDRNet-23；
从视觉效果方面来看，RTFormer-Base具有更佳的细粒度分解结果。

上表与图给出了ADE20K数据集上不同方案的性能对比，从中可以看到：

RTFormer-Base取得了42.1%mIoU指标且推理速度达71.4FPS，优于其他方案；
相比DDRNet-23-Slim，RTFormer-SLim取得了36.7%mIoU指标，同时保持相当的速度。
从视觉效果可以看到：相比DDRNet-23，所提方案的分割结果具有更好的细节与上下文信息。

Code解析

本文代码已经开源在PaddleSeg仓库，笔者对其核心code进行摘录如下：
class RTFormer(nn.Layer):
def __init__(self, ...):
      super().__init__()
      ...
def forward(self, x):
      x1 = self.layer1(self.conv1(x))
      x2 = self.layer2(self.relu(x1))

      # stage-3
      x3 = self.layer3(self.relu(x2))
      x3_ = x2 + F.interpolate(self.compression3(x3), size=paddle.shape(x2)[2:], mode=&#39;bilienar&#39;)
      x3_ = self.layer3_(self.relu(x3_))

      # stage-4与stage-5的计算流程类似stage-3
      x4_ = ...
      x5_ = ...

      # SegHead, DAPPM
      x6 = self.spp(x5)
      x6 = x6 + F.interpolate(x6, size=paddle.shape(x5_), mode=&#39;bilinear&#39;)
      x_out = self.seghead(paddle.concat([x5_, x6], axis=1))
      return F.interpolate(x_out, paddle.shape(x)[2:], mode=&#39;bilinear&#39;)

class ExternalAttention(nn.Layer):
def __init__(self, ...)
   super().__init__()

def _act_sn(self, x):
   x = x.reshape([-1, self.inter_channels, 0, 0]) * (self.inter_channels ** -0.5)
   x = F.softmax(x, axis=1)
   x = x.reshape([1, -1, 0, 0])
def _act_dn(self, x):
   x_shape = paddle.shape(x)
   h, w = x_shape[2], x_shape[3]
   x = x.reshape([0, self.num_heads, self.inter_channels //self.num_heads, -1])
   x = F.softmax(x, axis=3)
   x = x / (paddle.sum(x, axis=2, keepdim=True) + 1e-06)
   x = x.reshape([0, self.inter_channels, h, w])

def forward(self, x, cross_k=None, cross_v=None):
   x = self.norm(x)
   if not self.use_cross_kv:
         x = F.conv2d(x, self.k, bias=None, stride=2 if not self.same_in_out_chs else 1, padding=0)
         x = self._act_dn(x)  # n,c_inter,h,w
         x = F.conv2d(x, self.v, bias=None, stride=1, padding=0)
      else:
         B = x.shape[0]
         x = x.reshape([1, -1, 0, 0])  # n,c_in,h,w -> 1,n*c_in,h,w
         x = F.conv2d(x, cross_k, bias=None, stride=1, padding=0, groups=B)
         x = self._act_sn(x)
         x = F.conv2d(x, cross_v, bias=None, stride=1, padding=0, groups=B)
         x = x.reshape([-1, self.in_channels, 0, 0])
      return x
往期精选

数据集汇总：

人脸识别常用数据集大全
行人检测数据集汇总
10个开源工业检测数据集汇总
21个深度学习开源数据集分类汇总（持续更新）
小目标检测、图像分类、图像识别等开源数据集汇总
人体姿态估计相关开源数据集介绍及汇总
小目标检测相关开源数据集介绍及汇总
医学图像开源数据集汇总
自动驾驶方向开源数据集资源汇总
目标检测开源数据集汇总（二）
RGB-T 相关开源数据集资源汇总
图像去雾开源数据集资源汇总
图像分类方向优质开源数据集汇总（附下载链接）

顶会资源：

ECCV22 最新54篇论文分方向整理｜包含目标检测、图像分割、监督学习等（附下载）
CVPR 2022 全面盘点：最新250篇论文分方向汇总 / 代码 / 解读 / 直播 / 项目（更新中）
CVPR&#39;22 最新106篇论文分方向整理｜包含目标检测、动作识别、图像处理等32个方向
一文看尽 CVPR2022 最新 22 篇论文（附打包下载）
17 篇 CVPR 2022 论文速递｜涵盖 3D 目标检测、医学影像、车道线检测等方向CVPR 2021 结果出炉！最新500篇CVPR&#39;21论文分方向汇总（更新中）
CVPR 2021 结果出炉！最新600篇CVPR&#39;21论文分方向汇总（更新中）
CVPR 2020 Oral 汇总：论文/代码/解读（更新中）
CVPR 2019 最全整理：全部论文下载，Github源码汇总、直播视频、论文解读等
CVPR 2018 论文解读集锦（9月27日更新）
CVPR 2018 目标检测（object detection）算法总览
ECCV 2018 目标检测（object detection）算法总览（部分含代码）
CVPR 2017 论文解读集锦（12-13更新）
2000 ~2020 年历届 CVPR 最佳论文汇总

技术综述：

万字长文 | 手把手教你优化轻量姿态估计模型（算法篇）
工业应用中如何选取合适的损失函数（MAE、MSE、Huber）-Pytorch版
综述：图像处理中的注意力机制
搞懂Transformer结构，看这篇PyTorch实现就够了
熬了一晚上，我从零实现了Transformer模型，把代码讲给你听
YOLO算法最全综述：从YOLOv1到YOLOv5
图像匹配大领域综述！涵盖 8 个子领域，近 20年经典方法汇总
一文读懂深度学习中的各种卷积
万字综述｜核心开发者全面解读Pytorch内部机制
19个损失函数汇总，以Pytorch为例
一文看尽深度学习中的15种损失函数
14种异常检测方法总结
PyTorch常用代码段合集
神经网络压缩综述

理论深挖：

深入探讨：为什么要做特征归一化/标准化？
令人“细思极恐”的Faster-R-CNN

论文盘点：

图像分割二十年，盘点影响力最大的10篇论文
2020年54篇最新CV领域综述论文速递！涵盖14个方向：目标检测/图像分割/医学影像/人脸识别等

实践/面经/求学：

如何配置一台深度学习工作站？
国内外优秀的计算机视觉团队汇总
50种Matplotlib科研论文绘图合集，含代码实现
图像处理知多少？准大厂算法工程师30+场秋招后总结的面经问题详解
深度学习三十问！一位算法工程师经历30+场CV面试后总结的常见问题合集（含答案）
深度学习六十问！一位算法工程师经历30+场CV面试后总结的常见问题合集下篇（含答案）
一位算法工程师从30+场秋招面试中总结出的目标检测算法面经（含答案）
一位算法工程师从30+场秋招面试中总结出的语义分割超强面经（含答案）

Github优质资源：

25个【Awsome】GitHub 资源汇总（更新中）
超强合集：OCR 文本检测干货汇总（含论文、源码、demo 等资源）
2019-2020年目标跟踪资源全汇总（论文、模型代码、优秀实验室）

最是忆江南 · 发表于 2025-4-28 05:28:32

站位支持

郑科荣 · 发表于 2025-5-5 06:17:42

LZ敢整点更有创意的不？兄弟们等着围观捏~

		自动登录	找回密码
密码			立即注册