CABiNet：用于低延迟语义分割的高效上下文聚合网络-CSDN博客

摘要

随着自主机器需求的不断增加，视觉场景理解的像素级语义分割不仅需要准确，而且需要高效，以满足任何潜在的实时应用需求。在本文中，我们提出了CABiNet（Context Aggregated Bi-lateral Network，上下文聚合双边网络），这是一种双分支卷积神经网络（CNN），与最先进的网络相比，其计算成本显著降低，同时保持了具有竞争力的预测精度。在现有的用于高速语义分割的多分支架构的基础上，我们设计了一个廉价的高分辨率分支用于有效的空间细节描述，以及一个包含轻量级全局聚合和局部分布模块的上下文分支，该分支能够以较低的计算开销捕获精确语义分割所需的远程和局部上下文依赖。具体而言，我们在NVIDIA RTX 2080Ti上以76帧每秒（FPS）的速度和Jetson Xavier NX上以8帧每秒（FPS）的速度，在Cityscapes验证集和测试集上分别实现了 $76.6\%$ 和 $75.9\%$ 的平均交并比（mIOU）。

一、引言

视觉场景理解在现代机器人系统中具有深远的意义。然而，自主机器对实时性的要求需要在计算密集型设计（如像素级图像语义分割）中进行关键权衡。因此，由于难以在准确性与效率（即计算复杂度、内存占用和执行速度）之间实现最佳平衡，低延迟语义分割成为一项具有挑战性的任务。传统的实时语义分割架构通常只关注上述视角之一，从而导致高精度设计的计算成本高昂，而高速模型的准确性相对较低。这些高速模型往往预测精度相对较低，例如[24]、[25]，而更准确的模型往往推理速度较慢且计算开销较高，例如[21]、[37]。在计算成本和执行速度方面，高速架构和高精度架构之间存在显著差距（见表I）。
在这里插入图片描述

实时分割设计通常面临几个挑战。首先，高精度设计如[21]、[37]严重依赖密集特征提取器，如ResNet-18 [6]。其次，在相对高速的算法如[24]、[25]中使用的浅层提取器虽然降低了计算成本，但无法提取足够的特征以实现精确分割。第三，尽管计算成本高昂的模型具有很高的准确性，但在推理过程中会出现一些局部和全局的不一致性。上述不一致性通常不会出现在非实时方法如[43]、[13]中，但在设计低延迟架构时，权衡有时并不理想。在这方面，受流行的双分支架构启发，我们提出了CABiNet——上下文聚合双边网络，其中我们设计了两个分支，一个用于快速有效的空间细节描述，另一个用于密集上下文嵌入。我们还通过为实时应用重新设计全局聚合和局部分布（GALD）块[13]，解决了局部和全局不一致性的问题。我们的速度-准确性权衡以及有效的空间-上下文特征融合使我们能够在Cityscapes数据集上以76帧每秒（FPS）的速度和 $75.9\%$ 的平均交并比（mIOU）分数超越之前最先进的实时语义分割方法。代码和训练模型将公开提供。

二、相关工作

实时语义分割已通过多种方法得到解决。Romera等人[27]建议使用具有残差连接的分解卷积来保持准确性和执行速度之间的平衡。Poudel等人[24]提出了一个带有瓶颈的双分支网络，以有效捕获局部和全局上下文以实现快速分割。后来，他们在[25]中提出了一个改进的学习降采样模块，以改善执行速度和准确性之间的权衡。Yu等人[37]提出了精确的双分支分割网络，其中提出了新颖的特征融合和注意力细化模块，用于精确的语义分割任务。在这方面，[44]还研究了具有多尺度跳跃连接的多个编码器-解码器对。与[37]、[40]相比，这种浅层和深层路径的集合被视为多个网络的集合，允许使用较浅的骨干网络（如ResNet-34）进行有效特征表示。实时分割的另一种方法是使用深度不对称瓶颈[12]，这从理论上提供了足够的感受野并捕获了密集上下文。
在这里插入图片描述

注意力模块具有建模长距离依赖性的能力，并且几位作者在各项工作中采用了注意力的概念[15]、[17]、[34]、[29]。注意力被引入机器理解首先是在[17]中实现的，其中学习了输入的全局依赖性，然后将其应用于自然语言处理。自那时以来，许多工作都利用这一概念来执行单尺度和多尺度下的多个场景理解任务[4]、[13]、[26]、[41]、[42]、[33]、[9]，从而超越了之前的传统上下文嵌入方法。

Jiang等人[10]发表了另一项关注上下文的工作，其中他们引入了上下文细化和上下文集成模块，以实现高效的场景分割。在[18]中提出了轻量级特征金字塔编码模型，这是对常规编码器-解码器架构的改进，使用了深度扩张卷积。在另外两种方法[31]、[38]中提出了多尺度上下文聚合，其中[31]使用类边界监督来处理某些相关的边界信息，而[38]使用优化的级联分解ASPP[2]模块来平衡准确性和执行速度之间的权衡。Orsic等人[21]开发了一种方法，该方法利用轻量级上采样和横向连接，以及以残差网络作为主要识别引擎，用于实时场景理解。该特定算法被视为当前在Cityscapes数据集上进行实时语义分割的最先进网络。

III. 方法

我们在图3中展示了CABiNet的架构，该架构包含两个分支，一个用于快速有效的空间细节描述，另一个用于密集上下文嵌入。空间和上下文分支允许以极低的计算量进行多尺度特征提取。然后，这两个分支在融合块（FFM）中融合，以进行最终的对象类别预测。
在这里插入图片描述

A. 空间分支

传统的实时设计通常要么将图像的分辨率缩小[39]，要么使用轻量级简化模型[1]、[23]来加速整体架构。然而，缩小图像会导致空间信息丢失，而轻量级模型则由于不断的通道剪枝而容易破坏感受野。这个问题在[37]中得到了解决，但代价是计算量显著增加，从而在移动和嵌入式平台上降低了执行速度。基于这些观察，我们提出了一个浅层分支，该分支能够编码丰富的空间信息并保持足够的感受野，同时从全分辨率图像中保持极低的计算成本。具体来说，该路径有四层，第一层是卷积层（大核尺寸），随后是批量归一化和ReLU，接着是两个深度卷积层。深度卷积的战略使用可以产生与传统卷积相同的结果，但计算量减少，并且可以通过增加特征表示的数量来补偿特征的边际损失。最后，最后一层是另一个核大小为1的卷积层。前三层的步长固定为2，而最后一层的步长为1。因此，该分支（图3）生成的输出是输入分辨率的 $\left(\frac{1}{8}\right)^{t h}$ [37]，从而在显著降低计算量的同时保持了所需的空间信息。

B. 上下文分支

如前所述，详细的空间信息结合足够的感受野，对语义分割的准确性有显著影响[37]。虽然浅层分支负责空间细节，但我们设计了一个新的基于注意力的上下文分支，包含轻量级全局聚合[13]和局部注意力[13]块，以提供足够的感受野并捕获全局和局部上下文。我们在此分支中使用预训练的MobileNetV3-Small[7]作为轻量级特征提取器，它可以有效且高效地对输入图像进行下采样，以提供丰富的高级语义特征。然而，这些特征尚未完善，因此需要传递给一个称为上下文聚合块的细化阶段，该块由简化的全局注意力[13]和局部注意力子模块组成。

MobileNetV3-Small：MobileNetV3[7]结合了MobileNetV2[28]和MnasNet[32]中建议的层，为移动应用构建了最有效和高效的神经网络。使用了改进的swish非线性函数来提高层的性能，以及用于squeeze-and-excitation模块的硬sigmoid。主干后的最终特征向量大小为 $64 \times 32 \times 576$ 。
上下文聚合块：特征提取器表示输出中的长距离和局部依赖对于准确的语义分割至关重要[43]、[13]、[35]。所提出的上下文聚合块旨在有效且高效地捕获这种跨通道和通道内映射。以前的一些工作已经提出了有效获取此类语义的模块[13]、[4]、[43]、[14]。在我们的工作中，我们采用了全局聚合和局部分布（GALD）[13]中的全局注意力（GA）块。该模块足够强大，可以捕获对于准确语义分割至关重要的长距离依赖，但计算量大，执行时需要大量的GPU内存。

简化全局注意力块。图3显示了全局注意力模块的流程。仔细观察该管道表明，[13]中提出的全局注意力模块可能存在两个可能的限制。首先，原始设计建议直接从主干输出中使用三个并行卷积层提取上下文信息，从而增加了所需参数的数量。其次，键（Key）和值（Value）卷积的矩阵乘法在softmax激活阶段之后的下一个乘法过程中进行，由于这些计算是在相对较大的矩阵上进行的[43]，因此增加了时间复杂度。

由于输入特征向量 $A$ 的大小，矩阵乘法运算量很大，如果将其改为一个更小的值 $M$ （其中 $M << A$ ），则可能有助于减轻部分计算量。尽管如此，更改必须以向量输出大小保持不变的方式进行。因此，我们在全局注意力模块中采用空间金字塔池化（SPP）模块[40]来有效减小特征向量的大小（图3）。与其将所有空间点输入到乘法过程中，不如对点进行采样，并将仅某些代表点输入到该过程中，这样做更为可行。根据[43]，我们在四个尺度上使用四个自适应最大池化，并将池化结果展平并拼接，作为下一层的输入。在我们的实验中，稀疏表示的数量可以表示为 $S=\sum_{n \in 1,3,5,8} n^{2}=110$ ，从而将复杂度降低到 $\mathscr{O}(\hat{N} A M)$ ，这远低于 $\mathscr{O}\left(\hat{N} A^{2}\right)$ 。具体来说，对于输入到全局注意力（GA）块的 $64 \times 32=2048$ ，这种非对称乘法节省了大约 $\frac{64 \times 32}{110} \approx 18$ 倍的计算成本。此外，池化模块捕获的特征统计量足以提供有关全局场景语义的线索。

接下来，该模块采用三个并行的 $\times 1$ 卷积层，这导致参数数量相对较大。这可能不会对整体执行速度产生直接影响，但参数较少的神经网络设计证明了模型的效率。常规卷积层具有多个可学习的滤波器，这些滤波器在输入特征向量上进行卷积。在[5]中建议，可以用一种称为廉价线性运算（CLO）的概念来替换这些常规卷积层，这在图3中以图形方式进行了描绘。与常规卷积相比，这些线性变换显著减少了参数和计算量[5]。

局部注意力。我们计算了每组的全局统计量，然后将它们乘回组内的特征。这里需要注意的是，计算统计量的窗口相对较大，因此统计线索可能会偏向于较大的模式，因为其中有更多的样本，这可能会进一步导致较小模式的过度平滑。

在这方面，[13]提出了一个局部注意力（LA）模块，以自适应地使用特征，同时考虑先前全局注意力块编码的每个位置的模式。我们的消融研究表明，该模块高效且快速，因此无需进行额外的改进。从根本上讲，LA块通过重新计算空间范围来预测局部权重，这主要是针对避免先前GA模块中存在的粗特征表示问题。在这里，预测的局部权重在全局信息和局部上下文之间添加了逐点权衡。因此，局部注意力块被建模为三个深度卷积层的集合，这允许对来自先前GA模块的特征表示进行微调。

瓶颈：受先前工作[7]、[6]的启发，我们设计了一个简单的下采样模块，以在深度维度上限制精细特征的表示。这种限制后来使我们能够监督注意力块和上下文分支的表示学习。

C. 特征融合模块（FFM）

值得注意的是，从两个分支提取的特征处于不同的表示尺度，并且需要尺度归一化以进行有效融合。因此，为了节省计算量而简单地将两个特征相加[24]、[25]，不太可能产生理想的准确性。因此，在这项工作中，我们实现了[37]中提出的特征融合技术，并进行了某些调整。为了充分利用来自两个分支的向量表示，我们首先拼接两个特征，然后是一个下采样瓶颈。拼接后，最终的特征表示具有较大的维度，这增加了所需计算量。添加一个下采样瓶颈可以在特征选择（加权注意力）的后期阶段显著减少这些计算量，而不会损害整体准确性（表VI）。受[37]、[8]启发，添加了加权注意力部分，以根据特征对整体预测准确性的贡献来选择性地加权特征。这些选定的特征随后被上采样以生成与[37]中相同数量的表示，但计算量显著减少。上采样瓶颈之后的最后两层生成最终的输出预测。在这种情况下，我们仅使用两层，因为对于简单的类别分离，多层变得不必要，因此一个深度可分离卷积和一个逐点卷积就足够了。详细的示意图如图4所示。
在这里插入图片描述

D. 损失函数

为了训练模型，我们使用三个带有在线难例挖掘[30]的交叉熵损失函数，其中一个（主要）用于最终输出，两个（辅助）用于上下文分支。辅助损失函数允许对上下文分支和注意力模块的学习进行深度监督。我们模型的总体联合损失表示 $L (X; W)$ 可以表示为：

$W)=l_{p}(X ; W)+l_{c 1}\left(X_{1} ; W\right)+l_{c 2}\left(X_{2} ; W\right)$

其中， $l_{p}$ 是监控总体输出的主要损失， $l_{c 1}$ 是整个上下文分支的辅助损失， $l_{c 2}$ 是上下文聚合块的辅助损失， $W$ 是网络参数， $p$ 是网络预测的最终输出。如[43]、[37]所示，使用联合损失可以更容易地优化模型。

IV. 实验

我们在Cityscapes数据集[3]上对我们提出的方法进行了基准测试。Cityscapes是一个城市场景理解数据集，其中共包含5000张完全标注的图像，其中2975张用于训练，500张用于验证，剩余的1525张用于测试。该数据集包含35个类别，其中19个用于城市场景理解，图像大小为 $1024 \times 2048$ 。

A. 训练设置

为了优化网络，我们使用随机梯度下降（SGD）[11]，并将Cityscapes的初始学习率设置为 $e^{-4}$ 。我们采用多项式学习率策略，即在训练过程中，学习率乘以 $KaTeX parse error: Expected 'EOF', got '_' at position 42: … }}{\text { max_̲iter }}\right)^…$ ，其中power等于0.9。对于Cityscapes图像，我们在训练期间从原始输入图像中随机裁剪出大小为 $[1024, 1024]$ 的图像块。我们对两个数据集都使用了数据增强技术，如随机水平翻转、随机缩放和颜色抖动。缩放比例范围从 $(0.75, 1.0, 1.5, 1.75, 2.0)$ 。Cityscapes的批大小设置为6。此外，Cityscapes的训练迭代次数设置为160k。所有实验均在单个NVIDIA RTX 2080Ti和Jetson Xavier NX上使用PyTorch 1.4.0进行。
在这里插入图片描述

B. 与最先进的方法比较

表I基于GPU内存占用、MAdd/GLOPs计数、执行速度（RTX 2080Ti）以及在Cityscapes验证集和测试集上的整体mIOU分数，提供了我们的方法与其他架构之间的详细比较。从表中可以看出，我们的模型在实时场景理解方面优于以前的方法，在验证集和测试集上分别获得了最高的mIOU分数76.6%和75.9%。与内存效率最高的模型SINet[22]相比，CABiNet的mIOU高出7.7%，推理速度更快7.8 FPS。与最快的模型Fast-SCNN[25]相比，CABiNet的mIOU高出7.5%。与最近的工作GAS[16]相比，我们的模型在假设全分辨率图像输入的情况下，在速度上具有竞争力，同时mIOU高出4.1%。定性结果如图5所示。与[21]相比，我们的模型在检测如电线杆、交通标志等代表性不足的物体方面表现更好。得益于高效的全局和局部语义聚合，我们的模型不会出现局部或全局不一致的情况。此外，与以前建立的最先进算法[25]、[21]、[37]相比，我们的改进在准确性和速度上都更有优势。我们架构中的计算开销（如参数计数、GFLOPs等）显著低于现有的准确实时架构，同时提高了准确性。优化的GALD块与高效的空间细节和轻量级密集提取器相结合，使我们的方法在多个方面优于传统的实时语义分割架构。Cityscapes测试集上的更多定性结果如图6所示。
在这里插入图片描述

C. 消融研究

基线定义为一个简单的双分支网络，其中空间分支有两个卷积层，第二分支有一个未训练的特征提取器。基线没有注意力和瓶颈模块，其结构与[24]相似。为了融合来自两个分支的特征，我们简单地将它们相加，然后由一个小分类器块将其区分为相应数量的类别。与[24]不同，两个分支都以相同的分辨率接收图像，并且所有消融实验都是在这个基线上进行的。
在这里插入图片描述

上下文聚合块：上下文聚合块（CAB）是专为有效且高效地捕获局部和全局上下文而设计的。如果我们从设计中移除CAB，同时保持所有其他模块以及训练/推理参数不变，我们会发现整体mIOU分数下降了 $2.1\%$ ，同时推理时间减少了近3毫秒。这表明，添加上下文块可以增强特征表示，同时对整体执行速度和复杂性的影响微乎其微。表III进一步证明了上下文聚合块的有效性，它可以作为即插即用模块与其他双分支语义分割架构结合使用。

有趣的是，[43]中建议使用SPP模块[40]作为注意力模块，但添加廉价线性运算（CLO）[5]不仅减少了所需的计算量，还提供了略高的准确性（表IV）。这可能是因为在这些线性变换中，可以存在多个核大小[5]，从而允许多尺度特征聚合。

主干网络选择：许多之前的实时语义分割架构[37]、[21]、[44]、[39]采用了强大的特征提取器，如ResNet-18[6]。尽管这种选择对于准确的语义分割是合理的，但对执行速度和计算复杂度的影响很大。因此，为了进行有效比较，我们将我们的MobileNetV3主干网络替换为ResNet-18，并研究了结果（表V）。从表中我们确认，即使我们使用像ResNet-18这样昂贵的特征提取器，我们的分割头相比SwiftNet[21]和BiSeNet[37]仍然更轻、更快、更准确。此外，表V中CABiNet-R18和CABiNet-MV3的比较显示，尽管ResNet-18和MobileNetV3-Small提供了相似的mIOU分数，但ResNet-18增加的计算开销相比MobileNetV3-Small更大。
在这里插入图片描述
特征融合模块：文献中提出了多种融合技术，设计合适的融合技术对最终结果有显著影响。请参考表VI，了解各种融合技术之间的定量比较。在所有变体中，带加权注意力和瓶颈的特征拼接提供了最优的mIOU-FLOPs平衡。
嵌入式设备上的结果：在全尺寸GPU（Titan X或RTX20系列）上进行推理不太可能提供现实世界的分析，因为自动驾驶汽车、无人机[20]和无人地面车辆更可能配备具有有限内存的低功耗模块。因此，我们在全分辨率Cityscapes图像上，使用小型系统级模块Jetson Xavier NX，对我们的算法和其他算法进行了进一步的基准测试。结果如表I最后一列所示。