PP-LiteSeg: 来自baidu的实时语义分割模型

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

---

导读

来自baidu，提出了一种名为PP-LiteSeg的实时且手工设计的网络。

摘要

现实世界的应用对语义分割方法有很高的需求。虽然借助深度学习，语义分割取得了显著的进步，但实时方法的表现并不令人满意。在本工作中，我们提出了一种新的轻量级模型 PP-LiteSeg，用于实时语义分割任务。具体来说，我们提出了一个灵活且轻量级的解码器（FLD）来减少先前解码器的计算开销。为了增强特征表示，我们设计了一个统一注意力融合模块（UAFM），该模块利用空间和通道注意力产生权重，然后将输入特征与该权重融合。此外，我们还提出了一种简单的金字塔池化模块（SPPM），以较低的计算成本聚合全局上下文。广泛的评估表明，与其它方法相比，PP-LiteSeg在准确性和速度之间实现了更优的权衡。在Cityscapes测试集上，PP-LiteSeg在NVIDIA GTX 1080Ti上分别达到了72.0% mIoU/273.6 FPS和77.5% mIoU/102.6 FPS。源代码和模型可在PaddleSeg: https://github.com/PaddlePaddle/PaddleSeg 获取。

1、介绍

语义分割旨在精确预测图像中每个像素的标签。它已被广泛应用于实际场景中，例如医疗影像、自动驾驶、视频会议、半自动标注等领域。作为计算机视觉中的基础任务之一，语义分割吸引了众多研究者的关注。

随着深度学习的显著进展，许多基于卷积神经网络的语义分割方法被提出。FCN是第一个采用端到端和像素到像素方式训练的全卷积网络。它还提出了语义分割领域中原始的编码器-解码器架构，这一架构被后续的方法广泛采纳。为了实现更高的准确性，PSPNet利用金字塔池化模块来聚合全局上下文信息；SFNet则提出流对齐模块以加强特征表示。

然而，这些模型由于计算成本较高而不适用于实时应用。为了加快推理速度，Espnetv2使用轻量级卷积从扩展的接收域中提取特征。BiSeNetV2提出了双边分割网络，并分别提取细节特征和语义特征。STDCSeg设计了一个名为STDC的新骨干网络来提高计算效率。不过，这些模型在准确性和速度之间的平衡方面并没有达到令人满意的水平。

在本工作中，我们提出了一种名为PP-LiteSeg的实时且手工设计的网络。如图2所示，PP-LiteSeg采用了编码器-解码器架构，并包含三个新颖的模块：灵活且轻量级的解码器（FLD）、统一注意力融合模块（UAFM）以及简单的金字塔池化模块（SPPM）。以下介绍这些模块的设计动机和详细内容。

语义分割模型中的编码器提取层次化的特征，而解码器则融合并上采样特征。对于编码器中从低层到高层的特征，通道数逐渐增加，空间尺寸逐渐减小，这是一种高效的设计。对于解码器中从高层到低层的特征，空间尺寸逐渐增加，而最近的模型中通道数保持不变。因此，我们提出了一个灵活且轻量级的解码器（FLD），它逐渐减少通道数并增加空间尺寸。此外，所提出的解码器的容量可以根据编码器轻松调整。这种灵活的设计平衡了编码器和解码器的计算复杂度，使得整体模型更加高效。

强化特征表示是提高分割准确性的一种关键途径。这通常是通过在解码器中融合低层和高层特征来实现的。然而，现有方法中的融合模块通常面临较高的计算成本。在本工作中，我们提出了一种统一注意力融合模块（UAFM）来高效地强化特征表示。如图4所示，UAFM首先利用注意力模块产生权重α，然后将输入特征与α融合。在UAFM中，有两种类型的注意力模块，即空间注意力模块和通道注意力模块，它们利用输入特征的空间间和通道间的相互关系。

上下文聚合是提高分割准确性的另一个关键因素，但之前的聚合模块对于实时网络来说耗时较长。基于PPM的框架，我们设计了一个简单的金字塔池化模块（SPPM），它减少了中间和输出通道数，去除了短路连接，并用加法操作替换了拼接操作。实验结果表明，SPPM在较低的计算成本下提高了分割精度。

我们在Cityscapes和CamVid数据集上进行了广泛的实验来评估所提出的PP-LiteSeg。如图1所示，PP-LiteSeg在分割精度和推理速度之间实现了优越的权衡。具体而言，在Cityscapes测试集上，PP-LiteSeg分别达到了72.0% mIoU/273.6 FPS和77.5% mIoU/102.6 FPS的成绩。

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

• 我们提出了一种灵活且轻量级的解码器（FLD），它减轻了解码器的冗余，并平衡了编码器和解码器的计算成本。

• 我们提出了一个统一注意力融合模块（UAFM），该模块利用通道和空间注意力来加强特征表示。

• 我们提出了一种简单的金字塔池化模块（SPPM）来聚合全局上下文。SPPM在额外少量的推理时间内提高了分割精度。

• 基于上述模块，我们提出了PP-LiteSeg，一种实时语义分割模型。广泛的实验证明了我们在准确性和速度方面的最新研究成果（SOTA）。

2、相关工作

2.1 语义分割

深度学习帮助语义分割取得了显著的进步。FCN是首个用于语义分割的全卷积网络。它采用端到端和像素到像素的方式进行训练。此外，FCN可以对任意大小的图像进行分割。在FCN的设计基础上，后来提出了多种方法。Segnet利用编码器中最大池化操作的索引来进行解码器中的上采样操作。因此，解码器中的信息得到了复用，并产生了细化的特征。PSPNet提出了金字塔池化模块来聚合局部和全局信息，这对提高分割精度非常有效。此外，近期的一些语义分割方法利用了Transformer架构来获得更好的准确性。

2.2 实时语义分割

为了满足语义分割的实时需求，提出了许多方法，例如轻量级模块设计、双分支架构、早期降采样策略、多尺度图像级联网络。ENet采用早期降采样的策略来降低处理大图像和特征图的成本。为了提高效率，ICNet设计了一个多分辨率图像级联网络。基于双边分割网络，BiSeNet分别提取细节特征和语义特征。双边网络较轻，因此推理速度很快。STDCSeg提出了通道减少和接收域扩大的STDC模块，并设计了一个高效的骨干网络，可以在较低的计算成本下加强特征表示。为了消除双分支网络中的冗余，STDCSeg使用详细的ground truth信息指导特征，因此效率进一步提高。Espnetv2使用分组逐点和深度可分离的空洞卷积来以计算友好的方式从扩展的感受野中学习特征。

2.3 特征融合模块

特征融合模块常用于语义分割中以加强特征表示。除了逐元素相加和拼接的方法外，研究人员还提出了以下几种方法。在BiSeNet中，BGA模块采用逐元素乘法来融合来自空间和上下文分支的特征。为了增强具有高级上下文的特征，DFANet以阶段性和子网级别的方式融合特征。为了解决对齐问题，SFNet和AlignSeg首先通过一个卷积神经网络模块学习变换偏移，然后将这些变换偏移应用于网格采样操作以生成细化后的特征。具体来说，SFNet设计了流对齐模块。AlignSeg设计了对齐特征聚合模块和对齐上下文建模模块。FaPN通过将变换偏移应用于变形卷积来解决特征对齐问题。

3、方法

在本节中，我们首先分别介绍灵活轻量解码器（FLD）、统一注意力融合模块（UAFM）和简易金字塔池化模块（SPPM）。然后，我们展示了实时语义分割的PP-LiteSeg架构。

3.1 灵活轻量解码器

编码器-解码器架构已被证明对语义分割非常有效。通常情况下，编码器利用一系列分组为多个阶段的层来提取层次化的特征。从低级到高级的特征，通道数逐渐增加而特征的空间尺寸减小。这种设计平衡了每个阶段的计算成本，从而保证了编码器的效率。解码器也有几个阶段，负责融合和上采样特征。虽然特征的空间尺寸从高级到低级增加，但最近的轻量化模型中的解码器在所有级别保持相同的特征通道数。因此，浅层阶段的计算成本远大于深层阶段，导致浅层阶段的计算冗余。为了提高解码器的效率，我们提出了一种灵活轻量解码器（FLD）。如图3所示，FLD从高级到低级逐步减少特征的通道数。FLD可以轻松调整计算成本以实现编码器与解码器之间更好的平衡。尽管FLD中的特征通道数在减少，我们的实验表明PP-LiteSeg相比其他方法仍能达到竞争性的准确性。

3.2 统一注意力融合模块

如上所述，融合多级特征对于实现高精度的分割至关重要。除了逐元素求和和拼接的方法外，研究人员还提出了几种方法，例如SFNet、FaPN 和AttaNet。在此工作中，我们提出了一种统一注意力融合模块（UAFM），它通过通道和空间注意力来丰富融合后的特征表示。

UAFM框架。如图4(a)所示，UAFM利用一个注意力模块来产生权重α，并通过乘法（Mul）和加法（Add）操作将输入特征与α融合。具体来说，输入特征记为和。是较深层次模块的输出，而则是来自编码器的对应特征。需要注意的是，它们具有相同的通道数。UAFM首先使用双线性插值操作将上采样到与相同的大小，上采样的特征记为。然后，注意力模块以和作为输入并产生权重α。需要注意的是，注意力模块可以是一个插件，例如空间注意力模块、通道注意力模块等。之后，为了获得带注意力权重的特征，我们分别对和应用逐元素乘法操作。最后，UAFM对带注意力权重的特征执行逐元素加法操作，并输出融合后的特征。我们可以将上述过程公式化为：

空间注意力模块。空间注意力模块的动机是利用像素间的空间关系来产生权重，该权重代表输入特征中每个像素的重要性。如图4(b)所示，给定输入特征，即和，我们首先沿通道轴执行均值和最大值操作来生成四个特征，这些特征的维度为。之后，这四个特征被拼接成一个特征。对于拼接后的特征，应用卷积和sigmoid操作来输出。空间注意力模块的公式表示如方程2所示。此外，空间注意力模块可以灵活调整，例如移除最大值操作以降低计算成本。

通道注意力模块。通道注意力模块的核心概念是利用通道间的相互关系来生成权重，该权重指示输入特征中每个通道的重要性。如图4(b)所示，提出的通道注意力模块利用平均池化和最大池化操作来压缩输入特征的空间维度。这一过程生成四个维度为的特征。然后，沿着通道轴拼接这四个特征，并执行卷积和sigmoid操作来产生权重α ∈ 。简而言之，通道注意力模块的过程可以用方程3表示。

3.3 简易金字塔池化模块

如图5所示，我们提出了一种简易金字塔池化模块（SPPM）。它首先利用金字塔池化模块来融合输入特征。金字塔池化模块包含三个全局平均池化操作，对应的池化窗口大小分别为1×1、2×2和4×4。之后，输出特征经过卷积和上采样操作。对于卷积操作，卷积核大小为1×1，输出通道数小于输入通道数。最后，我们将这些上采样的特征相加，并应用一次卷积操作来产生精炼后的特征。与原始的金字塔池化模块（PPM）相比，SPPM减少了中间和输出通道数，去除了短路连接，并用加法操作替换了拼接操作。因此，SPPM更加高效且适合实时模型。

3.4 网络架构

提出的PP-LiteSeg架构如图2所示。PP-LiteSeg主要由三个模块组成：编码器、聚合模块和解码器。

首先，给定一张输入图像，PP-Lite使用一种常见的轻量级网络作为编码器来提取层次化的特征。具体来说，我们选择了STDCNet，因为它表现出色。STDCNet有5个阶段，每个阶段的步长为2，因此最终的特征大小是输入图像的1/32。如表1所示，我们提出了两个版本的PP-LiteSeg，即PP-LiteSeg-T和PP-LiteSeg-B，其编码器分别为STDC1和STDC2。PP-LiteSeg-B实现了更高的分割精度，而PP-LiteSeg-T的推理速度更快。值得注意的是，我们在编码器训练中应用了SSLD方法来获得增强的预训练权重，这对分割训练的收敛是有益的。

其次，PP-LiteSeg采用SPPM来建模长距离依赖。以编码器的输出特征为输入，SPPM产生包含全局上下文信息的特征。

最后，PP-LiteSeg利用我们提出的FLD逐渐融合多级特征并输出结果图像。具体来说，FLD由两个UAFM和一个分割头组成。为了提高效率，我们在UAFM中采用了空间注意力模块。每个UAFM接受两个特征作为输入，即编码器阶段提取的一个低级特征和由SPPM或更深层次融合模块产生的一个高级特征。后一个UAFM输出的融合特征具有1/8的下采样比例。在分割头中，我们执行Conv-BN-ReLU操作来减少1/8下采样特征的通道数至类别数量。随后，执行上采样操作将特征大小扩展到输入图像大小，并使用argmax操作预测每个像素的标签。我们采用了带有在线难例挖掘的交叉熵损失来优化模型。

4、实验

4.1 实现细节

训练设置。遵循常见的设置，我们选择动量为0.9的随机梯度下降（SGD）算法作为优化器。我们也采用了预热策略和“多项式”学习率调度器。对于Cityscapes数据集，批大小为16，最大迭代次数为160,000次，初始学习率为0.005，优化器中的权重衰减为5e−4。对于CamVid数据集，批大小为24，最大迭代次数为1,000次，初始学习率为0.01，权重衰减为1e−4。对于数据增强，我们使用随机缩放、随机裁剪、随机水平翻转、随机颜色抖动以及归一化。随机缩放的比例范围为[0.125, 1.5]和[0.5, 2.5]，分别适用于Cityscapes和CamVid数据集。Cityscapes数据集的裁剪分辨率为1024×512，CamVid数据集的裁剪分辨率为960×720。

推理设置。为了公平比较，我们将PP-LiteSeg导出为ONNX格式，并利用TensorRT执行模型。类似于其他方法，Cityscapes的一张图像首先被调整到1024×512和1536×768的大小，然后推理模型接受调整后的图像并产生预测图像，最后将预测图像调整回输入图像的原始大小。这三个步骤的时间成本被计为推理时间。对于CamVid数据集，推理模型直接以原始图像为输入，分辨率设为960×720。所有推理实验都在CUDA 10.2、CUDNN 7.6、TensorRT 7.1.3环境下，在NVIDIA 1080Ti GPU上进行。我们使用标准mIoU来比较分割精度，并使用FPS来比较推理速度。

4.2 和其他方法的对比

4.2 消融实验

5、结论

本文着重于设计一种用于语义分割的新型实时网络。首先，我们提出了灵活轻量解码器（FLD），以提高先前解码器的效率。然后，我们提出了一种统一注意力融合模块（UAFM），该模块有效地增强了特征表示。此外，我们还提出了一种简易金字塔池化模块（SPPM），以较低的计算成本来聚合全局上下文信息。基于这些新颖的模块，我们提出了PP-LiteSeg，这是一种实时语义分割网络。广泛的实验证明，PP-LiteSeg在分割精度和推理速度之间达到了最先进的权衡。在未来的工作中，我们计划将我们的方法应用于更多任务，例如抠图和交互式分割。

—END—

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

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

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