RTFormer: Efficient Design for Real-Time Semantic Segmentation with Transformer概述
0.前言
-
发表时间：NeurlPS2022(2022.10.13)
1.摘要
最近，基于Transformer的网络在语义分割方面取得了令人印象深刻的结果。然而在实时语义分割方面，由于Transformer的计算机制耗时，纯基于cnn的方法仍然在这一领域占主导地位。我们提出了一种高效的实时语义分割双分辨率TransformerRTprorr，它比基于CNN的模型在性能和效率之间实现更好的权衡。为了在GPU这类设备上实现高推理效率，我们的RTformer利用了线性复杂度的GPU友好注意力，并抛弃了多头机制。此外，我们发现交叉分辨率注意通过传播从低分辨率分支获得的高级知识，可以更有效地收集高分辨率分支的全局上下文信息。在主流基准上的大量实验证明了我们提出的RTformer的有效性，它在Cityscapes，CamVid和COCOStuff上取得了最先进的水平，并在ADE20K上显示出了良好的结果。
2.主要贡献
•提出了一种新的RTFormer块，它在GPU类设备上实现了语义分割任务的性能和效率之间更好的权衡。
•提出了一种新的网络架构RTFormer，该架构能够在不降低效率的前提下，深度利用注意力，充分利用全局上下文来改进语义分割。
•RTFormer在Cityscapes、CamVid和COCOStuff上达到了最先进的水平，并在ADE20K上显示出有前景的性能。此外，它为实时语义分割任务的实践提供了一个新的视角。
3.方法
RTFrormerer块包括两个分支，一个分支使用GPU友好型注意力，另一个使用跨分辨率注意力。
GPU友好型注意力是由External Attention(EA)衍生而来的。
EA基于两个外部的，小的，可学习的和共享的存储器，只用两个级联的线性层和归一化层就取代了现有流行的学习架构中的Self-attention，具体来说，通过两个由线性层实现的包含整个数据集特征的记忆单元，维度为Mxd，替换K和V，从而既使得每个样本能与其他样本交互，又减少了计算量。公式表示如下：
图中H表示头数，EA使用了多头机制，虽然通过共享K'和V'加快了计算速度，但是矩阵乘法操作依旧存在。DN表示double-normalization，分别对行和列进行归一化。
为了减少多头机制带来的计算量，作者提出了GPU-Friendly Attention，公式为：
其中Kg，Vg ∈RMg×d，Mg = M × H, GDN表示分组双归一化，将原双归一化的第二次归一化分成H组，EA和GFA对比如下图：
GFA有两个主要的改进。一方面，它能集成矩阵乘法，这对于gpu类设备来说是相当友好的。得益于此，可以将外部参数的大小从(M, d')扩大到(Mg， d)，因此可以调优更多的参数来提高性能。另一方面，利用分组双归一化，在一定程度上保持了多头机制的优越性。为了直观理解，可以认为GFA也生成了H个不同的注意力图来捕捉不同的token之间的关系，但相似度计算涉及了更多的特征元素，且所有的注意力图都对最终的输出有贡献。
跨分辨率注意力 多分辨率融合已被证明对密集预测任务是有效的。而对于多分辨率架构的设计，可以直观地将GFA独立应用于不同分辨率的分支，并在执行卷积模块或注意力模块后交换特征。但在高分辨率分支中，像素更关注局部信息，而非高层次的全局上下文。因此，作者认为直接将注意力转移到高分辨率特征图上学习全局上下文是不够有效的。为了更有效地获取全局上下文，提出了跨分辨率注意力机制，旨在充分利用从低分辨率分支中学习到的高层语义信息。如下图所示，计算公式如下：
其中Xh, Xl分别表示高分辨率分支和低分辨率分支上的特征映射，φ是一组包括splitting，permutation和reshaping的矩阵操作，dh表示高分辨率分支的特征维数。Xc为交叉特征，θ函数由池化层和卷积层组成，Xc的空间大小表示从xl生成的token数。在实验上，只在注意力图的最后一个轴上使用softmax进行归一化，因为当key和value不是外部参数时，单个softmax比双归一化性能更好。此外，为了计算性能，这里没有使用多头机制。
Feed Forward Network 在以前的基于Transformer的分割方法中，前馈网络（FFN）通常由两个MLP层和一个深度3×3卷积层组成，其中深度3×3卷积用于补充位置编码或增强局部性。此外，两个MLP层将隐藏维度扩展为输入维度的2到4倍。这种类型的FFN可以用相对较少的参数获得更好的性能。但在应该考虑类GPU设备上的延迟的情况下，FFN的典型结构不是很有效。为了平衡性能和效率，在RTFormer块的FFN中采用了两个3×3卷积层，而没有进行维数扩展。与典型的FFN配置相比，它显示了更好的结果。
RTFrormerer块如下。对于低分辨率，采用GPU友好型注意来捕获高层次的全局上下文。对于高分辨率，使用跨分辨率注意力，将从低分辨率分支学习到的高层全局上下文传播到每个高分辨率像素，并通过stage布局将更有代表性的特征从低分辨率分支提供给交叉分辨率注意力。它从低分辨率分支中提取K和V。此外，用两个3×3卷积层组成FFN。
RTFrormer体系结构说明。将RTFrormer块放置在最后两个阶段，用粉色块表示，在早期阶段使用卷积块，用蓝色块表示。此外，利用DDRNet的成功经验，添加了一个针对分割头的DAPPM(Deep Aggregation Pyramid Pooling Module)模块。
为了提取高分辨率特征图所需的足够的局部信息，将卷积层与RTFormer块相结合来构造RTFormer。具体来说，让RTFormer从一个由两个3×3卷积层组成的stem块开始，用几个连续的基本残差块组成前两个阶段。然后，从第3阶段开始，使用双分辨率模块，实现高分辨率和低分辨率分支之间的特征交换。对于最后三个阶段的高分辨率分支，特征stride保持为8不变，而对于低分辨率分支，则特征stride分别为16、32、32。特别是，将双分辨率模块安排为阶梯式布局，以借助低分辨率分支的输出增强高分辨率特征的语义表示。最重要的是，用提出的RTFormer块构造了stage4和stage5，如图2所示，用于有效的全局上下文建模，而stage3仍然由基本残差块组成。
4.结论
在本文中，我们提出了可以有效地捕获全局上下文，以提高实时语义分割性能的RTFormer。大量的实验表明，我们的方法不仅在通用的实时分割数据集上取得了最新的结果，而且在具有挑战性的一般语义分割数据集上表现出了优越的性能。由于前者的效率，我们希望我们的方法可以鼓励新的实时语义分割设计。一个限制是，虽然我们的RTFormer Slim只有4.8M参数，但边缘设备芯片可能需要更高的参数效率。我们把它留给以后的工作吧。
5.补充
DAPPM(Deep Aggregation Pyramid Pooling Module)模块结构如下：
用于进一步从低分辨率的特征图中提取上下文信息。以层次-残差的方式融合不同尺度的上下文信息。
