STDCNet：对BiSeNet做语义分割的重新思考

ronghuaiyang

于 2024-08-23 12:01:27 发布

阅读量950

点赞数 25

本文链接：https://blog.csdn.net/u011984148/article/details/141482799

版权

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

导读

来自美团，对Bisenet进行了网络结构的优化，并利用了边界信息作为监督信号。

摘要

BiSeNet已被证明是一种用于实时分割的流行双流网络。然而，其通过额外路径编码空间信息的原则耗时较长，而且从预训练任务（如图像分类）借用的主干网络由于缺乏针对特定任务的设计，可能对图像分割不够高效。为了解决这些问题，我们提出了一种名为短期密集拼接网络（STDC网络）的新颖且高效的结构，通过去除结构冗余来实现这一点。具体来说，我们逐步减少特征图的维度，并使用它们的聚合来进行图像表示，这构成了STDC网络的基本模块。在解码器中，我们通过将空间信息的学习集成到单一流方式的低级层中，提出了一种细节聚合模块。最后，将低级特征和深层特征融合以预测最终的分割结果。在Cityscapes和CamVid数据集上的广泛实验证明了我们方法的有效性，实现了分割精度和推理速度之间的良好折衷。在Cityscapes数据集上，我们以250.4 FPS的速度在测试集上达到了71.9%的mIoU，比最新方法快45.2%，而在更高分辨率图像上的推理时，我们达到了76.8%的mIoU，速度为97.0 FPS。代码仓库：https://github.com/ MichaelFan01/STDC-Seg

1、介绍

语义分割是计算机视觉领域中的一个经典且基础的话题，其目标是在图像中为每个像素分配标签。深度学习的发展极大地促进了语义分割性能的进步，通过各种突破性进展，在自动驾驶、视频监控、机器人感知等领域的需求迅速增长。这些应用激励研究者探索有效的且高效的分割网络，特别是针对移动设备领域。

为了满足这些需求，许多研究者提出设计低延迟、高效率的CNN模型，同时保持令人满意的分割精度。这些实时语义分割方法在各种基准测试上取得了显著的成绩。对于实时推理，一些工作，如DFANet和BiSeNetV1选择了轻量级主干网络，并研究了特征融合或聚合模块的方法来弥补精度的下降。然而，这些从图像分类任务借用的轻量级主干网络可能并不完全适合于图像分割问题，因为缺乏针对特定任务的设计。除了选择轻量级主干网络之外，限制输入图像的大小是另一种常用的提高推理速度的方法。较小的输入分辨率看起来很有效，但它很容易忽略边界附近的详细外观和小物体。为了解决这个问题，如图2(a)所示，BiSeNet采用了多路径框架来结合低级细节和高级语义。然而，添加额外的路径来获取低级特征是耗时的，而且辅助路径通常缺乏低级信息的引导。

为此，我们提出了一种新的手工设计的网络，旨在实现更快的推理速度、可解释的结构以及与现有方法相当的性能。首先，我们设计了一种称为短期密集拼接模块（STDC模块）的新结构，以少量参数获得可变的可扩展感受野。然后，将STDC模块无缝集成到U-net架构中，形成STDC网络，大大提高了在网络在语义分割任务中的性能。

具体来说，如图3所示，我们从多个连续层拼接响应图，每一层都以不同的尺度和视场编码输入图像/特征，从而产生多尺度特征表示。为了加速，随着分割性能的微小损失，各层的滤波器大小逐渐减小。STDC网络的具体结构可以在表2中找到。

在解码阶段，如图2(b)所示，我们不是使用额外耗时的路径，而是采用了细节指导来引导低级层学习空间细节。我们首先使用细节聚合模块来生成细节的真实标签。然后，采用二元交叉熵损失和骰子损失来优化细节信息的学习任务，这被认为是侧信息学习的一种类型。值得注意的是，在推理时间中不需要这种侧信息。最后，将低级层中的空间细节和深层中的语义信息融合以预测语义分割的结果。我们方法的整体架构如图4所示。

我们的主要贡献总结如下：

我们设计了一个短期密集拼接模块（STDC模块）来提取具有可扩展感受野和多尺度信息的深层特征。这个模块以可承受的计算成本提高了我们的STDC网络的性能。
我们提出了细节聚合模块来学习解码器，使得在低级层中更加精确地保留空间细节，而不增加推理时间中的额外计算成本。
我们进行了广泛的实验来展示我们方法的有效性。实验结果表明，STDC网络在ImageNet、Cityscapes和CamVid数据集上取得了最新的最佳结果。具体来说，我们的STDC1-Seg50在Cityscapes测试集上以250.4 FPS的速度达到了71.9%的mIoU。在相同的实验设置下，我们的STDC2-Seg75以97.0 FPS的速度达到了76.8%的mIoU。

2、相关工作

2.1. 高效网络设计

模型设计在计算机视觉任务中扮演着重要角色。SqueezeNet使用了火模块和某些策略来减少模型参数。MobileNet V1利用了深度可分离卷积来减少推理阶段的浮点运算数（FLOPs）。ResNet采用了残差构建层来达到出色的表现。MobileNet V2和ShuffleNet使用了分组卷积来减少计算成本，同时保持了相当的准确性。这些工作特别针对图像分类任务设计，它们扩展到语义分割应用时应仔细调整。

2.2. 通用语义分割

传统的分割算法，如阈值选择、超像素等，利用手工设计的特征来为图像中的每个像素分配标签。随着卷积神经网络的发展，基于全卷积网络（FCN）的方法在各种基准测试上取得了显著的性能。Deeplabv3 采用了一个空洞空间金字塔池化模块来捕获多尺度上下文。SegNet利用了编码器-解码器结构来恢复高分辨率的特征图。PSPNet设计了一个金字塔池化模块来同时捕获局部和全局上下文信息。空洞主干和编码器-解码器结构可以同时学习低级细节和高级语义。然而，大多数方法由于高分辨率特征和复杂的网络连接而需要大量的计算成本。在本文中，我们提出了一种高效且有效的架构，实现了速度和准确性之间的良好权衡。

2.3 实时语义分割

近年来，实时语义分割的实际应用快速增长。在这种情况下，设计高效的分割方法主要有两种主流方向。（i）轻量级主干网络。DFANet采用了轻量级主干网络来减少计算成本，并设计了一个跨层级特征聚合模块来提升性能。DFNet利用了“部分顺序剪枝”算法来获得轻量级主干网络和高效的解码器。（ii）多分支架构。ICNet设计了多尺度图像级联来实现良好的速度-精度折衷。BiSeNetV1和BiSeNetV2提出了两条流路径，分别用于低级细节和高级上下文信息。在本文中，我们提出了一种高效的轻量级主干网络来提供可扩展的感受野。此外，我们设置了一个单路径解码器，该解码器使用细节信息指导来学习低级细节。

3、提出的方法

BiSeNetV1使用轻量级主干网络，例如ResNet18和空间路径作为编码网络来形成双流分割架构。然而，分类主干网络和双流架构可能由于结构冗余而不够高效。在本节中，我们首先介绍所提出的STDC网络的详细情况。然后，我们介绍具有细节指导的单流方法的整体架构。

3.1 编码网络的设计

3.1.1 短期密集拼接模块

我们提出的网络的关键组件是短期密集拼接模块（STDC模块）。图3(b)和(c)展示了STDC模块的布局。具体来说，每个模块被分为几个块，我们使用来表示第i个块的操作。因此，第i个块的输出可以按如下方式计算：

其中和分别是第i个块的输入和输出。ConvX 包括一个卷积层、一个批量归一化层和ReLU激活层，是卷积层的核大小。

在STDC模块中，第一个块的核大小为1，其余的块简单地设置为3。给定STDC模块输出的通道数N，第i个块中卷积层的滤波器数量为，除了最后一个卷积层的滤波器数量与前一个卷积层相同。在图像分类任务中，通常的做法是在较高层使用更多的通道。但在语义分割任务中，我们关注的是可扩展的感受野和多尺度信息。低级层需要足够的通道来编码具有较小感受野的更精细的信息，而具有较大感受野的高级层更专注于高级信息归纳，与低级层设置相同的通道可能会导致信息冗余。下采样只发生在Block2中。为了丰富特征信息，我们通过跳过路径将到到的特征图拼接作为STDC模块的输出。在拼接之前，STDC模块中不同块的响应图通过3×3的平均池化操作下采样到相同的大小，如图3(c)所示。在我们的设置中，STDC模块的最终输出为：

其中表示STDC模块的输出，F是我们方法中的融合操作，而是来自所有n个块的特征图。为了考虑效率，我们采用了拼接作为我们的融合操作。在我们的方法中，我们使用STDC模块的4个块。

表1列出了STDC模块中各块的感受野，因此从所有块中收集了多尺度信息。我们声称我们的STDC模块有两个优点：（1）我们精心调整了各块的滤波器大小，通过几何级数递减的方式显著降低了计算复杂度。（2）STDC模块的最终输出是从各个块的特征图拼接而成，这样既保留了可扩展的感受野，也保留了多尺度信息。

给定输入通道维度 𝑀和输出通道维度 𝑁，STDC 模块的参数数量可以计算如下：

如方程3所示，STDC模块的参数数量主要由预定义的输入和输出通道维度决定，而块的数量对参数大小的影响较小。特别是当n达到最大极限时，STDC模块的参数数量几乎保持不变，仅由M和N定义。

3.1.2 网络架构

我们的网络架构如图3(a)所示。它由6个阶段组成，不包括输入层和预测层。一般来说，第1至第5阶段分别以2的步长对输入进行下采样，而第6阶段通过一个ConvX、一个全局平均池化层和两个全连接层输出预测logits。

第1和第2阶段通常被视为用于外观特征提取的低级层。为了追求效率，我们在每个阶段的第1和第2阶段只使用一个卷积块，根据我们的经验，这已经被证明是足够的。在第3、第4和第5阶段中STDC模块的数量在我们的网络中被仔细调优。在这些阶段中，每个阶段的第一个STDC模块以2的步长下采样空间分辨率。每个阶段的后续STDC模块保持空间分辨率不变。

我们将阶段的输出通道数记为，其中l是阶段的索引。在实践中，我们经验性地设置了为1024，并仔细调优了其余阶段的通道数，直到达到准确性和效率之间的良好平衡。由于我们的网络主要由短期密集拼接模块组成，我们将其称为STDC网络。表2显示了我们的STDC网络的详细结构。

3.2 解码器的设计

3.2.1 分割架构

我们使用预训练的STDC网络作为编码器的主干，并采用BiSeNet的上下文路径来编码上下文信息。如图4(a)所示，我们使用第3、第4和第5阶段来生成下采样比例分别为1/8、1/16、1/32的特征图。然后，我们使用全局平均池化来提供具有大感受野的全局上下文信息。采用U形结构来上采样源自全局特征的特征，并在编码阶段将它们与最后两个阶段（第4和第5阶段）的对应特征相结合。按照BiSeNet的做法，我们使用注意力细化模块来细化每两个阶段组合的特征。对于最终的语义分割预测，我们采用BiSeNet中的特征融合模块来融合编码器中第3阶段的1/8下采样特征和解码器中的对应特征。我们声称这两个阶段的特征处于不同的特征表示级别。编码主干的特征保留了丰富的细节信息，而解码器的特征则含有上下文信息，这是由于全局池化层的输入导致的。具体来说，分割头包括一个3×3的Conv-BN-ReLU操作，后面跟着一个1×1的卷积来得到输出维度N，该维度被设置为类别数量。我们采用带有在线难例挖掘的交叉熵损失来优化语义分割学习任务。

3.2.2 低级特征的细节指导

我们在图5(b)中可视化了BiSeNet的空间路径特征。与具有相同下采样比率的主干低级层（第3阶段）相比，空间路径可以编码更多的空间细节，例如边界、角点。基于这一观察，我们提出了一种细节指导模块来引导低级层以单流方式学习空间信息。我们将细节预测建模为一个二分类分割任务。我们首先通过拉普拉斯算子从分割的真实标签生成细节真实标签，如图4(c)所示。如图4(a)所示，我们在第3阶段插入细节头来生成细节特征图。然后我们使用细节真实标签作为细节特征图的指导，以引导低级层学习空间细节的特征。如图5(d)所示，带有细节指导的特征图可以编码比图5(c)中提到的结果更多的空间细节。最后，学习到的细节特征与解码器深层块的上下文特征融合用于分割预测。

细节真实标签生成：我们通过我们的细节聚合模块从语义分割的真实标签生成二进制细节真实标签，如图4(c)中蓝色虚线框所示。可以通过名为拉普拉斯核的2-D卷积核和一个可训练的1×1卷积来执行此操作。我们使用图4(e)所示的拉普拉斯算子来生成不同步长的软细细节特征图，以获得多尺度细节信息。然后我们将细节特征图上采样到原始大小，并通过一个可训练的1×1卷积进行动态重加权。最后，我们采用阈值0.1将预测的细节转换为包含边界和角点信息的最终二进制细节真实标签。

细节损失：由于细节像素的数量远少于非细节像素，细节预测是一个类不平衡问题。由于加权交叉熵总是导致粗糙的结果，我们采用二元交叉熵和dice损失联合优化细节学习。dice损失衡量预测图与真实标签之间的重叠。此外，它对前景/背景像素的数量不敏感，这意味着它可以缓解类不平衡问题。因此，对于高度为H和宽度为W的预测细节图，细节损失可以表示为：

其中表示预测的细节，表示对应的细节真实标签。表示二元交叉熵损失，而表示骰子损失，它们可以表示为：

其中i表示第i个像素，是一个拉普拉斯平滑项，用于避免除零的情况。在本文中我们设置 = 1。

如图4(b)所示，我们使用一个细节头来生成细节图，该细节图指导浅层编码空间信息。细节头包括一个3×3的Conv-BN-ReLU操作，后跟一个1×1的卷积来获取输出细节图。在实验中，细节头被证明能够有效地增强特征表示。需要注意的是，在推理阶段这部分被丢弃。因此，这种侧信息可以在推理时不增加任何成本的情况下轻松提高分割任务的准确性。

4、实验

4.1 实现细节

图像分类。我们使用批大小为64的小批量随机梯度下降（SGD）来训练模型，动量为0.9，权重衰减为1e−4。我们采纳了三种训练方法，包括学习率预热、余弦学习率策略和标签平滑。总共训练300个epoch，在前5个epoch内采用预热策略，学习率从0.001线性增加到0.1。分类块之前的Dropout设置为0.2。我们没有使用其他特殊的数据增强。

语义分割。我们使用动量为0.9、权重衰减为5e−4的小批量随机梯度下降（SGD）。Cityscapes和CamVid数据集的批大小分别设置为48和24。按照常见配置，我们使用“poly”学习率策略，其中初始学习率乘以。设置power为0.9，初始学习率为0.01。此外，我们分别为Cityscapes和CamVid数据集训练60,000和10,000次迭代，在这些迭代中，我们分别在前1,000和200次迭代采用预热策略。

数据增强包括颜色抖动、随机水平翻转、随机裁剪和随机缩放。Cityscapes训练时的比例范围为[0.125, 1.5]，裁剪分辨率为1024×512。对于CamVid训练，比例范围为[0.5, 2.5]，裁剪分辨率为960×720。

在所有实验中，我们基于pytorch-1.1在一个docker容器中进行实验。我们在CUDA 10.0、CUDNN 7.6.4和TensorRT 5.0.1.5环境下使用NVIDIA GTX 1080Ti GPU进行所有实验，批大小设置为1来基准测试我们方法的计算能力。

4.2 消融实验

4.3 和其他方法对比

5、结论

在本文中，我们重新审视了经典的分割架构BiSeNet，并对其结构进行了优化。一般来说，BiSeNet的分类主干网络和额外的空间路径极大地阻碍了推理效率。因此，我们提出了一种新型的短期密集拼接模块来提取具有可扩展感受野和多尺度信息的深层特征。基于该模块，我们设计了STDC网络，并在图像分类任务中实现了高FPS的同时保持了竞争力的准确性。使用STDC网络作为主干网络，我们的细节指导型STDC-Seg在实时语义分割任务中实现了最先进的速度-精度折衷。广泛的实验和可视化结果表明了我们提出的STDC-Seg网络的有效性。未来的研究方向包括：（i）验证主干网络在更多任务中的有效性，例如目标检测。（ii）我们将在语义分割任务中进一步探索空间边界的利用。