最近,基于
Transformer
的网络在语义分割方面取得了令人印象深刻的成果。然而,对于实时语义分割,由于Transformer
的计算机制非常耗时,纯CNN
方法仍然是该领域的主流。本文提出了
RTFormer
,一种用于实时语义分割的高效双分辨率Transformer
,与基于CNN的模型相比,它在性能和效率之间取得了更好的平衡。为了在类似GPU
的设备上实现高推理效率,RTFormer
利用具有线性复杂性的GPU Friendly Attention
,并放弃了多头机制。此外,作者发现,通过传播从低分辨率分支学到的高层次知识,跨分辨率注意力更有效地为高分辨率分支收集全部上下文信息。
对主流基准测试的大量实验证明了提出的
RTFormer
的有效性,它在Cityscapes
、CamVid
和COCOStuff
上达到了最先进的水平,并在ADE20K
上显示出良好的结果。
1、简介
语义分割是一项基本的计算机视觉任务,通常作为自主驾驶、移动应用、机器人感知等领域的关键感知模块。随着这些应用的发展,实时语义分割的需求越来越强烈。现有的实时分割方法主要侧重于利用CNN
架构,包括通过手工设计高效主干网络和解码器,以及探索神经架构搜索方法,以在准确性和效率之间找到更好的平衡。到目前为止,这些优秀的作品已经取得了显著的进步。
近年来,视觉Transformer
因其强大的视觉识别能力而备受关注。在继承它们的基础上,提出了一系列基于Transformer
的体系结构,并在一般的语义分割任务中表现出很好的性能。与基于CNN
的网络相比,这些基于Transformer
的架构的主要区别在于大量使用自注意力,而自注意力擅长捕获远程上下文信息,这在语义分割中是必不可少的。因此,作者认为注意力结构在实时语义分割任务中也应该有效。
但到目前为止,只有少数作品探索了注意力在这一领域的应用,目前的SOTA
水平仍然由基于CNN
的架构主导。作者认为,在实时环境中应用注意力的主要障碍可能来自以下两个方面。
- 一是现有的大多数注意力类型的计算特性对于类似
GPU
的设备来说都不是很容易推理的,例如二次复杂度和多头机制。二次复杂度在处理高分辨率特征时引入了巨大的计算负担,特别是在语义分割等密集预测任务中。尽管有几部作品缩小了key
和value
的大小,但二次复杂性的特性仍然存在。而多头机制将矩阵乘法拆分为多个组,这使得注意力操作在类似GPU
的设备上耗费时间,类似于执行组卷积的情况。 - 另一个原因是,仅关注高分辨率特征图本身可能不是用高级别语义信息捕获远程上下文的最有效方法,因为来自高分辨率特征图的单个特征向量的感受野有限。
本文提出了一种新的Transformer
块,名为RTFormer
块,如图2所示,其目的是在具有Transformer
的类似GPU
的设备上实现性能和效率之间的更好平衡。对于低分辨率分支,采用了新提出的GPU Friendly Attention
,该注意力源自External Attention
。它继承了External Attention
的线性复杂性特性,并通过丢弃矩阵乘法操作中的通道分裂,缓解了类GPU
设备的多头机制的弱点。相反,它扩大了外部参数的数量,并将External Attention
提出的双范数运算中的第二个归一化分解为多个组。这使得GPU Friendly Attention
能够在一定程度上保持多头机制的优势。对于高分辨率分支,采用交叉分辨率注意力,而不是仅在高分辨率特征本身内进行注意力。此外,与多分辨率融合的并行公式不同,作者将两个分辨率分支排列成阶梯布局。因此,通过辅助从低分辨率分支中学习到的高级别全局上下文信息,可以更有效地增强高分辨率分支。
基于提出的RTFormer
块,构建了一个新的实时语义分割网络,称为RTFormer
。为了了解足够的局部上下文,RTFormer
仍然在早期阶段使用卷积块,并在最后两个阶段放置RTFormer
块。通过大量的实验发现RTFormer
可以更有效地利用全局上下文,并比以前的工作取得更好的折衷效果。图1显示了RTFormer
与CamVid
上其他方法的比较。最后,我们将RTFormer
的贡献总结为以下三个方面:
- 提出了一种新的
RTFormer
块,它在类GPU
设备上实现了语义分割任务的性能和效率之间的更好平衡。 - 提出了一种新的网络体系结构
RTFormer
,它可以充分利用全局上下文,通过深度利用注意力而不损失效率来改进语义分割。 RTFormer
在Cityscapes
、CamVid
和COCOStuff
上达到了最先进的水平,在ADE20K
上表现良好。此外,它为实时语义分割任务的实践提供了一个新的视角。
2、本文方法
2.1、RTFormer block
RTFormer
块是继承多分辨率融合范式的双分辨率模块。与之前的工作相比,RTFormer
块由两种类型的注意力及其前馈网络组成,并按阶梯式布局排列,如图2所示。在低分辨率分支中使用GPU Friendly Attention
来捕获高级全局上下文。在高分辨率分支中,引入了交叉分辨率注意力,将从低分辨率分支中学习到的高级别全局上下文广播到每个高分辨率像素,并使用阶梯式布局将低分辨率分支的更具代表性的特征引入交叉分辨率注意力。
1、GPU-Friendly Attention
通过比较现有的不同类型的注意力发现External Attention
(EA)由于其令人满意的线性复杂度特性,可以成为在类似GPU
的设备上执行的潜在选择,GPU-Friendly Attention
(GFA)就是从中派生出来的。因此,在详细说明GFA
之前,首先说明EA。让表示一个输入特征,其中是元素(或图像中的像素)的数量,是特征维数,那么EA的公式可以表示为:
其中是可学习的参数,是参数维,DN
是双重规范化操作。EA
的多头版本可以解释为:
其中,是头数,而是的第个头。如图3的上半部分所示,多头机制生成个注意力图,以提高EA的容量,这使得矩阵乘法被分成几个组,类似于组卷积。虽然EA对不同的头使用共享的和,这可以大大加快计算速度,但分裂矩阵乘法仍然存在。
为了避免类似GPU的设备由于多头机制而延迟减少,作者提出了一种简单有效的GPU-Friendly Attention
。它是由方程式1所表示的基本External Attention
演变而来的,可公式化为:
其中,GDN表示分组双重归一化,它将原始双重归一化的第二次归一化分解为组,如图3左下部分所示。从方程3中可以发现GFA有两个主要改进。
- 一方面,它可以集成矩阵乘法,这对于类似GPU的设备非常友好。得益于此,可以将外部参数的大小从扩大到。因此,可以调整更多参数以提高性能。
- 另一方面,它利用分组双重归一化在一定程度上保持了多头机制的优势。为了直观理解,可以认为
GFA
还生成了个不同的注意力图,用于捕获Token
之间的不同关系,但计算相似性需要更多的特征元素,所有注意力图都有助于最终输出。
2、Cross-resolution Attention
多分辨率融合已被证明对密集预测任务是有效的。而对于多分辨率架构的设计,可以直观地将GFA
独立应用于不同分辨率的分支,并在执行卷积模块或注意力模块后交换特征。但在高分辨率分支中,像素更关注局部信息,而非高级别的全局上下文。因此,作者认为直接关注高分辨率特征图来学习全局背景是不够有效的。为了更有效地获取全局上下文,作者提出了一种交叉分辨率注意力,旨在充分利用从低分辨率分支中学习到的高层语义信息。如图2所示,与GFA
不同,高分辨率分支采用交叉分辨率注意力来收集全局背景。交叉分辨率注意力的计算表示为:
其中分别表示高分辨率分支和低分辨率分支上的特征映射,是一组矩阵运算,包括 splitting
、permutation
以及reshaping
,表示高分辨率分支的特征维数。值得说明的是,下文中表示为交叉特征的特征图是通过由池化层和卷积层组成的函数从计算得出的。的空间大小表示从生成的Token
数。
在实验上,只在注意力图的最后一个轴上使用softmax
进行归一化,因为当key
和value
不是外部参数时,单个softmax
比双重归一化性能更好。特别是,为了在类似GPU
的设备上进行快速推理,这里也放弃了多头机制。
3、Feed Forward Network
在以前的基于Transformer
的分割方法中,前馈网络(FFN
)通常由两个MLP
层和一个深度3×3卷积层组成,其中深度3×3卷积用于补充位置编码或增强局部性。此外,两个MLP
层将隐藏维度扩展为输入维度的2到4倍。这种类型的FFN
可以用相对较少的参数获得更好的性能。但在应该考虑类GPU
设备上的延迟的情况下,FFN
的典型结构不是很有效。为了平衡性能和效率,在RTFormer
块的FFN
中采用了两个3×3卷积层,而没有进行维数扩展。与典型的FFN
配置相比,它显示了更好的结果。
2.2、RTFormer
1、Backbone Architecture
为了提取高分辨率特征图所需的足够的局部信息,将卷积层与RTFormer
块相结合来构造RTFormer
。具体来说,让RTFormer
从一个由两个3×3卷积层组成的干块开始,用几个连续的基本残差块组成前两个阶段。然后,从第3阶段开始,使用双分辨率模块,实现高分辨率和低分辨率分支之间的特征交换。对于最后三个阶段的高分辨率分支,特征stride
保持为8不变,而对于低分辨率分支,则特征stride
分别为16、32、32。特别是,将双分辨率模块安排为阶梯式布局,以借助低分辨率分支的输出增强高分辨率特征的语义表示。最重要的是,用提出的RTFormer
块构造了stage4
和stage5
,如图2所示,用于有效的全局上下文建模,而stage3
仍然由基本残差块组成。
2、Segmentation Head
对于RTFormer
的分割头,在低分辨率输出特征之后添加了一个DAPPM
模块。然后将DAPPM
的输出与高分辨率特征融合,得到stride
为8的输出特征图。最后,将该输出特征传递到像素级分类头,用于预测密集语义标签。分类头由3×3卷积层和1×1卷积层组成,隐藏特征维数与输入特征维数相同。
3、Instantiation
用RTFormer Slim
和RTFormer Base
实例化了RTFormer
的架构,详细配置记录在表1中。对于通道数和块数,每个数组包含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)。