Semi-Supervised Semantic Segmentation With Pixel-Level Contrastive Learning From a Class-Wise Memory Bank
通过Class-Wise Memory Bank进行像素级对比学习的半监督语义分割
Paper:https://openaccess.thecvf.com/content/ICCV2021/html/Alonso_Semi-Supervised_Semantic_Segmentation_With_Pixel-Level_Contrastive_Learning_From_a_Class-Wise_ICCV_2021_paper.html
Code:https://github.com/Shathe/SemiSeg-Contrastive
Abstract
这项工作提出了一种半监督语义分割的新方法。这种方法的关键要素是我们的对比学习模块,它强制分割网络为整个数据集中的同类样本产生相似的像素级特征表示。为了实现这一目标,我们维护一个Memory Bank,该Memory Bank不断更新来自标记数据的相关且高质量的特征向量。在端到端训练中,标记和未标记数据的特征都被优化为与Memory Bank中的同类样本相似。我们的方法不仅优于当前半监督语义分割的最新技术,而且还优于众所周知的公共基准的半监督域适应,在最具挑战性的场景(即可用标记数据较少)上有更大的改进。代码可在 https://github.com/Shathe/SemiSeg-Contrastive 获取。
1 Introduction
语义分割的目标在于为图像中的每个像素分配语义类别标签。它是语义场景理解的一项重要计算机视觉任务,在医学成像 [30] 或自动驾驶 [2] 等许多应用中发挥着相关作用。与许多其他计算机视觉任务一样,深度卷积神经网络在语义分割方面表现出了显着的改进[2,19,1]。所有这些例子都遵循监督学习方法,需要大量带注释的数据才能很好地概括。然而,标记数据的可用性是监督学习中的常见瓶颈,特别是对于语义分割等任务,这需要繁琐且昂贵的每像素注释。
半监督学习假设只有一小部分可用数据被标记。它通过从未标记样本中提取知识来解决有限标记数据的问题。半监督学习已被应用于广泛的应用[37],包括语义分割[11,17,26]。以前的半监督分割工作主要基于每样本熵最小化[17,21,28]和每样本一致性正则化[11,36,28]。这些分割方法不会在学习的特征上强制执行任何类型的结构,以增加整个数据集的类间可分离性 。我们的假设是,克服这一限制可以带来更好的特征学习和性能,特别是当可用标记数据量较低时。
这项工作提出了一种半监督语义分割的新颖方法,遵循师生方案,其主要组成部分是新颖的表示学习模块(图 1)。该模块基于正向对比学习 [5, 14],并强制不同样本之间像素级特征的类可分离性。为了实现这一目标,教师网络仅根据标记数据生成候选特征,并将其存储在Memory Bank中。同时,学生网络学习从标记和未标记数据中生成类似的类别特征。存储在Memory Bank中的特征是根据其质量和学习到的对比优化的相关性来选择的。除了增加类间可分离性之外,该模块还强制在特征空间中对齐未标记和标记数据(Memory Bank),这是半监督语义分割中另一个未开发的想法。总之,我们提出了一种半监督语义分割的新颖框架,其主要贡献如下:
- 用于半监督语义分割的像素级对比学习方案,其中元素根据其相关性进行加权。
- 使用Memory Bank从标记数据中获取高质量像素级特征。
我们在著名的半监督语义分割基准上评估我们的方法,在不同的设置上达到最先进的水平。除此之外,我们表明我们的方法可以自然地解决半监督域适应任务,也获得最先进的结果。在所有情况下,对可比方法的改进随着未标记数据的百分比而增加。
2 Related Work
本节总结了半监督学习和对比学习的相关工作,特别强调与语义分割相关的工作。
2.1 半监督学习
Pseudo-Labeling
伪标签利用了通过现有模型 [21] 保留最有可能预测的类别来为未标记数据创建人工标签的想法 [24, 32]。伪标签的使用是由熵最小化驱动的[13],鼓励网络在未标记的数据上输出高度置信的概率。伪标签和直接熵最小化方法都常用于半监督场景[9,20,33,28],表现出良好的性能。我们的方法同时使用伪标签和直接熵最小化。
Consistency Regularization
一致性正则化依赖于这样的假设:模型应该对同一图像的扰动(例如数据增强)保持不变。这种正则化通常通过使用两种不同的方法来应用:分布对齐[3,31,35]或增强锚定[33]。虽然分布对齐强制扰动和未扰动样本的预测具有相同的类分布,但增强锚定强制它们具有相同的语义标签。为了在未标记数据上产生高质量的非扰动类别分布或预测,平均教师方法 [36] 提出了一种师生方案,其中教师网络是模型参数的指数移动平均值 (EMA),从而产生更稳健的预测。
2.2 半监督语义分割
半监督语义分割的一种常见方法是利用生成对抗网络(GAN)[12]。 Hung 等人[17]建议训练鉴别器来区分来自标记和未标记数据预测的置信图。 Mittal 等人 [26] 使用双分支方法,一个分支使用 GAN 方法强制执行低熵预测,另一个分支使用 Mean Teacher 方法 [35] 消除误报预测。 Feng等人[10]提出了类似的想法,最近的一项工作引入了动态互训练(DMT)。 DMT使用两个模型,模型的分歧用于重新加权损失。 DMT 方法还遵循 CBC [9] 的多阶段训练协议,其中伪标签以离线课程方式生成。其他工作基于用于一致性正则化的数据增强方法。 French等人[11]专注于应用CutOut[7]和CutMix[45],而Olsson等人[28]提出了一种专门用于语义分割的数据增强技术。
2.3 对比学习
对比学习[15]的核心思想是创建正负数据对,在特征空间中吸引正数据对并排斥负数据对。该技术已用于监督和自监督设置。然而,最近的自监督方法已经显示出与仅通过执行冗余减少[46]或通过蒸馏实现相似性最大化[5, 14]而使用正对的对比学习相似的性能水平。
至于语义分割,这些技术主要用作预训练[40,43,44]。最近,Wang 等人 [39] 展示了在像素和区域级别对同类监督样本应用标准对比学习的监督场景的改进。 Van 等人 [38] 展示了在无监督设置中对比学习的优势,将其应用于不同显着性掩模的特征之间。
在这项工作中,我们建议遵循基于相似性最大化和蒸馏的纯正对比学习 [5, 14]。通过这种方式,我们以比标准对比学习更简单、计算更高效的方式提高了半监督语义分割的性能[39]。与之前的工作不同,我们的对比学习模块解决了半监督场景,将标记和未标记数据中的类特征和每像素特征与存储在Memory Bank中的所有标记集中的特征进行对齐。与之前将图像级特征保存在Memory Bank中的对比学习工作 [42, 16] 相反,我们的Memory Bank保存不同语义类别的每像素特征。此外,由于所有数据集像素都没有无限的内存,因此我们建议仅保存最高质量的特征。
[16] Momentum Contrast for Unsupervised Visual Representation Learning
[42] Unsupervised Feature Learning via Non-Parametric Instance Discrimination
3 Method
半监督语义分割是一种逐像素分类任务,其中有两个不同的数据源可用:一小组标记样本 X l = { x l , y l } \mathcal{X}_l = \left\{ x_l , y_l \right\} Xl={xl,yl},其中 x l x_l xl 是图像, y l y_l yl 是其相应的注释,以及一大组未标记样本的数量 X u = x u \mathcal{X}_u = {x_u} Xu=xu。
为了解决这个任务,我们建议使用师生模式。教师网络 f ξ f_{\xi} fξ 从未标记的样本中创建强大的伪标签,并从标记的样本中创建Memory Bank条目,以教导学生网络 f θ f_\theta fθ 提高其分割性能。
Teacher-student scheme
学生网络 f θ f_\theta fθ 的学习权重 θ \theta θ 使用以下损失进行优化:
L s u p L_{sup} Lsup 是标记样本上的监督学习损失(第 3.1 节)。 L p s e u d o L_{pseudo} Lpseudo 和 L e n t L_{ent} Lent 分别处理伪标签(第 3.2 节)和熵最小化(第 3.3 节)技术,其中伪标签由教师网络 f ξ f_{\xi} fξ 生成。最后, L c o n t r L_{contr} Lcontr 是我们提出的纯正对比学习损失(第 3.4 节)。
教师网络 f ξ f_{\xi} fξ 的权重 ξ \xi ξ 是学生网络 f θ f_{\theta} fθ 的权重 θ \theta θ 的指数移动平均值,衰减率为 τ ∈ [ 0 , 1 ] \tau \in [0, 1] τ∈[0,1]。教师模型提供了更准确和稳健的预测[36]。因此,在每个训练步骤中,教师网络 f ξ f_{\xi} fξ 没有通过梯度下降进行优化,而是更新如下:
3.1 Supervised Segmentation: L s u p L_{sup} Lsup
我们的监督语义分割优化应用于标记数据 X l \mathcal{X_l} Xl ,遵循带有加权交叉熵损失的标准优化。令 H \mathcal{H} H 为 N N N 个每像素类概率分布 y 1 y_1 y1、 y 2 y_2 y2 的两个列表之间的加权交叉熵损失:
其中 C C C 是要分类的类数, N N N 是元素数,即 y 1 y_1 y1 中的像素数, α c \alpha^c αc 是每类权重, β n \beta^n βn 是每像素权重。 α c \alpha^c αc 和 β n \beta^n βn 的具体值详见 4.2 节。监督损失(参见图 2 的顶部)定义为:
其中 x l a x^a_l xla 是 x l x_l xl 的弱增强(增强细节请参见第 4.2 节)。
3.2 Learning from Pseudo-labels: L p s e u d o L_{pseudo} Lpseudo
半监督学习成功的关键是从未标记的数据中学习。我们的方法利用的一个想法是从伪标签中学习。伪标签由教师网络 f ξ f_{\xi} fξ 生成(见图 2)。对于每个未标记的样本 x u x_u xu,伪标签 y ^ u \hat{y}_u y^u 按照以下方程计算:
其中 f ξ f_{\xi} fξ 预测类别概率分布。请注意,伪标签是在每次训练迭代时计算的。
一致性正则化是通过增强锚定引入的,即为同一批次上的每个样本 x u x_u xu 计算不同的数据增强,帮助模型收敛到更好的解决方案[33]。未标记数据 X u \mathcal{X}_u Xu 的伪标签损失由交叉熵计算:
其中 x u a x^a_u xua 是 x u x_u xu 的强增强, A A A 是我们应用于样本 x u x_u xu 的增强数量(增强细节请参见第 4.2 节)。
3.3 Direct Entropy Minimization: L e n t L_{ent} Lent
将直接熵最小化应用于学生网络根据未标记样本 x u x_u xu 预测的类分布作为正则化损失:
其中 C C C 是要分类的类别数量, N N N 是像素数量, A A A 是增强数量。
Contrastive Learning: L c o n t r L_{contr} Lcontr
图 3 说明了我们提出的对比优化,其灵感来自于基于相似性最大化和蒸馏的纯正对比学习工作 [5, 14]。在我们的方法中,Memory Bank充满了来自教师 f ξ f_{\xi} fξ 的高质量特征向量(图 3 的右侧部分)。同时,学生 f θ f_{\theta} fθ 从 X l \mathcal{X}_l Xl 或 X u \mathcal{X}_u Xu中提取特征向量。以每个类的方式,每个特征都通过一个简单的自注意力模块,该模块充当对比损失中每个特征的权重。最后,损失强制来自学生的加权特征向量与来自Memory Bank的特征向量相似。由于Memory Bank包含来自所有标记样本的高质量特征,因此对比损失有助于在整个数据集的特征空间中创建更好的 类分离 ,并将未标记数据分布与标记数据分布对齐。
Optimization
令 f θ − f_{θ^-} fθ− 为没有分类层的学生网络, { x , y } \left\{x, y\right\} {x,y} 为来自 { X l , Y l } \left\{ \mathcal{X}_l , Y_l \right\} {Xl,Yl} 或 { X u , Y u } \left\{\mathcal{X}_u, Y_u \right\} {Xu,Yu} 的训练样本。第一步是提取所有特征向量: V = f θ − ( x ) V = f_{θ^−} (x) V=fθ−(x)。然后将特征向量 V V V 馈送到投影头 Z = g θ ( V ) Z = g_θ(V) Z=gθ(V) 和预测头 P = q θ ( Z ) P = q_θ(Z) P=qθ(Z),如下 [5, 14],其中 g θ g_θ gθ 和 q θ q_θ qθ 是两个不同的多层感知机(MLP)。接下来, P P P 按 y y y 中的不同语义类别进行分组。
令 P c = { p c } P_c = \left\{p_c\right\} Pc={pc} 为来自 c c c 类 P P P 的预测向量集。 Z c ′ = { z c ′ } Z^\prime_c = \left\{z^\prime_c\right\} Zc′={zc′} 表示教师从存储在Memory Bank中的标记示例中获得的 c c c 类投影向量的集合, Z ′ = g ξ ( f ξ − ( x ) ) Z^\prime = g_{\xi}(f_{\xi^−} (x)) Z′=gξ(fξ−(x))。
接下来,我们通过分配每个特征的学习权重(方程 8)来了解哪些特征向量( p c p_c pc 和 z c ′ z^{\prime}_c zc′ )对对比任务有益,该权重将用作对比损失函数(方程 11)的权重因子(方程 10)。这些每个特征的权重是使用特定于类的注意模块 S c , θ S_{c,\theta} Sc,θ(有关更多详细信息,请参阅第 4.2 节)计算的,该模块为每个 z c ′ z^{\prime}_c zc′ 和 p c p_c pc 特征生成单个值 ( w ∈ [ 0 , 1 ] w \in [0, 1] w∈[0,1])。根据[34],我们对这些权重进行 L 1 L_1 L1 标准化,以防止收敛到平凡的全零解。对于预测向量 P c P_c Pc 的情况,权重 w p c w_{p_c} wpc 计算如下:
其中 N P c N_{P_c} NPc 是 P c P_c Pc 中的元素数量。方程 8 也用于计算 w z c ′ w_{z^{\prime}_c} wzc′,将 Z c ′ Z^{\prime}_c Zc′ 和 z c ′ z^{\prime}_c zc′ 更改为 P c P_c Pc 和 p c ′ p^{\prime}_c pc′ 。
对比损失强制预测向量 p c p_c pc 与 [5, 14] 中提出的投影向量 z c ′ z^{\prime}_c zc′ 相似(在我们的例子中,投影向量位于Memory Bank中)。为此,我们使用余弦相似度作为相似度度量 C C C:
其中,预测和Memory Bank条目之间的加权距离计算如下:
并且,我们的对比损失计算如下:
Memory Bank
Memory Bank是维护每个类 c c c 的目标特征向量 z c ′ z^{\prime}_c zc′ 、 ψ \psi ψ 的数据结构,用于对比损失。由于没有无限的空间来保存标记数据的所有像素,因此我们建议仅存储来自具有最高质量的标记数据的特征向量的子集。如图 3 所示,每次训练迭代时都会使用教师生成的 z c ′ ∈ Z ′ z^{\prime}_c \in Z^\prime zc′∈Z′子集更新Memory Bank。为了选择 Z ′ 的哪个子集包含在Memory Bank中,我们首先执行特征质量过滤器 (FQF),其中我们仅保留在应用分类层时导致准确预测的特征, y = a r g m a x f ξ ( x l ) y = arg max f_{\xi}(x_l) y=argmaxfξ(xl),置信度高于阈值, f ξ ( x l ) > ϕ f_{\xi}(x_l) > \phi fξ(xl)>ϕ。剩余的 Z ′ Z^\prime Z′ 按类 Z c ′ Z^\prime_c Zc′ 分组。最后,我们不是随机选择每个 Z c ′ Z^\prime_c Zc′ 的子集来更新Memory Bank,而是使用特定于类的注意力模块 S c , ξ S_{c, \xi} Sc,ξ。我们得到排名分数 R c = S c , ξ ( Z c ′ ) R_c = S_{c, \xi}(Z^\prime_c ) Rc=Sc,ξ(Zc′) 对 Z c ′ Z^\prime_c Zc′ 进行排序,并且仅使用前 K K K 个得分最高的向量更新Memory Bank。Memory Bank是每个类的先进先出 (FIFO) 队列,以提高计算和时间效率。通过这种方式,它可以在计算和时间方面以非常高效的方式维护最新的高质量特征向量。有关超参数的详细信息包含在第 4.2 节中。
4 Experiments
本节描述评估中使用的数据集和实施细节。它还包含对我们的方法在半监督语义分割的不同基准上的评估,包括半监督域适应设置和详细的消融研究。
4.1 Datasets
Cityscapses [6]。它是一个真实的城市场景数据集,由 2975 个训练样本和 500 个验证样本组成,具有 19 个语义类别。
PASCAL VOC 2012 [8]。它是一个具有 21 个语义类别的自然场景数据集。该数据集分别有 10582 张和 1449 张图像用于训练和验证。
GTA5 [29]。它是从视频游戏中捕获的合成数据集,具有逼真的城市场景,总共有 24966 张图像。原始数据集提供了 33 个不同的类别,但根据 [41],我们仅使用与 Cityscapes 共享的 19 个类别。
4.2 Implementation details
Architecture
我们在实验中使用 DeepLab 网络 [4]。对于消融研究和大多数基准测试实验,使用具有 ResNet-101 主干的 DeepLabv2 与以前的作品进行公平比较(即相似的设置)[28,9,17,26]。具有 Resnet50 主干的 DeepLabv3+ 也用于与 [25] 进行同等比较。在(公式 2)的训练过程中,τ 设置为 0.995 到 1。
预测和投影头遵循 [14]:Linear → BatchNorm [18] → Relu [27] → Linear,隐藏维度和输出维度为 256。所提出的特定于类的注意模块遵循类似的架构:Linear → BatchNorm → LeakyRelu [23]→线性→Sigmoid,隐藏维度和输出维度分别为 256 和 1。我们使用 2 × N c l a s s e s 2 \times N_{classes} 2×Nclasses 注意力模块,因为它们是以类方式使用的。特别是,每个类使用两个模块,因为我们有不同的用于投影或预测特征向量的模块。
Data augmentation
我们使用两种不同的增强设置,一种针对标记样本的弱增强设置,一种针对未标记样本的强增强设置,遵循[28]并进行了较小的修改(表1描述了我们方法中的数据增强方案)。此外,我们设置 A = 2 A = 2 A=2(方程 6)作为每个样本的增强数量。
4.3 Benchmark Experiments
以下实验将我们的方法与不同半监督设置中最先进的方法进行比较,包括半监督域适应任务。
4.3.1 Semi-supervised Semantic Segmentation
Cityscapes、Pascal VOC
4.3.2 Semi-supervised Domain Adaptation
4.4 Ablation Study
以下实验研究了所提出方法的不同组成部分的影响。评估是在 Cityscapes 数据上完成的,因为与 Pascal VOC 相比,它提供了更复杂的场景。我们选择具有挑战性的标记数据比率 1 30 \frac{1}{30} 301 。
4.4.1 Loss Impact
Future direction
5 Conclusion
本文提出了一种半监督语义分割的新方法。我们的工作展示了结合纯正对比学习技术来解决这个半监督任务的好处。所提出的对比学习模块提高了这些设置中语义分割的性能。我们的新模块包含一个Memory Bank,该Memory Bank会根据教师网络根据标记数据生成的特征不断更新。这些特征是根据其质量和对比学习的相关性来选择的。我们的学生网络针对标记和未标记数据进行了优化,以学习与Memory Bank中类似的类特征。像素级对比学习的使用几乎没有得到开发,这项工作展示了它为半监督语义分割和半监督域适应带来的潜力和好处。我们的结果在几个公共基准上都优于最先进的结果,尤其是在更具挑战性的设置上,即当可用标记数据量较低时,有显着的改进。