LETNet：使用高效Transformer和CNN的轻量级实时语义分割网络

导读

LETNet将U形CNN与Transformer有效地结合在一起，以胶囊嵌入的方式弥补彼此的不足。同时，精心设计的轻量级空洞bottleneck（LDB）模块和特征增强（FE）模块对从头开始的训练产生了积极影响。在具有挑战性的数据集上进行的大量实验表明，LETNet在准确性和效率之间的平衡方面表现出色。

摘要

在过去十年里，卷积神经网络（CNNs）在语义分割任务中展现出了重要地位。尽管CNN模型表现出非常优秀的性能，但捕捉全局表示的能力仍然不足，导致结果不尽如人意。最近，Transformer在自然语言处理（NLP）任务中取得了巨大成功，展现了其在建模长距离依赖关系方面的优势。最近，Transformer也开始受到计算机视觉研究者的广泛关注，他们将图像处理任务重新定义为序列到序列预测问题，但这种方法导致局部特征细节的丢失。在本工作中，我们提出了一种轻量级实时语义分割网络，称为LETNet。LETNet将U形CNN与Transformer有效地结合在一起，以胶囊嵌入的方式弥补彼此的不足。同时，精心设计的轻量级空洞bottleneck（LDB）模块和特征增强（FE）模块对从头开始的训练产生了积极影响。在具有挑战性的数据集上进行的大量实验表明，LETNet在准确性和效率之间的平衡方面表现出色。具体而言，它仅包含0.95M参数和13.6G FLOPs，但在Cityscapes测试集上达到了120 FPS下的72.8% mIoU，在CamVid测试数据集上达到了250 FPS下的70.5% mIoU，使用单个RTX 3090 GPU。源代码在https://github.com/IVIPLab/LETNet。

1、介绍

语义分割的任务旨在为每个像素分配一个语义标签，这一任务广泛应用于增强现实设备、自动驾驶和视频监控等领域。自从全卷积网络（FCN）被提出以来，现有的语义分割模型一直将其作为一种原型来进行改进。然而，基于FCN的模型的感受野有限。因此，它们无法学习远程依赖关系，这对提取全局语义信息不利，而这对于密集型任务尤其重要，特别是语义分割任务。为了解决这一局限性，一些最近的方法提出了使用大型卷积核、空洞卷积和特征金字塔来扩大感受野。另一种方法是在FCN结构中集成Non-local，这是从自然语言处理（NLP）领域借鉴过来的，旨在建模特征图中所有像素的全局交互作用，但代价是较高的内存消耗和计算成本。另一方面，研究人员开始尝试完全去除卷积，探索只使用注意力模块的模型，即Transformer，该模型设计用于建模序列到序列的长程依赖关系，并捕捉序列中任意位置的关系。

与之前的基于CNN的方法不同，Transformer 不仅在全局上下文建模方面强大，而且在大规模预训练的情况下也能在下游任务上取得良好的结果。视觉Transformer（ViT），它使用带有位置嵌入的二维图像块作为输入来执行图像识别任务。然而，与CNN相比，ViT的缺点是必须在大数据集上进行预训练。图像分辨率远大于NLP中的词汇。许多计算机视觉任务，如语义分割，要求在像素级别进行密集预测，这对于高分辨率图像上的Transformers来说很难处理，因为其自注意力的计算复杂度与图像尺寸呈二次方关系。此外，当Transformer应用于图像处理领域时，二维图像被切片并作为一维序列输入到模型中，从而打破了局部结构之间的联系，并在所有阶段仅关注全局上下文。因此，低分辨率特征缺乏详细的定位信息，无法通过直接上采样至全分辨率来有效恢复，导致分割结果粗糙。

尽管Transformer能够实现全局信息建模，但它无法提取精细的空间细节。相反，CNN可以提供一种方式来提取低级视觉线索，很好地补偿了这种精细的空间细节。因此，一些方法试图将CNN与Transformer结合起来处理语义分割任务。例如，在医学图像分割领域，TransUNet、TransBTS和TransFuse已经取得了令人满意的结果。受此启发，我们也提出了一种轻量级实时语义分割模型，名为LETNet，它是基于CNN和Transformer的。如图1所示，我们的LETNet在模型性能、模型大小和推理速度之间实现了很好的平衡。本文的主要贡献有三个方面：

• 我们提出了一种轻量级空洞bottleneck（LDB）来提取重要的语义信息。LDB由空洞卷积和深度可分离卷积组成，在参数和计算量方面都大大减少。

• 我们提出了一种混合网络，即LETNet，用于语义分割。LETNet采用了最简洁的编码器-解码器结构，并将高效的Transformer视为胶囊网络来学习全局信息。同时，在跳跃连接中添加了一个特征增强（FM）模块，以帮助在恢复分辨率时补充边界细节信息。

• LETNet在单一RTX3090硬件平台上仅使用0.95M的参数量就在Cityscapes测试集上达到了72.8%的mIoU，在CamVid数据集上也获得了70.5%的良好性能。这些性能优于大多数现有模型。

2、相关工作

A. 基于CNN的语义分割方法

得益于卷积神经网络强大的特征表示能力，语义分割方法也取得了显著的进步。基于CNN的开创性文章是FCN，之后许多架构都在此基础上进行了改进。为了缓解图像分辨率与有限感受野之间的矛盾，DeepLab和 PSPNet使用了并行空洞卷积来构建空洞空间金字塔池化（ASPP）模块，该模块引入了对各种尺度上下文信息的良好描述符。此外，利用建模特征依赖性的优点，自注意力机制吸引了许多学者的兴趣。例如，基于SENet，ECANet提出了一种无需降维的局部跨通道交互策略以及一种自适应选择一维卷积核大小的方法。此外，还有一些在NLP [6]中常用的关注机制来建模长距离依赖关系。其中典型的例子是非局部神经网络，它利用两点之间的相似性来加权每个位置的特征。DANet使用ResNet作为主干网络，随后是一个由空间维度和平行的通道维度组成的注意力模块，用于捕获长距离的深层特征依赖关系，以提高分割结果的质量。CCNet进一步改进了计算像素与其所在行和列中所有像素之间的关联，从而减轻了计算负担。LRNNet提出了一个有效的简化非局部模块，该模块使用区域奇异向量来生成更简化的代表性特征，用于建模远程依赖关系和进行全局特征选择。

尽管这些方法取得了不错的效果，但它们并没有改变非局部本质上是一种像素级别的矩阵算法的事实，这仍然给计算机带来了巨大的计算挑战。因此，轻量级网络应运而生。例如，ICNet使用多尺度图像作为输入，利用高层语义信息和底层的空间细节。BiseNet和BiseNet-v2提出了双路径架构，一条分支负责提取深层语义信息，另一条高分辨率的浅层分支负责提供详细的补充信息。DFANet利用了特征复用策略，增强了不同层级特征之间的交互和聚合。此外，每个阶段的末尾使用逐点注意力来增强特征表达能力，同时确保计算量较小。ESPNet和ESPNet-v2通过将分解卷积集成到逐点卷积和空洞卷积中，减少了参数数量和计算量。此外，NRD使用动态卷积神经网络从图像中提取特征信息。然而，基于CNN的方法总有一个难以完全解决的问题，那就是建模长距离关系的局限性。虽然现有方法只是依赖于构建深编码器和下采样操作，但负面影响是冗余的参数和更多局部细节的损失。

B. 基于Transformer的语义分割方法

Transformer最初在中被提出，并在自然语言处理领域取得了巨大的成功。与CNN不同，Transformer不仅在全球上下文建模方面强大，而且在大规模预训练的情况下在下游任务上也取得了良好的效果。例如，ViT提出使用带有位置嵌入的2D图像块作为输入来执行图像识别任务。DETR及其变形版本利用Transformer的编码器-解码器结构来融合检测头部的上下文信息。SETR从序列到序列的角度解决了语义分割问题。它放弃了CNN，是一种完全基于Transformer的结构。同时，SegFormer提出了分层的编码器结构，输出多尺度特征，并在解码器中融合这些特征。之后，一些变形结构被提出用于医学图像分割。

尽管上述方法取得了不错的效果，但由于Transformer的计算复杂度与图像尺寸的平方成比例，这会增加大量的计算负担。此外，当在图像域中使用Transformer时，输入的图像块被视为一维序列输入到模型中，这破坏了局部结构之间的联系，并且只关注阶段的全局上下文建模。缺乏详细的定位信息导致分割结果粗糙。另一方面，CNN架构提供了一条路径来提取低级视觉线索，这些线索可以很好地补偿这种精细的空间细节。因此，我们的目标是探索一种结合CNN与Transformer的方法来处理分割任务，并使用高效的Transformer来降低内存消耗。

3、方法

A. 网络架构

如图2所示，LETNet包括一个编码器（Encoder）、一个解码器（Decoder）、一个高效Transformer以及三个长距离跳跃连接。具体来说，编码器和解码器是CNN结构，用于提取局部特征以获得更好的图像表示。Transformer通过自注意力和多层感知机（MLP）结构反映复杂的空间变换和长距离特征依赖关系，以获得全局特征表示。这三个长距离连接受到了UNet的启发，将低级空间信息与高级语义信息相结合，以实现高质量的分割。

B. 轻量级空洞bottleneck（LDB）

如图3所示，LDB的整体结构采用了ResNet的思想，并将模块设计为残差模块，以在尽可能减少网络层数的同时收集更多的特征信息。具体来说，在bottleneck处，通过1×1卷积将输入特征的通道数减少一半。减少通道数后，参数量和计算量大大减少。虽然这会损失一部分准确性，但堆叠两个这样的模块将更有益于弥补这一点的损失。同时，由于使用了1×1卷积，网络深度必须加深以获得更大的感受野。因此，在1×1卷积之后，添加了3×1和1×3的分解卷积，以扩展感受野来捕捉更广泛的上下文信息。此外，分解卷积也是基于考虑参数数量和计算量的。类似地，在接下来的双分支结构中，两个分支也都使用分解卷积，其中一个负责局部和短距离的特征信息，另一个使用空洞卷积，负责在不同的空洞率下从更大的感受野中提取特征信息。在这两个分支旁边是通道注意力，灵感来源于ECANet，旨在在通道维度建立注意力矩阵以增强特征表达并抑制噪声干扰，因为对于CNN来说，大部分的特征信息都包含在通道中。然后，将两个低维分支和中间特征融合，并输入到下面的1×1逐点卷积中，以恢复特征图的通道数与输入特征图的通道数相同。最后，使用通道混洗的策略来避免由深度卷积导致的信息独立和通道间无相关性的缺点，并促进不同通道之间语义信息的交换。完整的操作如下：

C. 高效Transformer (ET)

正如之前所提到的，尽管CNN在局部特征提取方面具有优势，但其捕捉全局表示的能力仍然不足，这对于许多高层次的计算机视觉任务非常重要。为了解决这个问题，我们引入Transformer来学习长距离依赖关系。然而，在图像处理任务中，由于输入图像的分辨率远大于自然语言处理领域的词汇。我们引入了高效Transformer (ET)，灵感来源于ETSR。与传统的Transformer不同，ET占用较少的计算资源。同时，为了避免过高的内存使用和计算负载，我们放弃了ETSR中多个ET的串联连接，而是仅使用一个ET作为胶囊网络，放置在整个网络的中间。众所周知，Transformer由两层归一化、一个多头注意力（MHA）和一个多层感知机（MLP）组成。ET与原始Transformer最大的区别在于MHA。在层归一化之后，ET设置了一个缩减层来将通道数减半，从而减少了部分计算。然后，一个线性层将特征图投影成三个矩阵，Q（查询）、K（键）和V（值）。具体来说，在EMHA中，Q、K和V首先被分成s段，然后对Qi、Ki和Vi执行缩放点积注意力。之后，我们将得到的Q1...Qs拼接起来得到整个输出Q。实际上，这依赖于组卷积的思想，将大矩阵分割成小矩阵，然后进行计算并最终合并，以达到减少计算量的目的。最后，使用扩展层来恢复通道数。EMHA的架构如图4(a)所示，缩放点积注意力（SDPA）的操作可以定义为：

D. 特征增强 (FE)

在神经网络中，较低层具有高分辨率和精确的空间信息（分辨率对应于空间位置），但语义信息较少。相比之下，较高层具有低分辨率，缺乏空间位置信息，但语义信息丰富。因此，在分割任务中，为了使高层信息也具有足够的空间信息，通常会将低层的空间信息与高层的语义信息相结合来实现高质量的分割。因此，我们采用UNet式的结构来融合相同分辨率的高低层特征图。同时，在三个长连接的过程中，我们提出了一种特征增强（FE）模块来提升特征表达能力。如图4(b)所示，特征依赖建模从两个维度进行，一个是通道维度，另一个是空间维度，并且这两个维度同时进行变换，最后融合以更有效地将低层信息传输到高层。操作可以定义为：

E. 像素注意力 (PA)

像素注意力（PA）机制根据不同像素位置特征的重要性学习权重。这意味着每个通道具有相同的权重，但对于同一通道，在不同像素位置的权重是不同的。像素注意力更加关注图像中物体的边缘和纹理，因此添加PA有助于恢复边缘细节信息，从而提高分割性能。操作如图5所示，公式如下：

4、实验

消融实验

如表II所示，我们设计了一些关于提出的模块的消融实验来证明这些模块的有效性。具体分为三组实验：A. 长连接，B. 特征增强，C. 高效Transformer。值得注意的是，基线模型没有引入任何提出的模块。基线模型的架构由LDB组成的编码器-解码器构成，在验证集上实现了69.85%的mIoU。

在组A中，我们逐渐加入L1、L2、L3来观察效果。可以看到，在加入L1之后，性能有了显著的提高，达到了0.58%，这证明了浅层信息对于深层语义信息的恢复是非常有益的。同时，当在模型中加入三个长跳跃连接时，性能增加了1.16%的mIoU。这进一步验证了长连接对于图像分割的重要性。

在组B中，我们在加入特征增强之后逐渐加入L1、L2、L3。比较B1和A1，可以看到加入FE仅增加了101K参数和0.0032G FLOPs，但却能够提升模型性能0.2%的mIoU，这是一个相当显著的提升。

在组C中，引入高效Transformer后，模型的性能提高了3.07%的mIoU，引入像素注意力机制也给模型带来了0.5%的提升。最终，C3*是提出的LETNet的最终版本，相对于基线模型，性能提高了5.15%的mIoU。

以上所有实验充分证明了提出的模块和策略的必要性和有效性。

和其他方法的对比

5、总结

在本文中，我们提出了一种结合Transformer和CNN的轻量级实时语义分割网络。我们结合了CNN的局部特征提取能力和Transformer的长距离依赖建模能力。具体而言，在模型中间引入了一个高效Transformer作为胶囊网络。与传统Transformer不同，我们使用了一个更为轻量级的多头注意力（MHA），可以显著减少GPU内存消耗。同时，设计的轻量级空洞bottleneck（LDB）模块可以在确保极端简单和轻量化的基础上学习更多特征。此外，为了弥补CNN在提取深层语义信息时丢失的浅层细节信息，我们在模型中使用了U形连接。在连接不同层次时，还设计了一个特征增强（FM）模块以改善有效特征表达并抑制噪声。广泛的实验结果显示，我们的模型在模型大小和性能之间取得了优秀的平衡。

—END—

论文链接：https://arxiv.org/pdf/2302.10484

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

喜欢的话，请给我个在看吧！