点击上方“AI公园”,关注公众号,选择加“星标“或“置顶”
导读
来自baidu,使用Transformer达到实时推理的语义分割。
摘要
最近,基于Transformer的网络在语义分割任务中展现了令人印象深刻的效果。然而,在实时语义分割领域,纯基于CNN的方法仍然占据主导地位,这是由于Transformer的时间消耗计算机制。我们提出了RTFormer,这是一种高效的双分辨率变换器,用于实时语义分割,它在性能和效率之间实现了比基于CNN的模型更好的权衡。为了在GPU等设备上实现高效的推理,我们的RTFormer利用了线性复杂度的GPU友好的注意力机制,并放弃了多头机制。此外,我们发现跨分辨率注意力更有效地通过传播从低分辨率分支学到的高级知识来为高分辨率分支收集全局上下文信息。在主流基准上的广泛实验表明了我们提出的RTFormer的有效性,它在Cityscapes、CamVid和COCOStuff上达到了最前沿的性能,并在ADE20K上显示出了有希望的结果。代码可在PaddleSeg获取:https://github.com/PaddlePaddle/PaddleSeg。
1、介绍
语义分割是一项基础的计算机视觉任务,通常作为自动驾驶、移动应用、机器人感知等关键感知模块。随着这些应用的发展,实时执行语义分割的需求越来越强烈。现有的实时分割方法主要集中在利用CNN架构,包括手工设计高效的骨干网络和解码器,以及探索神经架构搜索方法以寻找准确性和效率之间的更好平衡。迄今为止,这些优秀的工作已经取得了显著的进步。
最近,视觉Transformer因其强大的视觉识别能力而受到了广泛关注。继承自它们的一系列基于Transformer的架构被提出,并在一般的语义分割任务上显示出了非常有前景的表现。与基于CNN的网络相比,这些基于Transformer的架构的主要区别在于大量使用了自注意力机制,而自注意力擅长捕捉对于语义分割至关重要的长距离上下文信息。因此,我们认为注意力结构也应该在实时语义分割任务中有效。
但是到目前为止,只有少数工作探索了注意力在实时语义分割领域的应用,而最前沿的技术仍由基于CNN的架构主导。我们认为在实时设置中应用注意力的主要障碍可能来自以下两个方面。一是大多数现有类型的注意力计算特性对于GPU等设备来说不够友好,例如二次复杂度和多头机制。二次复杂度在处理高分辨率特征时会引入大量的计算负担,特别是在像语义分割这样的密集预测任务中。尽管一些工作如减小了key和value的尺寸,但二次复杂度的特性仍然存在。而多头机制将矩阵乘法分成多个组,这使得注意力操作在GPU等设备上耗时,类似于执行分组卷积的情况。另一个方面是在高分辨率特征图本身上执行注意力可能不是捕获包含高级语义信息的长距离上下文的最有效方式,因为高分辨率特征图中的单个特征向量具有有限的感受野。
我们提出了一种新型的Transformer块,命名为RTFormer块,如图 2 所示,旨在在GPU等设备上实现基于Transformer的性能和效率之间的更好权衡。对于低分辨率分支,我们采用了新提出的GPU-Friendly Attention,它源自外部注意力。它继承了外部注意力的线性复杂度特性,并通过放弃矩阵乘法运算中的通道分裂来缓解多头机制对于GPU等设备的弱点。相反,它增加了外部参数的数量,并将外部注意力提出的双重规范化操作中的第二个规范化分成多个组。这使得GPU-Friendly Attention能够在一定程度上保持多头机制的优势。对于高分辨率分支,我们采用了跨分辨率注意力,而不是仅仅在高分辨率特征自身内部进行注意力操作。此外,与多分辨率融合并行公式不同,我们将两个分辨率分支排列成阶梯式布局。因此,高分辨率分支可以通过辅助从低分辨率分支学到的高级全局上下文信息得到更有效的增强。基于提出的RTFormer块,我们构建了一个新的实时语义分割网络,命名为RTFormer。为了学习足够的局部上下文,我们仍然在早期阶段使用卷积块,并将RTFormer块放置在最后两个阶段。通过广泛的实验,我们发现RTFormer能够更有效地利用全局上下文,并且实现了比之前的工作更好的权衡。图 1 显示了RTFormer与其他方法在CamVid上的比较。最终,我们总结了RTFormer的贡献如下三个方面:
提出了一种新型的RTFormer块,它在GPU等设备上针对语义分割任务实现了性能和效率之间的更好权衡。
提出了一种新的网络架构RTFormer,它能够通过深度利用注意力来充分利用全局上下文,从而在不牺牲效率的情况下提高语义分割的质量。
RTFormer在Cityscapes、CamVid和COCOStuff上达到了最前沿的性能,并在ADE20K上显示出了有希望的结果。此外,它为实时语义分割任务的实际应用提供了一个新的视角。
2 相关工作
通用语义分割。传统的分割方法利用手工制作的特征来解决像素级别的标签分配问题,例如阈值选择、超像素和图算法。随着深度学习的成功,一系列基于全卷积神经网络(FCN)的方法在各种基准上实现了卓越的性能。这些方法从不同的方面改进了FCN。DeepLabv3和PSPNet通过引入空洞空间金字塔池化模块和金字塔池化模块扩大了感受野,并融合了不同层次的特征。SegNet通过编码器-解码器结构恢复了高分辨率地图。HRNet引入了一个多分辨率架构,该架构在整个网络中维持了高分辨率特征图。OCRNet通过查询全局上下文增强了主干网络输出的特征。
实时语义分割。为了解决实时分割问题,已经提出了多种方法。ICNet通过使用精心设计的多分辨率图像级联网络解决了这一问题。FasterSeg利用神经架构搜索(NAS)来达到高精度与低延迟之间的平衡。BiSeNetV1和BiSeNetV2采用双流路径网络和特征融合模块以实现速度与分割性能之间的良好平衡。STDC通过提出带有细节引导模块的单流结构重新思考并改进了BiSeNet。DDRNets通过设计具有多重双边融合和深度聚合金字塔池化模块的双深层分支网络实现了更好的性能。
注意力机制。注意力机制在计算机视觉领域得到了蓬勃发展。自注意力是近年来被广泛使用的注意力机制的一种特殊情况。然而,其二次复杂度限制了它的应用。一些工作对自注意力进行了改进,以实现线性复杂度。但它们仍然不太适合在GPU上进行推理。受外部注意力的启发,我们开发了一种适用于GPU的注意力模块,该模块在类似GPU的设备上具有低延迟和高性能。
Transformer在语义分割中的应用。最近,Transformer在语义分割中表现出优异的性能。DPT将Transformer作为编码器来提高密集预测任务的表现。SETR提出了一种序列到序列的方法,并取得了令人印象深刻的结果。SETR使用预训练的VIT作为其骨干网络,在空间分辨率上没有降采样。然而,由于其骨干网络较重且分辨率非常高,难以用于实时分割任务。SegFormer通过引入分层Transformer编码器和轻量级全MLP解码器提高了效率。与SETR相比,SegFormer既具有更高的效率也具有更高的性能。然而,与一些最先进的基于CNN的实时分割模型相比,SegFormer的效率仍然相对较低。通过引入我们的RTFormer模块,我们的方法可以在实现实时速度的同时利用注意力机制的优势。
方法
本节详细阐述我们所提出的方案。首先,我们描述RTFormer模块,然后介绍如何基于RTFormer模块构建RTFormer。
3.1 RTFormer模块
RTFormer模块是一个双分辨率模块,继承了多分辨率融合范式。与之前的工作不同,RTFormer模块由两种类型的注意力及其前馈网络组成,并以阶梯式布局排列,如图2所示。在低分辨率分支中,我们使用适用于GPU的注意力来捕捉高级别的全局上下文。而在高分辨率分支中,我们引入了跨分辨率注意力来将从低分辨率分支学到的高级别全局上下文广播到每个高分辨率像素,阶梯式布局旨在从低分辨率分支向跨分辨率注意力输入更具代表性的特征。
GPU-Friendly Attention
比较现有的不同类型注意力,我们发现外部注意力(EA)因其线性复杂度的良好特性,可以作为一种潜在的选择在类似GPU的设备上执行,而我们的GPU-Friendly Attention(GFA)就是基于它发展而来的。因此,在详细介绍GFA之前,我们先回顾一下EA。令 表示一个输入特征,其中 N 是元素的数量(或图像中的像素数量),d是特征维度,那么EA的公式可以表示为:
其中 是可学习的参数,𝑀是参数维度,𝐷𝑁是双重归一化操作。而EA的多头版本可以表示为:
其中 ,𝑑'=𝑑/𝐻,𝐻是头的数量,而 𝑋𝑖是 𝑋 的第 𝑖个头。如图3上部所示,多头机制为EA生成了 𝐻个注意力图,这提高了EA的容量,并使得矩阵乘法被分成多个组,类似于组卷积。尽管EA为不同的头共享 𝐾'和 𝑉',这可以大大加快计算速度,但分组矩阵乘法依然存在。
为了避免由于多头机制导致在类似GPU的设备上延迟降低的问题,我们提出了一种简单有效的GPU-Friendly Attention。它从基本的外部注意力(由方程1表示)演变而来,可以表示为:
其中, 而 GDN 表示分组双重归一化,它将原始双重归一化的第二次归一化分为 H 组,如图3左下部分所示。从方程3我们可以看出,GFA有两个主要改进。一方面,它使得矩阵乘法被集成在一起,这对类似GPU的设备非常友好。得益于这一点,我们可以将外部参数的大小从 (M, d') 扩大到(,d)。因此,可以有更多的参数进行调优以提高性能。另一方面,通过利用分组双重归一化,它在一定程度上保持了多头机制的优点。为了直观理解,可以认为GFA也生成了H个不同的注意力图来捕获tokens之间的不同关系,但是更多的特征元素参与了相似性计算,并且所有的注意力图都对最终输出做出了贡献。
跨分辨率注意力 多分辨率融合已经被证明对于密集预测任务是有效的。对于多分辨率架构的设计,我们可以直观地在不同分辨率分支独立应用GFA,并在执行卷积模块或注意力模块后交换特征。但在高分辨率分支中,像素更多关注局部信息而非高级别的全局上下文。因此,我们认为直接在高分辨率特征图上进行注意力以学习全局上下文的效果不够充分。为了更有效地获取全局上下文,我们提出了一种跨分辨率注意力,其目的是充分利用从低分辨率分支学到的高级语义信息。如图2所示,与GFA不同,跨分辨率注意力被应用于高分辨率分支以收集全局上下文。这种跨分辨率注意力的计算可以表示为:
其中,分别表示高分辨率分支和低分辨率分支上的特征图, 是一组矩阵操作,包括分割、排列和重塑,表示高分辨率分支的特征维度。值得注意的是,特征图 (在以下文本中称为跨特征)是从通过函数 计算得出的,该函数由池化和卷积层组成。特征图 的空间大小指示了从生成的token数量。实验表明,我们只在注意力图的最后一维上应用softmax进行归一化,当key和value不是外部参数时,单一softmax的表现优于双重归一化。特别是为了在类似GPU的设备上快速推理,这里同样放弃了多头机制。
前馈网络 在之前的基于Transformer的分割方法中,前馈网络(FFN)通常由两个MLP层和一个3x3深度卷积层组成,其中3x3深度卷积层用于补充位置编码或增强局部性。此外,两个MLP层将隐藏维度扩展为输入维度的两倍或四倍。这种类型的FFN可以用相对较少的参数实现更好的性能。但在考虑类似GPU设备上的延迟的情况下,典型的FFN结构并不是非常高效。为了平衡性能和效率,我们在RTFormer模块的FFN中采用了两个3x3卷积层而不进行维度扩展。结果显示,这种方式甚至比典型的FFN配置更好。
3.2 RTFormer
图4展示了RTFormer的整体架构。
骨干架构 为了提取高分辨率特征图所需的足够局部信息,我们将卷积层与我们提出的RTFormer模块结合起来构建RTFormer。具体来说,我们让RTFormer从包含两个3x3卷积层的块开始,并用几个连续的基本残差块构成前两个阶段。然后,从第三阶段开始,我们使用双分辨率模块来实现高分辨率和低分辨率分支之间的特征交换。对于最后三个阶段的高分辨率分支,特征步幅保持不变为8;而对于低分辨率分支,特征步幅分别为16、32、32。特别地,我们将双分辨率模块以阶梯式布局排列,以便借助低分辨率分支的输出来加强高分辨率特征的语义表示。最重要的是,我们用图2中所示的RTFormer模块构建第四阶段和第五阶段,以实现高效的全局上下文建模,而第三阶段仍由基本残差块组成。
分割头 对于RTFormer的分割头,我们在低分辨率输出特征之后添加了一个DAPPM模块。在融合了DAPPM的输出与高分辨率特征之后,我们获得了步幅为8的输出特征图。最后,这个输出特征被传递给像素级别的分类头以预测密集语义标签。分类头由一个3x3卷积层和一个1x1卷积层组成,其中隐藏特征维度与输入特征维度相同。
实例化 我们通过RTFormer-Slim和RTFormer-Base实例化RTFormer的架构,详细的配置记录在表1中。对于通道和块,每个组包含5个元素,分别对应5个阶段。特别地,含有两个数字的元素对应于双分辨率阶段。例如,64/128意味着高分辨率分支的通道数为64,低分辨率分支的通道数为128。而1/2意味着高分辨率分支的基本卷积块数量为1,低分辨率分支的基本卷积块数量为2。值得注意的是,最后两个元素表示RTFormer模块的数量,对于RTFormer-Slim和RTFormer-Base而言,这两个数字都是1。跨特征的空间大小分别为RTFormer-Slim设置为64(8x8),RTFormer-Base设置为144(12x12)。
4 实验
本节在Cityscapes、Camvid、ADE20K和COCOStuff数据集上验证RTFormer。我们首先介绍这些数据集及其训练细节。然后,我们在Cityscapes和CamVid上将RTFormer与最先进的实时方法进行比较。此外,我们还总结了在ADE20K和COCOStuff上的更多实验,以进一步证明我们方法的通用性。最后,我们提供了在ADE20K上关于RTFormer模块内部不同设计模块的消融研究。
4.1 实现细节
在进行语义分割微调之前,所有模型都在ImageNet上进行了预训练。关于ImageNet的训练细节将在补充材料中提供。我们使用mIoU作为性能指标,FPS作为效率指标,并默认情况下在RTX 2080Ti上测量FPS(不使用tensorrt加速)。
4.2 在不同数据集上和其他方法的对比
4.3 消融实验
5 结论
本文提出了RTFormer,它能够高效地捕捉全局上下文以提高实时语义分割的性能。广泛的实验表明,我们的方法不仅在常见的实时分割数据集上达到了新的最佳结果,而且在更具挑战性的通用语义分割数据集上也表现出优越的性能。由于RTFormer的高效性,我们希望我们的方法能够鼓励使用Transformer设计新的实时语义分割方案。一个局限性是,虽然我们的RTFormer-Slim只有4.8M参数,但在边缘设备的芯片上可能还需要更高的参数效率。我们留待未来的研究解决这个问题。
—END—
论文链接:https://arxiv.org/pdf/2210.07124v1
请长按或扫描二维码关注本公众号
喜欢的话,请给我个在看吧!