《RTFormer: Efficient Design for Real-Time Semantic Segmentation with Transformer》

期刊：NIPS

年份：2022

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

摘要

最近，基于Tramsformer的网络在语义分割方面显示出令人印象深刻的结果。然而，对于实时语义分割，由于Tramsformer的耗时计算机制，纯基于 CNN 的方法在该领域仍然占主导地位。我们提出了 RTFormer，一种用于实时语义分割的高效双分辨率Tramsformer，与基于 CNN 的模型相比，它在性能和效率之间实现了更好的权衡。为了在类似 GPU 的设备上实现高推理效率，我们的 RTFormer 利用具有线性复杂度的 GPU-Friendly Attention 并丢弃多头机制。此外，我们发现通过扩展从低分辨率分支中学习到的高级知识，跨分辨率注意力更有效地为高分辨率分支收集全局上下文信息。在主流基准上的广泛实验表明了我们提出的RTFormer的有效性，它在cityscape、CamVid和COCOStuff上实现了最先进的性能，并在ADE20K上展示了有希望的结果。

Introduction

现有方法的局限性：

• 当前的实时语义分割方法主要集中在基于CNN的架构上，包括手工设计的高效骨干网络和解码器，以及通过神经架构搜索方法来寻找准确率和效率之间的平衡。
• 尽管这些方法取得了显著的改进，但它们在处理高分辨率特征时仍面临着计算量大的挑战，尤其是在密集预测任务如语义分割中。

Transformer的潜力：

• 近期，基于Transformer的视觉模型因其强大的视觉识别能力而受到关注。Transformer架构通过自注意力机制能够有效捕获长距离上下文信息，这对于语义分割是有益的。

主要贡献：

• 提出了一种新颖的 RTFormer 块，它在类似 GPU 的设备上实现了更好的性能和效率之间的权衡，用于语义分割任务。
• 提出了一种新的网络架构 RTFormer，它可以通过深度利用注意力来提高语义分割，而不会损失效率。
• RTFormer 在 Cityscapes、CamVid 和 COCOStuff 上实现了最先进的性能，并在 ADE20K 上展示了有希望的性能。此外，它为实时语义分割任务的实践提供了新的视角。

Related Work

基于深度学习的方法：

• 随着深度学习的发展，基于全卷积网络（FCN）的方法在各种基准测试中取得了显著的性能提升。这些方法从不同角度改进了FCN，例如：
  • Deeplabv3和PSPPNet通过引入空洞空间金字塔池化模块和金字塔池化模块来扩大感受野并融合不同级别的特征。
  • SegNet通过编码器-解码器结构恢复高分辨率图。
  • HRNet引入了一种多分辨率架构，在整个网络中保持高分辨率特征图。

实时语义分割：

• 为了解决实时分割问题，提出了多种方法，例如：
  • ICNet使用多分辨率图像级联网络。
  • FasterSeg利用神经架构搜索（NAS）来平衡高准确性和低延迟。
  • BiSeNetV1和BiSeNetV2采用双流路径网络和特征融合模块，在速度和分割性能之间取得平衡。
  • STDC通过提出单流结构和细节引导模块来改进BiSeNet。
  • DDRNet通过设计双深度分支网络和深度聚合金字塔池化模块来提高性能。

注意力机制：

• 注意力机制在计算机视觉领域得到了广泛应用，例如：
  • SE块将注意力函数应用于通道，提高了网络的表示能力。
  • 非局部操作能够捕获长距离依赖关系，在视频分类任务上表现出色。
  • 自注意力是注意力机制的一种特殊形式，近年来得到了广泛应用，但其二次复杂度限制了其使用。

Transformer在语义分割中的应用：

• 最近，Transformer在语义分割中表现出了潜力：
  • DPT将Transformer作为编码器来提高密集预测任务的性能。
  • SETR提出了一种序列到序列的方法，取得了令人印象深刻的结果，但由于其重型骨干网络和非常高的分辨率，难以用于实时分割任务。
  • SegFormer通过引入分层Transformer编码器和轻量级全MLP解码器来提高效率，但与一些最先进的基于CNN的实时分割模型相比，其效率仍然相对较低。

Method

3.1 RTFormer block

它采用了双分辨率模块设计，用于有效捕获全局上下文信息以提高语义分割的性能。对于低分辨率，应用了 GPU-Friendly Attention。对于高分辨率，引入了一个交叉分辨率注意，将低分辨率分支学习到的高级全局上下文广播到每个高分辨率像素，并利用阶梯布局将低分辨率分支更具代表性的特征输入到交叉分辨率注意中。此外，我们用两个 3 × 3 卷积层组成 FFN。

GPU-Friendly Attention： 

Multi-Head 外部注意力和 GPU-Friendly Attention 之间的比较。多头外部注意力将矩阵乘法分成几个组，而我们的 GPU-Friendly Attention 使矩阵乘法集成起来，这对类似 GPU 的设备更友好。

GFA从外部注意力（External Attention, EA）演化而来，具有线性复杂度，适合在GPU上高效执行。

 GDN表示分组双归一化，将原始双归一化的第二个归一化分成H组 。

评价：

• 它使得矩阵乘法可以被集成，这对于类似gpu的设备来说是非常友好的。
• 利用分组双归一化，在一定程度上保持了多头机制的优越性。为了直观的理解，可以认为GFA也生成了h个不同的注意图来捕捉不同的token之间的关系，但是涉及更多的特征元素来计算相似度，所有的注意图都对最终的输出有贡献。

Cross-resolution Attention：

在高分辨率分支中，像素专注于超过高级全局上下文的局部信息。因此，我们假设直接对高分辨率特征图进行注意力来学习全局上下文是不够的。为了更有效地获得全局上下文，我们提出了一种交叉分辨率注意力，旨在充分利用从低分辨率分支中学习到的高级语义信息。与 GFA 不同，高分辨率分支中采用交叉分辨率注意力来收集全局上下文。并且这种交叉分辨率注意力的计算表示为：

其中

• 表示高分辨率分支和低分辨率分支的特征 
• 表示为交叉特征，由池化层和卷积层组成的函数 θ 从  计算
• 是一组矩阵运算，包括分裂、排列和重塑

Feed Forward Network:

在之前的基于Transformer的分割方法中，前馈网络(FFN)通常由两个MLP层和一个深度3 × 3卷积层组成，其中深度3 × 3层用于补充位置编码或增强局部性。此外，两个 MLP 层将隐藏维度扩展为输入维度的 2 倍或四倍。这种类型的 FFN 可以在相对更少的参数下获得更好的性能。但是在应考虑类似 GPU 的设备上的延迟的情况下，FFN 的典型结构不是很有效。为了平衡性能和效率，在 RTFormer 块的 FFN 中采用了两个没有维度扩展的 3×3 卷积层。它显示出比典型的 FFN 配置更好的结果。

3.2 RTFormer

将 RTFormer 块放置在最后两个阶段，这些阶段由粉红色块指示，并在由蓝色块指示的早期阶段使用卷积块。此外，添加了一个用于分割头的DAPPM模块

Backbone Architecture： 

为了提取高分辨率特征图所需的足够的局部信息，我们将卷积层与我们提出的RTFormer块结合起来构建RTFormer。具体来说

• 让RTFormer从一个由两个3 × 3卷积层组成的干块开始，并由几个连续的基本残差块组成前两个阶段[。
• 然后，从stage3开始，我们使用双分辨率模块，使高分辨率和低分辨率分支之间的特征交换成为可能，灵感来自。
• 对于后三个阶段的高分辨率分支，特征步长保持8不变，而对于低分辨率分支，特征步长分别为16、32、32。特别地，我们将双分辨率模块排列成阶梯式布局，借助低分辨率分支的输出来增强高分辨率特征的语义表示。
• 最重要的是，我们使用我们提出的RTFormer块构建stage4和stage5，如图所示，用于高效的全局上下文建模，而stage3仍然由基本残差块组成。

Segmentation Head：

对于RTFormer的分割头，我们在低分辨率输出特征后添加了一个DAPPM模块，。在将DAPPM的输出与高分辨率特征融合后，我们得到了stride=8的输出特征图。最后，将该输出特征传递到像素级分类头中，用于预测密集语义标签。分类头由 3 × 3 卷积层和 1 × 1 卷积层组成，隐藏特征维度与输入特征维度相同。

Instantiation：

我们用RTFormer-Slim和RTFormerBase实例化RTFormer的体系结构，详细配置如表所示。对于通道数和块数，每个数组包含5个元素，分别对应5个阶段。特别是，两个数字的元素对应于双分辨率阶段。例如，64/128 表示高分辨率分支的通道数为 64，低分辨率分支的通道数为 128。虽然 1/2 表示高分辨率分支的基本卷积块数为 1，低分辨率分支为 2。值得注意的是，块数数组中最后两个元素表示RTFormer块的数量，RTFormer-Slim和RTFormer-Base均为1。RTFormer-Slim 和 RTFormer-Base 的交叉特征空间大小分别设置为 64(8 × 8) 和 144(12 × 12)。

Result

与 Cityscapes 和 CamVid 上其他最先进的实时方法的比较。城市景观和CamVid的性能分别用1024 × 2048、720 × 960的单作物进行测量。#Params 是指参数的数量。FPS 的计算方式与性能测量相同。在此表中，* 表示我们重新训练此方法遵循其原始训练设置，u 表示我们测量单个 RTX 2080Ti GPU 上的 FPS。

 Conclusion

在本文中，我们提出了 RTFormer，它可以有效地捕获全局上下文以提高实时语义分割性能。大量实验表明，我们的方法不仅在用于实时分割的公共数据集上取得了新的最先进的结果，而且在用于一般语义分割的挑战性数据集上也显示出卓越的性能。由于 RTFormer 的效率，我们希望我们的方法可以鼓励使用 Transformer 进行实时语义分割的新设计。一个限制是，虽然我们的 RTFormer-Slim 只有 4.8M 参数，但在边缘设备芯片中可能需要更多的参数效率。我们将其留给未来的工作。