论文阅读笔记：RTFormer: Efficient Design for Real-Time Semantic Segmentation with Transformer

代码：https://github.com/PaddlePaddle/PaddleSeg

论文：https://arxiv.org/pdf/2210.07124v1.pdf

1 背景

与基于CNN的网络相比，基于Transformer的网络的主要区别在于自注意力的大量使用，而自注意力擅长捕捉语义分割中必不可少的长距离上下文信息。因此，作者认为注意力结构在实时语义分割任务中也应该是有效的。

但是，注意力在实时场景中应用的主要障碍可能来自以下两个方面：

1）现有的大多数类型的注意力的计算特性对于类似GPU的设备来说不易于推理，例如二次复杂度和多头机制。二次复杂度在处理高分辨率时引入了较大的计算负担，尤其是在语义分割等密集预测任务中。而多头机制将矩阵拆分成多个组，使得注意力操作在类GPU设备上耗时增加，类似于执行组卷积。

2）由于高分辨率特征图的单个特征向量的感受野有限，仅对高分辨率特征图本身进行关注可能不是捕获具有高层语义信息的长距离上下文的最有效方法。

2 创新点

1）提出了一种新的RTFormer模块，在语义分割任务的GPU类设备上实现了性能和效率的较好折中；

2）提出了一种新的网络架构RTFormer，在不损失效率的情况下，充分利用全局上下文，通过深度利用注意力来改进语义分割；

3） RTFormer在Cityscapes、CamVid和COCOStuff上取得了最先进的水平，在ADE20K上表现出良好的性能，如图1。此外，为实时语义分割任务的实践提供了新的视角。

3 方法

为了在带有Transformer的类GPU设备上实现性能和效率的更好折中，作者提出一种新的transformer block，命名为RTFormer block，如图2。

对于低分辨率分支，采用了一种新提出的基于外部注意力的GPU友好注意力，继承了外部注意力的线性复杂度特性，并通过丢弃矩阵乘法运算中的通道分裂，缓解了类GPU设备多头注意力机制的弱点。

外部注意力（External Attention，EA）这一概念出自论文《Beyond self-attention: External attention using two linear layers for visual tasks》。这篇论文认为，如果希望注意力机制可以考虑到来自其他样本的影响，那么就需要一个所有样本共享的特征。为此，作者引入一个外部的记忆单元M，来刻画所有样本最本质的特征，并用M来表示输入特征：

$A=Norm(F{M}_k^T)$

$F_{out}=A{M}_v$

其中，$F$ 表示输入特征，$M_k$ 和 $M_v$ 是记忆单元。代码中分别使用两个线性层的权重对 $F$ 进行注意力计算。

外部注意力机制如下：

伪代码如下：

多头外部注意力机制如下：

伪代码如下：

对于高分辨率分支，作者采用交叉分辨率注意力，而不是只在高分辨率特征本身进行注意力。将两个分辨率分支排列成阶梯式布局，借助于从低分辨率分支学习到的全局上下文信息，可以更有效的增强高分辨率分支。

4 模块

4.1 GPU友好注意力

在介绍GPU友好注意力（GFA）之前，先回顾一下外部注意力EA。$X\in R^{N\times d}$ 表示一个输入特征，其中 $N$ 为元素个数(或图像中的像素)，$d$ 为特征维数，则EA的表达式可以表示为：

其中 $K,V\in R^{M\times d}$ 是可学习的参数，$M$ 是参数维度， $DN$ 是双归一化操作。

双归一化是指分别对列和行进行归一化，在第一维上使用softmax算子，第二维上使用 l1-norm 来标准化注意图。

原文说这么做可以使不同查询点的注意力图差异增加并且在语义上是有意义的，让不同查询点把更多的注意力放在各自不同的区域。

多头版的EA可以表示为：

其中 $K’,V’\in R^{M\times d’},d^{\prime }=d/H$，$H$ 是多头注意力的头的数量，$X_i$ 是输入特征 $X$ 的第 $i$ 个头，如图所示。

多头机制产生H个注意力图以提高EA的容量，这使得矩阵相乘被分成若干组，相当于组卷积。虽然EA对不同的头使用共享的 $K^{\prime }$ 和 $V^{\prime }$，可以大大的加快计算速度，但分裂后的矩阵乘法依然存在。

为了避免多头机制导致的GPU设备速度降低，作者提出了一种简单有效的GPU友好注意力机制。由式（1）表示的外部注意力演化而来，可表示为：

其中 $K_g,V_g\in R^{M_g\times d},M_g=M\times H$，GDN表示分组双重归一化（Grouped Double Normalization），即将原本的双重归一化的第二个归一化拆分成H组，如图所示。

从式（3）可以发现GFA有两个主要的改进。一方面它使得矩阵乘法得以集成，这对于类似GPU的设备来说是相当友好的，得益于此。作者将外部参数的大小$(M,d)$ 扩大到了 $(M_g,d)$，因此可以调整更多的参数来提高性能。另一方面，利用GDN的优势，一定程度上保持了多头机制的优越性。为了直观理解，可以认为GFA也生成了H个不同的注意力图，用于捕获不同的查询之间的关系，但需要更多的特征元素来计算相似度，所有的注意力图对最终的输出都有贡献。

4.2 交叉分辨率注意力

多分辨率融合已被证明对密集预测任务是有效的。而对于多分辨率架构的设计，可以直观地将GFA独立地应用在不同分辨率分支中，并在卷积模块或注意力模块执行后，进行特征交换。但在高分辨率分支中，像素更关注局部信息，而非全局上下文。因此作者认为直接在高分辨率特征上进行注意力学习来学习全局上下文是不够有效的。为了更有效地获取全局上下文，作者提出了交叉分辨率注意力机制，旨在充分利用低分辨率分支中学习到的高层语义信息，如图2所示。

计算过程可表示为：

其中 $X_j,X_l$ 表示高分辨率分支和低分辨率分支的特征图，$\Phi$ 是一组矩阵运算，包括split、permute和 reshape，$d_h$ 表示高分辨率分支的特征维数。特征图 $X_c$ 是由 $X_l$ 通过由池化层和卷积层组成的函数 $\theta$ 计算得到的，在后文中记为交叉特征。特征图 $X_c$ 的空间大小表示从 $X_l$ 生成的token的个数。

交叉注意力中的softmax仍然按原来的方式计算，作者通过实验发现，当键值不是外部参数时，单个softmax比双重归一化的表现更好。为了在类GPU设备上进行快速推理，交叉注意力中也舍弃了很多多头机制。

4.3 FFN

在以往的transformer-based的分割方法中，前馈网络FFN通常由两个MLP层和一个卷积层组成。考虑到类GPU设备场景，FFN的结构效率不是很高。为了平衡性能和效率，作者在RTFormer的FFN中采用了两个3×3的卷积层，并且没有进行维度扩展，与典型的FFN结构实验相比，本算法取得了更好的效果。

4.4 RTFormer

图4为RTFormer的全部结构。

DAPPM参考自论文《Deep dual-resolution networks for real-time and accurate semantic segmentation of road scenes 》

BackBone：为了提取高分辨率特征所需的局部信息，作者将卷积层与提出的RTFormer Bock相结合来构造RTFormer。具体来说，RTFormer由两个3×3的卷积开始，用几个连续的基础残差块组成前两个阶段，然后从stage 3开始，使用双分辨率模块来实现高分辨率分支和低分辨率分支之间的特征交换。对于三个阶段的高分辨率分支，特征步长保持为8不变。对于低分辨率分支，特征步长分别为16,32,32。特别地，作者将双分辨率模块排列成阶梯式分布，借助低分辨率分支的输出来提升高分辨率特征的语义表示。

Segmentation Head：对于RTFormer的分割头，作者在低分辨率输出特征后添加DAPPM模块，并将DAPPM的输出与高分辨特征进行融合，得到了下采样8倍的输出特征图。最后将输出特征传递给像素分类头，用于预测稠密语义标签。分类头由一个3×3的卷积层和一个1×1的卷积层组成。

Instantiation：RTFormer的体系结构分为RTFormer-Slim 和 RTFormer-Base。详细的配置记录在表1中。有斜杠的，斜杠前是高分辨率分支的配置，斜杠后是低分辨率分支的配置。

5 实验

5.1 和SOTA方法的比较

在单张RTX 2080Ti上，分别使用 1024×2048 和 720×960 的尺寸推理 Cityscapes 和 CamVid 数据集，与SOTA方法的对比如表2。

在 CamVid 测试集上，与DDRNet-23的效果对比，如图5。

在单张RTX 2080Ti上，使用 512×2048 的尺寸推理 ADE20K 数据集，与SOTA方法对比效果如表3。

在 ADE20K 验证集上，与DDRNet-23的效果对比，如图6。

使用 640 × 640 的尺寸推理 COCOStuff 数据集，与SOTA方法对比效果如表4。

5.2 消融实验

1. 为了验证提出的GPU友好注意力的有效性，将RTFormer块中使用的注意力替换成不同的类型和组合。如表5a所示，给出了多头自注意力，多头外注意力，GPU友好注意力和交叉分辨率注意力的不同组合的结果，GFA+CA是指在低分辨率分支中使用GFA，在高分辨率分支中使用CA，此外，通过 $M=d\times r$ 来调整对头外部注意中的超参M，其中 r 为衰减比例。可以看到GFA优于多头外部注意力的所有设置。且当引入最佳交叉分辨率注意力后，性能进一步提升，而FPS下降才不到2。

2. 表5b为两层MLP加一层卷积的经典FFN结构和本文提出的两层卷积FFN的对比实验。可以看到本文提出的FFN在速度和效果上都占优。

3. 表5c为分组双重归一化中不同的分组数的对比实验，第一个数字为低分辨率中的配置，第二个数字为高分辨率中的设置，发现分组对速度几乎没有影响（因为矩阵乘法的计算维度不变，只是归一化次数变多），且8+2的参数使性能达到最佳。

4. 表5d为交叉分辨率注意中交叉特征的空间大小（即低分辨率输入经过特征映射后的token个数）对比实验，根据FPS和mIoU之间的折衷，RTFormer - Slim的交叉特征空间在8 × 8大小是最好的。这在一定程度上说明低分辨率的交叉特征接近高分辨率特征维数的空间尺寸是合适的。
   

6 结论

虽然RTFormer-Slim只有4.8M的参数，但在边缘设备中可能还需要更高的参数效率。