SCTNet:使用了Transformer语义信息的单分支CNN实时分割方法

点击上方“AI公园”,关注公众号,选择加“星标“或“置顶”


导读

来自美团,SCTNet既享受到了Transformer语义分支提供的丰富语义表示,又保持了轻量级单分支CNN的高效率。

f28b6b35ae6f1e4c46941b0457dfb05a.png

摘要

最近的实时语义分割方法通常采用额外的语义分支来追求丰富的长距离上下文。然而,这个额外的分支会导致不必要的计算开销,并减慢推理速度。为了解决这一困境,我们提出了SCTNet,这是一种具有Transformer语义信息的单分支卷积神经网络(CNN),适用于实时分割。SCTNet既享受到了无推理语义分支提供的丰富语义表示,又保持了轻量级单分支CNN的高效率。SCTNet利用Transformer作为训练专用的语义分支,因为Transformer具有出色的提取长距离上下文的能力。借助提出的类似Transformer的CNN块CFBlock和语义信息对齐模块,SCTNet能够在训练过程中捕捉到Transformer分支中的丰富语义信息。在推理阶段,只需部署单分支CNN。我们在Cityscapes、ADE20K和COCO-Stuff-10K数据集上进行了广泛的实验,结果表明我们的方法达到了新的最先进性能水平。代码和模型可在https://github.com/xzz777/SCTNet获取。

392b897a263c5cfb4c799502190a0a3f.png

1、介绍

作为计算机视觉中的一个基础任务,语义分割旨在为输入图像中的每个像素分配一个语义类别标签。它在自动驾驶、医学图像处理、移动应用以及许多其他领域发挥着重要作用。为了实现更好的分割性能,近期的语义分割方法追求丰富的长距离上下文。已提出了多种方法来捕获和编码丰富的上下文信息,包括大感受野、多尺度特征融合、自注意力机制等。其中,自注意力机制作为Transformer的一个核心组成部分,已被证明具有显著的建模长距离上下文的能力。尽管这些工作取得了显著的进步,但它们通常导致较高的计算成本。值得注意的是,基于自注意力的工作甚至具有与图像分辨率平方级别的计算复杂度,这在处理高分辨率图像时会显著增加延迟。这些限制阻碍了它们在实时语义分割中的应用。

许多近期的实时方法采用了双边架构来快速提取高质量的语义信息。BiSeNet提出了一种双边网络,用于在早期阶段分离详细的空域特征和丰富的上下文信息,并并行处理它们,如图 2(a) 所示。继 BiSeNet之后,BiSeNetV2和 STDC进一步加强了提取丰富长距离上下文的能力或减少了空间分支的计算成本。为了平衡推理速度和准确性,DDRNet、RTFormer和 SeaFormer采用了特征共享架构,在深层阶段划分空域和上下文特征,如图 2(b) 所示。然而,这些方法在两个分支之间引入了密集融合模块以增强提取特征的语义信息。总之,所有这些双边方法由于额外的分支和多个融合模块的存在,都面临着有限的推理速度和高计算成本的问题。

13ff459840916f95ecd0d36c8afaf04b.png

为了解决上述困境,我们提出了一种具有Transformer语义信息的单分支卷积神经网络(CNN),用于实时分割(SCTNet)。它可以高效地提取语义信息,而不会因双边网络造成大量的计算负担。具体来说,SCTNet从训练专用的Transformer语义分支向CNN分支学习长距离上下文。为了减轻Transformer与CNN之间的语义差距,我们精心设计了一个类似Transformer的CNN块,称为CFBlock,并在对齐之前使用了一个共享解码器头。通过训练过程中的语义信息对齐,单分支CNN可以同时编码语义信息和空间细节。因此,SCTNet可以在维持轻量级单分支CNN架构的高效率的同时,从Transformer架构的大有效感受野中对齐语义表示。总体架构如图 2(c) 所示。在三个具有挑战性的数据集上的广泛实验结果表明,所提出的SCTNet在准确性和速度之间有着比以前的工作更好的折衷。图 1 直观地展示了SCTNet与其他实时分割方法在Cityscapes验证集上的比较。

所提出的SCTNet的主要贡献可以总结为以下三个方面:

  • 我们提出了一种新颖的单分支实时分割网络,称为SCTNet。通过从Transformer到CNN利用语义信息对齐来学习提取丰富的语义信息,SCTNet既享有Transformer的高准确性,又保持了轻量级单分支CNN的快速推理速度。

  • 为了缓解CNN特征与Transformer特征之间的语义差距,我们设计了CFBlock(ConvFormer块),它仅使用卷积操作就能像Transformer块一样捕获长距离上下文。此外,我们提出了SIAM(语义信息对齐模块),以更有效的方式对齐特征。

  • 广泛的实验结果表明,所提出的SCTNet在Cityscapes、ADE20K和COCO-Stuff-10K数据集上的实时语义分割方面超越了现有的最先进方法。SCTNet为实时语义分割的速度提升和性能改进提供了新的视角。

2、相关工作

语义分割。FCN引领了利用卷积神经网络(CNN)进行语义分割的趋势。继FCN之后,一系列改进的基于CNN的语义分割方法被提出。DeepLab通过空洞卷积扩大了感受野。PSPNet、U-Net以及 RefineNet通过融合不同层次的特征表示来捕获多尺度上下文。一些方法提出了各种注意力模块来提高分割性能。近年来,Transformer被应用于语义分割,并显示出有前景的表现。SETR首次直接将视觉Transformer应用于图像分割。SegViT尝试使用纯视觉Transformer进行语义分割。PVT将典型的层级结构从CNN引入到基于Transformer的语义分割模型中。SegFormer提出了一种高效的基于多尺度Transformer的分割模型。

实时语义分割。早期的实时语义分割方法通常通过压缩通道或快速下采样来加速推理。ICNet首次引入了多分辨率图像级联网络以加快速度。BiSeNetV1和 BiSeNetV2采用了双分支架构和特征融合模块,以实现速度和准确性之间的更好折衷。STDC重新考虑了BiSeNet的双分支网络,去除了空间分支,并添加了一个详细指导模块。DDRNets通过在早期阶段共享分支实现了更好的折衷。非常近的一些针对实时分割的有效Transformer方法已经被提出,但它们仍然存在未解决的问题。TopFormer仅在特征图的1/64尺度上使用Transformer,导致准确性较低。RTFormer和 SeaFormer需要两个分支之间频繁的交互。这种额外的计算会减缓推理速度。也有一些单分支或多分支的方法。

注意力机制。近年来,注意力机制在计算机视觉领域得到了广泛应用。许多方法致力于减少注意力机制的计算复杂度。虽然其中一些方法已经实现了线性复杂度,但它们包含频繁的移位或重塑操作,这带来了大量的延迟。MSCA显示了有前景的表现,但是较大的核尺寸不太易于使用,并且注意力的多尺度设计进一步增加了推理速度的损耗。外部注意力形式非常简单。它使用外部参数作为键和值,并通过两个线性层实现注意力机制。GFA(GPU-Friendly Attention)通过用组双重规范化替换EA中的头部分割来改进外部注意力,这对于GPU设备更加友好。

3、方法

0ed363700fd57a8e2bbb75feccab9709.png

动机

去除双边网络中的语义分支可以显著加快推理速度。然而,这会导致浅层的单分支网络缺乏长距离的语义信息,从而导致准确性降低。虽然使用深度编码器和强大的解码器或复杂的增强模块可以恢复准确性,但这会减慢推理过程。为了解决这个问题,我们提出了一种仅在训练阶段使用的对齐方法,这种方法可以丰富语义信息而不牺牲推理速度。具体来说,我们提出了SCTNet,这是一种具有训练专用语义提取变换器的单分支卷积网络,它拥有Transformer的高准确性以及CNN的快速推理速度。SCTNet的概览如图 3 所示。

Conv-Former 块

作为不同类型的网络,CNN 和Transformer提取的特征表示显著不同。直接对齐CNN和Transformer之间的特征会使学习过程变得困难,从而导致性能提升有限。为了让CNN分支更容易地学会如何从Transformer分支提取高质量的语义信息,我们设计了Conv-Former 块。Conv-Former 块尽可能模拟Transformer块的结构,以便更好地学习Transformer分支的语义信息。同时,Conv-Former 块仅使用高效的卷积操作来实现注意力功能。

25bff1454ea094374527bca281b5a586.png

Conv-Former 块的结构类似于典型的Transformer编码器,如图 4 的左侧所示。其过程可以描述如下:

dd6333f86fe6dd524d09565f33d89421.png

卷积注意力。用于实时分割的注意力机制应具有低延迟和强大的语义提取能力。正如相关工作中讨论的那样,我们认为GFA是一个潜在的选择。我们的卷积注意力源自GFA。GFA与所提出的卷积注意力之间有两个主要区别。首先,我们将GFA中的矩阵乘法替换为逐像素的卷积操作。点卷积相当于逐像素的乘法,但没有特征展平和重塑操作。这些操作不利于保持固有的空间结构,并会带来额外的推理延迟。

此外,卷积提供了一种更灵活的方式来扩展外部参数。然后,由于Transformer与CNN之间的语义差距,仅仅通过计算几个可学习向量与每个像素之间的相似性,并根据相似性图和可学习向量增强像素,不足以捕获丰富的上下文。为了更好地对齐Transformer的语义信息,我们将可学习向量扩大为可学习的核。一方面,这将像素与可学习向量之间的相似性计算转换为带有可学习核的像素块之间的相似性计算。另一方面,使用可学习核的卷积操作在一定程度上保留了更多的局部空间信息。卷积注意力的操作可以总结如下:

8f8237411d48e0d3093691ba6d6fb793.png

考虑到效率问题,我们使用条带卷积而不是标准卷积来实现卷积注意力。更具体地说,我们利用一个 1×𝑘 和一个 𝑘×1卷积来近似一个 𝑘×𝑘 卷积层。图 4 展示了卷积注意力的实现细节。

前馈网络(Feed Forward Network)。典型的前馈网络在提供位置编码和嵌入通道方面起着至关重要的作用。近期Transformer模型中的典型前馈网络由一个扩展点卷积、一个 3x3 深度可分离卷积和一个压缩点卷积组成。不同于典型的前馈网络,我们的前馈网络由两个标准的  3x3 卷积层组成。与典型的前馈网络相比,我们的前馈网络更高效,并提供了更大的感受野。

语义信息对齐模块

我们提出了一种简单而有效的对齐模块来进行训练过程中的特征学习,如图 3 所示。它可以分为骨干特征对齐和共享解码器头对齐。

骨干特征对齐。得益于Conv-Former块的类似Transformer的结构,对齐损失可以轻松地将Conv-Former块的特征与Transformer的特征对齐。简而言之,骨干特征对齐首先对来自Transformer和CNN分支的特征进行下采样或上采样以进行对齐。然后将CNN的特征投影到Transformer的维度。该投影可以:1) 统一通道数量;2) 避免直接对齐特征,这样做会损害训练过程中CNN的真实监督。最后,对投影后的特征应用语义对齐损失以对齐语义表示。

共享解码器头对齐。Transformer解码器通常使用多个阶段的特征进行复杂的解码,而SCTNet解码器仅选择第2阶段和第4阶段的特征。考虑到它们在解码空间上的显著差异,直接对解码特征和输出logits进行对齐只能获得有限的改进。因此,我们提出了共享解码器头对齐。具体来说,将单分支CNN的第2阶段和第4阶段特征进行拼接后输入到一个点卷积以扩展维度。然后将高维特征通过变换器解码器。使用Transformer解码器的新输出特征和logits与变换器解码器的原始输出计算对齐损失。

总体架构

为了减少计算成本同时获得丰富的语义信息,我们将流行的双分支架构简化为一个快速的CNN分支用于推理,以及一个仅用于训练时语义对齐的Transformer分支。

骨干网络。为了提高推理速度,SCTNet采用了典型的层级CNN骨干网络。SCTNet从一个由两个连续的3x3卷积层组成的stem块开始。前两个阶段由堆叠的残差块组成,而后两个阶段包括所提出的类似Transformer的块,称为Conv-Former 块(CFBlocks)。CFBlock运用了几种精心设计的卷积操作来执行与Transformer块类似的长距离上下文捕获功能。我们在第2至第4阶段的开始处应用了一个由步幅卷积、批量归一化和ReLU激活组成的convdown层来进行下采样,为了清晰起见,图 3 中省略了这一点。

解码器头。解码器头由一个DAPPM和一个分割头组成。为了进一步丰富上下文信息,我们在第4阶段的输出之后添加了一个DAPPM。然后将该输出与第2阶段的特征图进行拼接。最后,将这个输出特征传递给分割头。具体来说,分割头由一个3x3 Conv-BN-ReLU算子后跟一个1x1卷积分类器组成。

训练阶段。众所周知,Transformer在捕获全局语义上下文方面表现出色。另一方面,CNN已经被广泛证明在建模层级局部信息方面优于Transformer。受到Transformer和CNN优点的启发,我们探索了为实时分割网络配备这两种优势的方法。我们提出了一种单分支CNN,它学习将其特征与强大的Transformer的特征对齐,如图 3 中蓝色虚线框所示。这种特征对齐使得单分支CNN既能提取丰富的全局上下文,又能提取详细的空域信息。具体来说,在训练阶段中有两个流。SCTNet采用一个仅用于训练的Transformer作为语义分支来提取强大的全局语义上下文。语义信息对齐模块监督卷积分支以对齐来自Transformer的高质量全局上下文。

推理阶段。为了避免双分支带来的大量计算成本,在推理阶段只部署CNN分支。借助与Transformer对齐的语义信息,单分支CNN可以在没有额外的语义提取或昂贵的密集融合的情况下生成准确的分割结果。具体来说,输入图像被送入单分支层级卷积骨干网络。然后解码器头选取骨干网络中的特征,并进行简单的拼接后跟随逐像素分类。

对齐损失

为了更好地对齐语义信息,需要一种专注于语义信息而非空间信息的对齐损失。在实现中,我们使用CWD Loss(通道级蒸馏损失)(Shu 等人,2021)作为对齐损失,它比其他损失函数显示出更好的结果。CWD Loss可以总结如下:

f0b8813e3ab1a660f0c5511f4b2c49ce.png

其中 𝑐=1,2,…,𝐶表示通道索引,𝑖=1,2,…,𝐻⋅𝑊表示空间位置,和 分别是Transformer分支和CNN分支的特征图。𝜙将特征激活转化为通道级概率分布,消除了Transformer和紧凑型CNN之间的尺度影响。为了最小化 ,当较大时,应该也较大。但当较小时,的值无关紧要。这迫使CNN学习前景显著性的分布,其中包含了语义信息。𝑇表示一个超参数,称为温度。𝑇越大,概率分布越软。

4、实验

数据集和实施细节

我们在三个数据集上进行了SCTNet的实验,即Cityscapes、ADE20K和COCO-Stuff-10K,以证明我们方法的有效性。为了公平比较,我们构建了基准模型SCTNet-B,其规模与RTFormer-B/DDRNet-23/STDC2相当。此外,我们也介绍了一个较小的变体,称为SCTNet-S。我们首先在ImageNet上预训练我们的CNN骨干网络,然后在语义分割数据集上对其进行微调。在训练阶段,语义Transformer分支可以是任何层级Transformer网络。在我们的实现中,我们选择了SegFormer作为所有实验中的Transformer分支。我们在单个NVIDIA RTX 3090上测量所有方法的推理速度。除非另有说明,所有报告的FPS结果都是在相同的输入分辨率下获得的,以保证公平的性能比较。对于Cityscapes,我们使用torch和tensor-RT两种方式来测量速度。

和其他方法的对比

961efeba3f89e435f474d1b887bc9f38.png

消融实验

不同类型blocks的对比:

2c111c0d88750f950af7cc6c4f2a6122.png

语义信息对齐模块有效性:

ee7c917dc1c6cc8d836582a1ead72ea4.png

各个组件的有效性:

6d72f571943beba0814cbcf97e06e9ed.png

结果可视化:

c6725588a0cddae5833c801631a2e46e.png

总结

在这篇论文中,我们提出了SCTNet,这是一种新颖的单分支架构,能够在不增加额外推理计算成本的情况下提取高质量的长距离上下文。广泛的实验表明,SCTNet取得了新的前沿成果。此外,通过展示SCTNet的高效性,我们为双边分支网络中的语义分支提供了一种新的见解,并提供了一种新的途径来推动实时分割社区的发展,不仅采用了Transformer的结构,还利用了它的知识。

4b875e2ad5617788483cda38930a2cfe.png

—END—

论文链接:https://arxiv.org/pdf/2312.17071v2

f24d74e93627116cf9524c4e1241c5a8.jpeg

请长按或扫描二维码关注本公众号

喜欢的话,请给我个在看吧

评论
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值