Region-aware_Contrastive_Learning_for_Semantic_Segmentation

Hanzhe Hu1,2, Jinshi Cui1*, Liwei Wang1*
1北京大学电子工程与计算机科学学院,国家机器感知重点实验室
2浙江实验室

摘要

​  近期研究通过在单张图像中局部或全局地利用上下文信息,并采用像素级交叉熵损失监督模型,已在语义分割领域取得显著成果。然而从整个数据集的整体视角来看,语义关联不仅存在于单个图像内部,更普遍存在于整个训练数据中,这使得单纯考虑图像内部相关性已不足以满足需求。受无监督对比学习最新进展的启发,我们提出了一种区域感知对比学习方法(RegionContrast),为监督式语义分割提供创新解决方案。为提升语义相似像素的相似度并保持与其他像素的区分度,我们采用对比学习方法实现这一目标。借助记忆库技术,我们将探索将所有代表性特征存储到记忆库中。在不损失泛化能力的前提下,为高效整合训练数据至记忆库同时避免消耗过多计算资源,我们提出通过构建区域中心来表征每张图像中不同类别的特征。因此,我们提出的区域感知对比学习方法在训练阶段对所有数据均采用区域级处理,相比探索像素级关联的方法,这种设计能显著节省内存资源。该方法在训练过程中计算成本极低,测试时也无需额外开销。大量实验表明,我们的方法在Cityscapes、ADE20K和COCO Stuff三个基准数据集上均取得了业界领先的表现。

引言

​  语义分割是计算机视觉领域一个基础且具有挑战性的问题,其目标是对图像中的每个像素分配一个类别标签。该技术已广泛应用于自动驾驶、场景理解、图像编辑等多个领域。近年来,得益于ImageNet [11]和Cityscapes [10]等大规模数据集的普及,语义分割领域取得了显著进展。特别是基于全卷积网络(FCN)[30],涌现出许多前沿方法,这些方法专注于利用上下文信息。DeepLabV3 [5]提出了ASPP模块,该模块通过不同膨胀率对空间规律采样的像素进行聚合;而PSPNet [52]则提出了金字塔池化模块,在池化前将特征图划分为多个区域。非局部网络[40]采用自注意力机制,使每个像素都能接收来自所有其他像素的信息,从而获得更完整的像素级表征。
​  尽管上述方法在多数情况下能获得令人满意的分割效果,但仍存在明显缺陷。具体而言,当前主流方法大多侧重于从图像内部多维度挖掘上下文信息,却忽视了其他图像中潜在的关联信息。如图1所示,类间关系同样值得深入探索。对于图像中的某个区域而言,仅通过膨胀卷积或自注意力机制只能获取周边同类特征的部分信息,而现实中这类区域本应能接触到更多样化的特征信息。因此,单纯探索图像内部关系仍不够全面,这就催生了从整体数据集视角进行特征学习的需求。
​  近年来,无监督对比学习在预训练图像分类或目标检测等下游任务的强特征提取器方面备受关注。简而言之,大多数研究采用图像级对比学习:将数据集中所有其他图像视为负样本,而经过随机增强处理的原始图像则作为正样本。得益于记忆库的利用,大量负样本可被引入辅助对比学习,从而获得更优质的特征表征。值得注意的是,多数无监督对比学习方法聚焦于分类问题,而语义分割则需要比分类多得多的语义信息。直观来说,为适应分割任务,我们可以将图像级对比学习调整为像素级对比——即对原始图像内外的像素进行对比,如文献[41]所示。然而这种处理方式存在严重缺陷:不同图像的像素可能属于同一类别,这会削弱后续特征学习效果。因此,我们不再局限于无监督设置,而是探索完全监督下的对比学习方法,以获取丰富的类别信息。
​  在本研究中,我们提出了一种全新的全监督式对比学习范式,专门针对语义分割问题。通过将像素类别作为先验知识引入,对比学习得以实现更高效的执行。首先说明最直接的方法:当模型预测确定每个像素的类别时,会为每个类别构建对应的内存库来存储不同类别的像素嵌入数据。对于图像中的每个像素,其对应的正负样本均可从内存库中提取,从而完成对比学习过程。虽然这种方法简单有效,但由于整个数据集包含的像素数量庞大,会导致内存占用过高,进而显著拖慢训练速度。
​  为解决上述问题并恢复足够的嵌入特征,我们提出了区域感知对比学习(RegionContrast)。具体而言,由于图像中某一类别的区域特征由该类别所有像素特征构成,因此我们可以构建同一图像内不同类别的区域中心。通过这种方式,我们无需将所有像素嵌入推入记忆库,而是仅从不同类别中选取若干区域中心进行存储。虽然一张图像可能包含多个同类区域,但其嵌入空间中的特征具有相似性。因此为简化处理流程,我们为单张图像的每个类别生成一个区域中心。为促进难以分类像素的特征学习,我们在生成区域中心时提出动态采样方法,以增强对困难样本的关注度。在构建好不同类别的记忆库后,即可实施区域感知对比学习。具体而言,对于图像中的某个区域中心,其对应的正样本来自同类记忆库中的嵌入向量,而负样本则取自其他记忆库的嵌入向量。在获得正负样本后,即可执行对比学习流程。
​  我们的区域对比方法(RegionContrast)整体框架如图2所示。传统交叉熵损失作为像素级监督机制,而区域对比方法则专注于图像间关系学习。最关键的是,该方法可轻松适配各类分割模型,在训练时仅需少量计算资源,测试时也无需额外开销。
​  总之,我们的贡献总结如下:

  • 我们提出了一种全新的全监督式对比学习框架,专门针对特定的语义分割问题。
  • 为了以内存高效的方式适应分割场景,我们设计了一种有效的区域感知对比学习(RegionContrast),从整个数据集的整体视角探索语义关系。
  • 我们在多个公开数据集上进行了大量实验,并在Cityscapes、ADE20K和COCO Stuff等三个语义分割基准上取得了最先进的性能。

方法

image-20250904221906967

图2.我们提出的区域感知对比学习框架(RegionContrast)的整体架构。所有类别的记忆库均包含当前迭代前从图像生成的区域中心。对于当前迭代的输入图像,我们首先将其输入特征提取器(骨干+ ASPP模块)以获取其深度特征。通过所提出的动态采样方法,我们从图像中构建区域中心,并根据对应类别将其推入记忆库。在记忆库内执行区域感知对比学习时,红色箭头表示推送力,绿色箭头表示拉力(不同颜色的圆点代表不同类别的区域中心)。该模型采用交叉熵损失与所提出的RegionContrast进行联合监督。

1.区域感知对比学习

​  在本研究中,我们以监督学习的方式探索对比学习。得益于现有标注数据的优势,与以往方法不同,本研究将对比学习应用于类别层面而非实例层面,从而更有效地提升特征表示能力。

1.1.整体框架

​  如图2所示,我们展示了所提出的区域对比度(RegionContrast)的整体框架。我们选用DeepLabV3 [5]作为基础分割网络,并以在ImageNet数据集上预训练的ResNet作为主干网络。通过替换最后两个下采样操作,并在后续卷积层中引入膨胀卷积,有效提升了特征图的分辨率和感受野,使输出步幅从16倍优化至8倍。该模型采用常规像素级交叉熵损失与我们提出的区域感知对比度损失相结合的监督机制。
​  具体而言,要实现区域级别的对比学习,首先需要提取区域特征。给定输入图像I∈R C×H×W,我们将其通过主干网络和ASPP模块处理以获得特征图F∈R C×H×W。在神经网络预测的引导下,我们从特征图F中进一步提取区域特征——这通过在特征图上添加分割头来实现。为简化处理,我们选择用一个区域中心来表征图像中的每个类别,该中心编码了该类别的核心信息。实践中,生成的区域中心是类别i的向量Ri∈R C。随后,同类别的区域中心作为正样本,而不同类别的则作为负样本。本文提出的监督式区域感知对比学习的核心在于:在保持不同类别间区分能力的同时,将同一类别的区域特征进行整合。

1.2.区域中心

​  直观来说,类别i的区域中心可以定义为单个图像中所有属于该类别的像素特征的平均值。具体而言,给定特征图F∈R C×H×W(其中C、H和W分别表示特征图的维度、高度和宽度),类别i的区域中心的直接定义可表述为: \[R_{i}=\frac{\sum_{x,y}F_{(x,y)}{\mathbb 1}[L_{(x,y)}=i]}{\sum_{x,y}{\mathbb 1}[L_{(x,y)}=i]},\] ​  其中\(L_{(x,y)}\)是基础分割网络预测的像素标签,\({\mathbb 1}(\cdot)\)是表示该像素属于类别i的二元指示器。
​  然而,由于网络预测中存在误判,上述构建方法中的区域中心会覆盖像素的模糊特征,这将误导区域中心的学习过程。为了更关注难以分类的像素,我们进一步提出一种动态采样方法来构建区域中心。
​  除了特征图\(F\in\mathbb{R}^{C\times H\times W}\)和预测图\(P\in\mathbb{R}^{N\times H\times W}\)外,我们还引入了真实值图\(G\in\mathbb{R}^{H\times W}\)来挖掘困难像素(其中C表示特征维度,N为类别数)。在真实值的引导下,可以有效过滤困难负样本并获取困难正样本。为了更关注困难样本,我们在生成区域中心时应对不同像素特征赋予差异化权重——困难样本的像素需要比简单正样本获得更高权重。为此,我们利用预测的置信度图来为容易的正样本分配权重,其中每个位置的置信度图值\(c_{i,j}\in[0,1]\)。因此,简单正像素\((i,j)\)的权重为\(1-c_{i,j}\),而困难正像素的权重为1。从形式上讲,类别i的简单正样本可表示为\(E P_{i}=\sum_{x,y}{\mathbb 1}[L_{(x,y)}=i]\cap{\mathbb 1}[G_{(x,y)}=i]\),而困难正样本可表示为\(H P_{i}=\sum_{x,y} {\mathbb 1}[G_{(x,y)}=i] - {\mathbb 1}[L_{(x,y)}=i]\cap{\mathbb 1}[G_{(x,y)}=i]\)。类别i区域中心的最终定义可定义为, \[R_{i}=\frac{\sum_{x,y}F_{(x,y)}((1-c_{(x,y)})\cdot E P_{i}+H P_{i})}{\sum_{x,y}\,{\mathbb 1}[G_{(x,y)}=i]},\]

1.3.区域对比

​  在为每张图像构建所有类别的区域中心后,来自类别i的区域中心的区域感知对比损失可直接定义为, \[L_{i}^{N C E}=\frac{1}{|M_{i}|}\sum_{k_{+}\in M_{i}}-\log\frac{\exp(q\cdot k_{+}/\tau)}{\exp(q\cdot k_{+}/\tau)+\sum_{k_{-}}\exp(q\cdot k_{-}/\tau)},\] ​  其中Mi表示从整个训练集的类别i中收集区域中心的记忆库,而k−来自其他类别的记忆库。
​  包含负样本的记忆库对学习优质特征表征至关重要。先前在语义分割中应用对比学习的方法[41,54,39]需要维护大量像素级嵌入存储的内存库,这不仅导致对大容量内存库的迫切需求,还因内存负担过重而造成训练速度缓慢。相比之下,我们的方法通过引入区域中心点的优势,在训练过程中所需的内存资源显著减少。具体来说,对于包含N个类别中D张图像的训练集,我们会构建N个记忆库,每个记忆库的容量为D×C(其中C表示嵌入特征的维度,这些嵌入也被称为区域中心)。在训练过程中,我们会将这N个记忆库作为不同的队列进行维护。针对每个小批量数据,系统会生成不同类别的区域中心,并将其放入对应的队列中,在下一个训练周期内完成更新。
​  综上所述,语义分割的最终监督机制可概括为以下内容。我们采用了提出的区域感知对比损失函数,并结合传统像素级交叉熵损失与辅助损失,这与现有前沿研究[52,47,20]的做法一致。具体而言,将主干网络ResNet第三阶段的输出结果输入辅助层,通过同样采用交叉熵损失的辅助损失进行监督预测。简言之,该损失函数可表示为: \[L=L_{C E}+\alpha L_{A U X}+\beta L_{R C},\] ​  其中α和β用于平衡主分割损失LCE、辅助损失LAUX以及区域感知对比损失LRC。

结论

​  在本文中,我们提出了区域感知对比学习(RegionContrast),将对比学习引入语义分割问题。与以往的无监督对比学习方法不同,我们提出了一种全新的全监督式对比学习框架,专门针对图像分割任务。通过利用标签数据,我们能够深入挖掘更丰富的语义关联。此外,我们创新性地提出了不同类别区域中心的概念——这些区域中心被存储在内存中,并参与后续的对比学习过程。采用区域级嵌入而非像素级嵌入来存储整体训练集的信息,使得对比学习能够以内存高效的方式实现。消融实验验证了RegionContrast各组件的有效性。我们的RegionContrast模型在Cityscapes、ADE20K和COCO Stuff三个基准数据集上取得了业界领先的表现。