Squeeze-and-Attention Networks for Semantic Segmentation

0.摘要

        最近,将注意力机制整合到分割网络中可以通过更重视提供更多信息的特征来提高它们的表征能力。然而,这些注意力机制忽视了语义分割的一个隐含子任务,并受到卷积核的网格结构的限制。在本文中,我们提出了一种新颖的squeeze-and-attention网络(SANet)架构,利用一个有效的squeeze-and-attention(SA)模块来考虑分割的两个独特特征:i)像素组注意力,和ii)像素级预测。具体来说,所提出的SA模块通过引入“注意力”卷积通道,将像素组注意力施加于传统的卷积上,从而以高效的方式考虑空间-通道的相互依赖关系。最终的分割结果是通过合并SANet的四个层次阶段的输出来进行的,以获得增强的像素级预测的多尺度上下文。对两个具有挑战性的公共数据集进行的实证实验验证了所提出的SANet的有效性,其在PASCAL VOC数据集上达到了83.2%的mIoU(无COCO预训练),并在PASCAL Context数据集上实现了54.4%的state-of-the-art mIoU。

1.引言

        分割网络已成为自动驾驶、医学图像分析、机器人导航和虚拟现实等领域的关键识别元素。分割方法的进展主要是通过改进像素级表示以实现准确的标注。然而,语义分割并不完全等同于像素级预测。在本文中,我们认为语义分割可以分解为两个独立的维度:像素级预测和像素组合。具体而言,像素级预测地址了每个像素的预测,而像素组合则强调像素之间的连接。以前的分割工作主要关注从像素级别改进分割性能,但很大程度上忽视了像素组合的隐含任务。通过将语义分割分解为两个子任务,可以发现被广泛忽视的像素组合任务。如图1所示,第一个子任务要求精确的像素级注释,为图像分类引入了空间约束。最近的分割模型通过使用金字塔池化和扩张卷积层来聚合上下文特征,以进行像素级标注,取得了显著的进展。然而,这些卷积核的网格结构限制了在分割网络中学习到的空间特征的形状。特征聚合策略增强了像素级预测结果,但对于全局图像理解的视角仍未充分利用。

        为此,我们引入了像素组合的第二个子任务,直接鼓励将属于同一类的像素组合在一起,而不受空间限制。像素组合涉及将从一系列电磁波谱采样的图像转化为在任务特定语义谱中定义的像素组,其中语义谱的每个条目对应一个类。受到设计一个考虑像素组合的模块的启发,我们设计了一种新颖的squeeze-and-attention(SA)模块来缓解卷积核的局部约束。SA模块包含下采样但未完全压缩的注意力通道,以有效产生非局部的空间注意力,同时避免在输出头中使用大幅度扩张的卷积。具体而言,我们使用注意力卷积来生成注意力蒙版,因为每个卷积核都会在输入特征图上扫过。与增强骨干网络的SE模块不同,SA模块将空间注意力整合到网络的头部单元中,它们的输出被汇聚以提高分割性能。由SA模块引入的空间注意机制强调了不同空间尺度上属于同一类的像素组的关注。此外,压缩通道可以作为全局注意力蒙版。

        我们设计了包含四个SA模块的SANets来处理分割的上述两个任务。SA模块学习多尺度的空间特征和非局部的光谱特征,从而克服卷积层在分割中的限制。我们使用扩张ResNet和Efficient Nets作为骨干网络,以充分发挥它们在图像识别方面的强大能力。为了聚合多阶段的非局部特征,我们在骨干网络的多个阶段输出上采用SA模块,从而得到更好的物体边界和场景解析结果。这种简单但有效的创新使得将SANets推广到其他相关的视觉识别任务更加容易。我们使用两个具有挑战性的分割数据集(PASCAL context和PASCAL VOC 2012)验证了SANets的性能。本文的贡献有三个方面:

•我们将语义分割分解为两个子任务:像素级的密集预测和像素组合。

•我们设计了一种称为squeeze-and-attention(SA)模块,它同时考虑了个体像素的多尺度密集预测和像素组的空间注意力。

•我们提出了一种multi-level heads的squeeze-and-attention网络(SANet),以利用SA模块的表示增强能力,并集成多尺度的上下文特征和图像级别的类别信息。

Squeeze-and-Attention Networks for Semantic Segmentation_第1张图片图1:语义分割可以分解为两个子任务:明确的像素级预测和隐式的像素分组。这两个任务将语义分割与图像分类分开。受到像素分组的设计的启发,我们设计了一个新颖的压缩注意力(SA)模块以及一个SANet,以提高稠密预测的性能并解决被忽视的像素分组问题。Squeeze-and-Attention Networks for Semantic Segmentation_第2张图片图2:(a)残差块;(b)压缩和激励(SE)模块;(c)压缩和注意(SA)模块。为了简化,我们展示了卷积(CONV)、全连接(FC)、平均池化(Avg.Pool)层,但省略了归一化和激活层。SA模块具有与SE模块类似的结构,其中包含一个额外的路径来学习重新校准输出特征图Xout的通道权重。不同之处在于SA模块的注意力通道使用平均池化来降采样特征图,而不是像SE模块中那样完全压缩。因此,我们将该通道称为注意卷积(ACONV)通道。

2.相关工作

        多尺度上下文。近年来,在语义分割方面取得的最新进展主要得益于将多尺度上下文特征纳入分割模型以提取判别性特征。一种拉普拉斯金字塔结构被引入以结合多尺度特征。RefineNet使用多路径方式明确地整合从多尺度输入中提取的特征以增强分割结果。编码器-解码器结构已经被用于融合具有不同语义意义水平的特征。最流行的方法使用池化操作从不同尺度收集空间信息。类似地,EncNet采用编码模块将不同上下文投影到高斯核空间以编码多尺度上下文特征。图形模型如CRF和MRF被用来施加光滑性约束以获得更好的分割结果。最近,设计了gather-excite模块,通过聚集远距离上下文中的特征来缓解经典卷积的局部特征约束。我们通过合并骨干残差网络的不同阶段输出来改进多尺度密集预测。

        通道注意力。有选择地对特征图的通道进行加权有效地增加了传统残差模块的表示能力。一个很好的例子是squeeze-and-excitation(SE)模块,因为它强调在特征图的选定通道上的注意力。这个模块通过将相关的类别进行组合显著提高了残差网络的分类准确性。EncNet也使用了SE模块的类别识别能力。Discriminative Feature Network(DFN)在其平滑子网络中利用了通道加权范式。虽然重新校准特征图通道的光谱权重被证明对改善卷积层的表示能力有效,但是实现(例如squeeze-and-excitation模块)会导致过多的模型参数。与SE模块相比,我们设计了一个新颖的squeeze-and-attention(SA)模块,其中包含一个经过下采样但不完全挤压的卷积通道,用于生成灵活的模块。具体而言,这个额外的通道为像素分组生成特定类别的软注意力掩码,同时在经典卷积通道之上添加了缩放的空间特征,用于像素级别的预测。

        像素组注意力。注意机制在神经语言处理中的成功推动了其在语义分割中的应用。空间变换网络明确地学习空间注意力,以仿射变换的形式增加特征的不变性。由于机器翻译和图像翻译具有许多相似之处,RNN和LSTM被用于通过将语义标签与翻译相连接进行语义分割。与此同时,[7]采用了一种尺度敏感的注意力策略,使网络能够关注不同尺度的对象。[42]设计了一种特定的空间注意力传播机制,包括一个收集通道和一个扩散通道。[35]通过计算相关指标使用自注意力掩码。[18]通过收集局部特征设计了一种聚集-激励操作来生成图像分类的硬掩码。此外,[36]证明了具有边际计算成本的完全挤压模块对于图像分类是有效的。由于由空间非对称重校准(SAR)模块生成的权重是向量形式,因此不能直接用于分割。与现有的注意力模块不同,我们使用由池化层实现的下采样通道来聚合多尺度特征并同时生成软全局注意力掩码。因此,SA模型增强了像素级密集预测的目标,并考虑到了迄今为止被大部分忽视的像素组注意力。

3.框架

        经典卷积主要关注空间局部特征编码,而通过使用全局图像信息有选择地重新加权特征图通道,挤压和激活(SE)模块增强了其性能[19]。受到这种简单但有效的SE模块对图像级别分类的启发,我们设计了一个挤压和注意(SA)模块,它融合了用于像素级密集预测的全卷积层的优点,并额外增加了一种更局部的特征图重新加权形式,我们称之为像素组注意力。类似于提高分类性能的SE模块,SA模块是专门设计用于改善分割结果的。

3.1.压缩和激励模块

        由于残差网络(ResNets)在图像识别方面表现出色,因此它们被广泛用作分割网络的骨干,已经证明,在大型图像数据集ImageNet上预训练的ResNets能够很好地迁移到其他视觉任务,包括语义分割[41,5]。由于经典卷积可以被视为一种空间注意力机制,我们从作为ResNets基本组件的残差块开始。如图2(a)所示,传统的残差块可以表示为:

        其中F(·)表示残差函数,由Θ参数化,Ω表示两个卷积层的结构。Xin ∈ R^C'×H'×W'和Xout ∈ R^C×H×W 分别是输入和输出特征图。SE模块通过重新校准特征图通道来改善残差块,值得注意的是,我们采用了SE模块的更新版本,在[19]中与原始版本的性能相当。如图2(b)所示,SE模块可以表示为:        其中,用于重新校准输入特征图Xin通道的学习权重w计算如下:        其中,Φ(·)表示sigmoid函数,σ(·)表示ReLU激活函数。首先,使用平均池化层对输入特征图Xin进行“挤压”。然后,采用由W1和W2参数化的两个全连接层来获取“激发”权重。通过添加这样一个简单的重新加权机制,SE模块有效地增加了残差块的表示能力。Squeeze-and-Attention Networks for Semantic Segmentation_第3张图片图3:压缩注意力网络。SANet聚合了来自多个分层SA头部的输出,生成多尺度的按类别分组掩码,解决语义分割中被忽视的像素分组任务。这些掩码的训练由地面真实注释中的相应分类区域监督。此外,这些掩码也用于引导像素级预测,即FCN头部的输出。通过这种方式,我们同时利用了SA模块的像素组注意力提取能力和多尺度的上下文特征。

3.2.压缩和注意力模块

        语义分割需要在图像的全局和局部层次上具有有用的表示。在像素级别,卷积层根据局部信息生成特征图,因为卷积是根据每个像素周围的局部进行计算的。像素级别的卷积为所有语义分割模块奠定了基础,并且以各种方式增加卷积层的感受野可以提升分割性能,表明更大的上下文对于语义分割是有用的。在图像的全局层面上,可以利用上下文来确定特征图的哪些部分被激活,因为上下文特征表明哪些类别有可能在图像中同时出现。此外,[40]还表明,全局上下文提供了更广阔的视野,对于语义分割是有益的。全局上下文特征对这些区域进行整体编码,而不是为图像的每个部分独立学习重新加权。然而,对于以更精细的尺度对上下文进行编码的研究还很少,这是需要的,因为同一图像的不同区域可能包含完全不同的环境。

        为了实现这一目标,我们设计了一个“挤压-注意力”(SA)模块,通过一种重新加权机制来考虑局部和全局方面,为语义分割任务学习更具代表性的特征。SA模块扩展了SE模块的重新加权通道,如图2(b)所示,同时保留了空间信息,以适应场景解析的需求。因此,如图2(c)所示,提出了一个简单的挤压-注意力模块,可以表示为:        其中,Xattn = Up (σ (Xhattn)),Up(·)是一个上采样函数,用于扩展注意力通道的输出:        其中,Xˆattn表示注意力卷积通道Fattn (·)的输出,该通道由Θattn和注意力卷积层的结构Ωattn参数化。一个平均池化层AP ool(·)被用来执行不完全挤压操作,然后注意力通道的输出Xˆattn被上采样以匹配主卷积通道Xres的输出。这样,SA模块扩展了SE模块并保留了空间信息,注意力通道的上采样输出Xattn在主通道上聚合了非局部提取的特征。

3.3.压缩和注意力网络

        我们在SA模块的基础上构建了一个用于语义分割的SA网络(SANet)。具体而言,我们使用SA模块作为头部从骨干网络的四个阶段提取特征,以充分利用它们的多尺度特性。如图3所示,总损失包括三个部分:密集损失(交叉熵损失)、掩膜损失(交叉熵损失)和分类损失(二进制交叉熵损失)。“ynj”是Yden的平均池化结果。因此,SANet的总损失可以表示为:其中,α和β是分类损失和辅助损失的权重参数。总损失的每个组成部分可以表示如下:Squeeze-and-Attention Networks for Semantic Segmentation_第4张图片其中,N是每个epoch的训练数据大小,M表示空间位置的数量,C表示数据集的类别数。Yˆnij和Ynij是SANets的预测结果和真实标签,yˆnj和ynj是用于计算分类损失Lcat的分类预测和目标。Lcat采用二元交叉熵形式。Lmask和Lden是典型的交叉熵损失。辅助头部类似于深度监督策略,但其输入来自ResNet的第四阶段,而不是常用的第三阶段。SANets的预测集成了逐像素的预测,并通过第四个SA特征图进行正则化。因此,SANet的正则化密集分割预测是Yˆden + YˆSA4。SA网络的骨干以膨胀的FCN作为背景。假设输入图像的大小为3×512×512。SA模块的主通道具有与其注意力对应部分相同的通道数,并且与输入特征具有相同的空间大小。经验上,我们将输入的通道大小在主通道和注意力通道中都减小到四分之一,在注意力通道中设置下采样(最大池化)和上采样的比例为8,并在两个数据集的SE模块的中间全连接层中设置通道数为4。我们使用两个组进行主通道和注意力通道的第一个卷积操作。此外,我们将SA头部的输出调整为分割数据集的类别数。Squeeze-and-Attention Networks for Semantic Segmentation_第5张图片图4:对α和β的削弱研究,它们分别加权类别损失和稠密预测损失。我们使用ResNet50作为背骨架构对SANet进行测试,每种情况训练20个时期。左侧:具有固定β=0.8的SANet的mIoU。右侧:具有固定α=0.2的SANet的mIoU。Squeeze-and-Attention Networks for Semantic Segmentation_第6张图片图5:PASCAL Context验证集上的样本语义分割结果。PASCAL VOC验证集上的语义分割结果示例。(a)原始图像。(b)地面真实图像。(c)FCN基准结果。(d)SANet结果。SANet生成更准确的结果,特别是对物体边界更准确。最后一行显示了一个具有相对复杂背景的失败案例,这给分割模型带来了挑战。表1:在PASCAL Context数据集上对SANet进行的削弱研究结果(没有背景的59个类别)。SA:压缩注意力头部。Cat:类别损失。Den:稠密预测损失。PAcc:像素准确性(%)。mIoU:平均交叉联合(%)。Squeeze-and-Attention Networks for Semantic Segmentation_第7张图片

4.实验结果

        在本节中,我们首先将SA模块与SE模块进行比较,然后使用PASCAL Context [28]数据集进行消融研究,以测试总训练损失的每个组件的有效性,并在具有挑战性的PASCAL VOC数据集 [12]上验证SANets。按照场景解析的惯例[5,40],我们在PASCAL Context上同时评估均值交并比(mIoU)和像素精度(PAcc),并且仅在PASCAL VOC数据集上评估mIoU,以评估分割模型的有效性。

4.1.实现

        我们使用Pytorch [30]实现了SANets并进行了消融研究。对于训练过程,我们采用了与先前研究 [41,40]相同的多项式学习率衰减策略。PASCAL Context和PASCAL VOC的初始学习率分别为0.001和0.0001。对于两个数据集,我们采用随机梯度下降和多项式学习率退火策略。对于PASCAL Context数据集,我们训练SANets进行80个epoch。至于PASCAL VOC数据集,我们在COCO数据集上预训练模型,然后在验证集上训练网络50个epoch。我们将ResNet50和ResNet101作为SANets的骨干网络,因为这些网络已被广泛用于主流分割基准测试。对于所有训练情况,我们将批次大小设置为16,并最近根据[40]的实现,在多个GPU上使用同步批次归一化。我们连接四个SA头部的输出以利用不同阶段的多尺度特征,并对深度网络的训练进行正则化。

表2:在带有背景的PASCAL Context数据集上的平均交叉联合(%)结果(60个类别)。Squeeze-and-Attention Networks for Semantic Segmentation_第8张图片表3:基于PASCAL Context数据集上的基线膨胀FCN、带有SE模块的膨胀FCN(FCN-SE)和使用ResNet50或ResNet101作为骨架的SANet的像素准确性(PAcc)和mIoU。SANet显著优于其SE对应物和基线模型。每个模型训练20个时期。Squeeze-and-Attention Networks for Semantic Segmentation_第9张图片

4.2.PASCAL Context数据集上的结果

        Pascal Context数据集共包含59个类别、4998个训练图像和5105个测试图像。由于该数据集相对较小,我们将其作为基准来设计模块架构并选择超参数,包括α和β。为了进行消融研究,我们探索了SA模块的每个组成部分对于改善SANets分割结果的影响。消融研究包括三个部分。首先,我们测试了总训练损失中权重α和β的影响。如图4所示,我们将α从0到1.0进行测试,发现α=0.2时的SANet效果最好。类似地,我们将α固定为0.2,发现β=0.8可以产生最佳的分割性能。其次,我们使用选择的超参数来研究方程(7)中分类损失和密集预测损失的影响。表1显示,使用ResNet50作为骨干的包含四个双重使用的SA模块的SANet相比于FCN基线显著改善了分割性能(PAcc提高了2.7%,mIoU提高了6.0%)。此外,分类损失和辅助损失也提升了分割性能。

        我们将SANets与最先进的模型进行比较,以验证其有效性。如表2所示,使用ResNet101作为骨干的SANet达到了53.0%的mIoU。当包括背景类时,mIoU为52.1%,超过了其他竞争对手。此外,我们使用最近发布的EfficientNet(EffNet)[32]作为骨干。然后,EffNet版本的SANet在PASCAL Context数据集上实现了54.4%的mIoU,刷新了纪录。图5展示了使用相同骨干的扩张ResNet50 FCN和SANet的分割结果。在前三行中,SANets生成了更好的目标边界和更高的分割准确性。然而,对于像最后一行这样复杂的图像,两个模型都无法生成干净的解析结果。总体而言,定性评估与定量结果一致。我们还通过将SA模块与基线的扩张FCNs(包括ResNet50和ResNet101)上的SE模块进行比较,验证了SA模块的有效性。如表3所示,在两种设置下,SANet取得了最佳的准确率,分割结果显著提升(mIoU提高了4.1%和4.5%),而FCN-SE模型几乎没有改善分割结果。

4.3.注意i和特征图

        经典卷积已经通过每个卷积核在输入特征图上的空间位置进行扫描,产生了固有的全局注意力。因此,为了更好地理解SA模块中注意力通道的效果,我们可视化了PASCAL VOC数据集中一个示例的注意力和特征图,并在SANet中的Head1和Head4之间进行了比较。我们使用L2距离来显示SA模块内的注意力通道的注意力图,并选择最激活的特征图通道作为同一SA模块内主通道的输出。SA模块输出特征图的激活区域(红色)可以视为所选点的像素组。为了可视化,我们将图6中的所有特征图缩放到相同大小。在这个示例中,我们选择了三个点(红色,蓝色和品红色)来表明注意力通道强调像素组的注意力,这与SA模块的主通道专注于像素级预测相辅相成。有趣的是,如图6所示,低级(SA head1)和高级(SA head4)的注意力通道发挥着不同的作用。对于低级阶段,注意力通道的注意力图具有广泛的视场,主通道的特征图专注于保留对象边界的局部特征提取。相反,对于高级阶段,注意力通道的注意力图主要关注所选点周围的区域,主通道的特征图比head1更加均匀,语义含义更清晰。Squeeze-and-Attention Networks for Semantic Segmentation_第10张图片图6:在PASCAL VOC数据集上对经过训练的SANet的SA头部1和头部4的注意力和特征图可视化。对于每个头部,展示了主通道、注意力通道和输出的特征图。(a)原始图像及其地面真实值;(b)蓝色点的像素组可视化;(c)黄色点的像素组可视化;(d)洋红点的像素组可视化。Squeeze-and-Attention Networks for Semantic Segmentation_第11张图片图7:在PASCAL VOC验证集上的样本语义分割结果。 (a)原始图像。 (b)地面真实图像。 (c)FCN基准结果。 (d)一个SANet。 SANet相对于基线生成更准确的分割结果。

4.4.PASCAL VOC上的结果

        PASCAL VOC数据集[12]是最广泛研究的分割基准数据集,包含20个类别,由10582张训练图像、1449张验证图像和1456张测试图像组成。我们像之前的工作[26,10]一样,使用增强数据对SANet进行了80个epoch的训练。首先,我们在没有COCO预训练的情况下测试了SANet。如表4所示,SANet达到了83.2%的mIoU,比竞争对手更高,并在多个类别上占据主导地位,包括飞机、椅子、奶牛、桌子、狗、植物、绵羊和电视监视器。这一结果验证了双用途SA模块的有效性。模型[9,6]使用除PASCAL VOC或COCO之外的其他数据集,比如JFT [31],不包含在表4中。然后,我们测试了带有COCO预训练的SANet。如表5所示,使用COCO数据进行预训练的SANet达到了86.1%的mIoU,与包括PSPNet [41]在内的排名靠前的模型相媲美,并且优于基于重型ResNet152骨干的RefineNet [22]。我们的SA模块比EncNet [40]的编码模块计算效率更高。如图6所示,与基线模型相比,SANet的预测结果边界更清晰,定性结果更好。

4.5.复杂度分析

        我们的目标不是在不考虑计算成本的情况下追求SOTA,而是设计受到这种直觉启发的用于分割的轻量级模块。我们使用MACs和模型参数来分析SANet的复杂性。如表6所示,Deeplab V3+(我们的实现)和SAN都使用ResNet101骨干,在PASCAL VOC数据集上进行评估,以进行公平比较。在没有使用COCO数据集进行预训练的情况下,我们的SANet的mIoU增加了1.7%,超过了Deeplab V3+。与SDN(238.5M参数)等重型模型相比,SANet以不到四分之一的参数(55.5M参数)稍微表现较差。比较结果表明,SANet既有效又高效。

表4:在没有在COCO数据集上进行预训练的情况下,PASCAL VOC数据集上各类别的IoU和mIoU。SANet的mIoU达到了83.2%,超过了其他模型,并在多个类别上占据主导地位。每列的最佳两个结果被突出显示。为了公平比较,使用了额外数据集(如JFT)的模型没有被包括在内。Squeeze-and-Attention Networks for Semantic Segmentation_第12张图片表5:在在COCO数据集上进行预训练的情况下,PASCAL VOC数据集上的平均交叉联合(%)结果。SANet达到了86.1%的mIoU,与最先进的模型相当。Squeeze-and-Attention Networks for Semantic Segmentation_第13张图片

表6:使用ResNet101作为骨架对PASCAL VOC测试集进行评估的mIoUs(%)、每秒的乘加操作(MACs)和网络参数(Params)。我们重新使用膨胀ResNet101实现Deeplab V3+作为其骨干,以进行公平比较。Squeeze-and-Attention Networks for Semantic Segmentation_第14张图片

5.总结

        在本文中,我们从像素级预测和像素分组这两个独立的维度重新思考了语义分割问题。我们设计了一个SA模块来解决像素分组的隐含子任务。SA模块增强了像素级的密集预测,并考虑了被大多数人忽视的像素组注意力。更重要的是,我们提出了在两个具有挑战性的基准测试上取得有希望的分割性能的SANet。我们希望简单而有效的SA模块以及构建在SA模块之上的SANet能够促进其他研究团队在分割领域的研究。

你可能感兴趣的:(深度学习,人工智能,网络,目标检测,计算机视觉)