全监督语义分割训练新范式「像素对比学习」，苏黎世联邦理工等出品

机器之心专栏

机器之心编辑部

在这篇论文中，研究者提出了一种新的、全监督语义分割训练范式「像素对比学习」，强调利用训练集中、跨图像的像素 - 像素对应关系来学习一个结构化的特征空间，用来替代传统的、基于图像的训练范式。

• 论文地址：https://arxiv.org/pdf/2101.11939.pdf
• 代码地址：https://github.com/tfzhou/ContrastiveSeg

当前，语义分割算法的本质是通过深度神经网络将图像像素映射到一个高度非线性的特征空间。然而，现有算法大多只关注于局部上下文信息（单个图像内、像素之间的位置和语义依赖性），却忽略了训练数据集的全局上下文信息（跨图像的、像素之间的语义相关性），因而难以从整体的角度对习得的特征空间进行约束，进而限制了语义分割模型的性能。

最近，苏黎世联邦理工学院及商汤研究院的研究者提出了一种新的、全监督语义分割训练范式：像素对比学习（pixel-wise contrastive learning），强调利用训练集中、跨图像的像素 - 像素对应关系（cross-image pixel-to-pixel relation）来学习一个结构化（well structured）的特征空间，用来替代传统的、基于图像的（image-wise）训练范式。

该训练策略可直接应用于主流的语义分割模型，并在模型推理阶段不引入额外计算开销。下图展示了在 Cityscapes 验证集上主流分割算法的性能，可以看出，在 DeepLabV3、HRNet、OCR 上引入像素对比学习后，取得了较为显著的性能提升。

当前语义分割领域忽略了什么问题？

图像语义分割旨在为图像中的每个像素预测一个语义标签，是计算机视觉领域的一个核心问题。自 Fully Convolutional Network（FCN）[1] 提出以后，主流的语义分割算法强调图像内部的上下文信息（intra-image context）。主要从两方面入手: 1) 提出不同的上下文聚合模块（context aggreation module），如 dilated convolution, spatial pyramid pooling, encoder-decoder 及 non-local attention 等经典模型，核心思想是利用额外的模型参数或特殊的操作来建模提取图像内部的上下文信息；2）传统算法将语义分割视为一个像素级的分类任务，因而逐像素独立计算交叉熵损失（cross-entropy loss），但完全忽略了像素间的依赖关系。因而有研究者提出结构化的（structure-aware）损失函数，如 pixel affinity loss [2]，lovasz loss [3] 等，直接在训练目标函数中对分割结果的整体结构信息进行显式约束。

然而，以上工作都只关注图像内部的上下文信息，却忽略了跨图像的、全局的上下文信息：在训练集中，来自不同图像的像素之间也具有极强的相关性，如图 1（b），相同颜色的像素表示他们具有相同的语义。

图1：语义特征空间学习示意图。

更进一步，当前语义分割算法的本质是通过深度神经网络，将图像像素映射到一个高度非线性的特征空间（如图 1（c）），在这个过程中，只利用了上下文聚合模块或结构化的损失函数来强调局部的像素间的依赖关系，但忽略了一个本质问题：一个理想的语义分割特征空间究竟是什么样的？

研究者认为，一个好的分割特征空间应同时具备两个性质：

1. 强判别能力：在该特征空间中，每个像素的特征应具有较强的分类能力（strong categorization ability of individual pixel embeddings）；
2. 高度结构化：同类像素的特征应高度紧致（intra-class compactness），不同类像素的特征尽量分散（inter-class dispersion）。

然而，当前语义分割方法一般只关注性质 1，却忽略了 2。此外，很多表征学习（representation learning）的工作 [4, 5] 也验证了：通过强调性质 2，有助于更好的增强性质 1。因此我们大胆假设，尽管当前的语义分割算法已经取得了极佳性能，但是通过同时考虑性质 1 和 2，有可能习得一个更好的、结构化的分割特征空间，进而进一步提高语义分割算法的性能。

无监督对比学习引发的思考

近年来，无监督学习领域迎来了巨大发展，源头是对比学习（contrastive learning）[6, 7] 在海量无标注训练样本下的成功应用。假设

是一张没有标注的训练样本图像 I 的特征向量，

为图像 I 的一个正样本（positive sample）特征，这个正样本往往是对 I 施以某种变换得到的（如 flipping、cropping 操作等），

为一个负样本（negative sample）特征，训练集中其它非 I 的图像均被视为负样本。之后通过对比学习损失函数，如下面的 InfoNCE loss [8] ，进行无监督训练：

目标是将正样本从大量负样本中识别出来。无监督训练得到的图像特征表现出了极强的泛化能力，能够为下游任务提供极佳的网络初始化权重，或仅经过少量有标签样本的微调（finetuning），即可获得接近全监督训练性能的图像分类模型。

无监督对比学习的成功也带来了启发，对比学习属于度量学习（metric learning），本质是利用了数据集的整体信息，学习了一个具有极强表达能力的图像表征空间。在图像语义分割的全监督训练条件下，训练图像的每个像素的标签已经给出，我们可以将正样本视为同属于相同语义类别的像素，将负样本视为不属于同一语义类别的像素，而不论它们是否来源于同一个训练图像。之后即可以利用度量学习或对比学习，对传统的交叉熵损失加以改进，进而挖掘所有训练图像中、像素与像素之间的全局语义关系，进而获得一个高度结构化的分割特征空间，从而同时强调性质 1 和 2。由此研究者提出了一个基于像素对比学习的、全监督的、语义分割训练范式，像素对比学习（pixel-wise contrastive learning），强调利用训练数据集的全局上下文信息，从整体的角度对习得的特征空间进行显式约束，使其在全局结构上具备良好的性质（intra-class compactness 和 inter-class dispersion）。

如图 1（d）所示，给定一个训练样本中的像素 i，也称之为锚点（anchor point），研究者将 i 与其它像素在分割特征空间上进行比较，尽可能拉近 i 与其同类的像素（正样本）的距离，并迫使 i 尽可能远离其它不同类像素（负样本）。因此该训练范式能够考虑整个训练集中所有像素的全局语义相似度，使得模型能够利用更具多样性的、大规模样本提升表征学习的能力，从而获得更优的语义特征空间（如图 1（e））。

基于 Pixel-Wise 交叉熵的经典语义分割损失函数有何问题？

下面以语义分割领域经典的 Pixel-wise 交叉熵作为出发点，进一步探讨将度量学习或对比学习引入语义分割训练的必要性。

如前所述，当前语义分割算法将该任务视作一个逐像素的语义分类问题，即为图像中的每个像素 i 预测一个语义标签 c。因此使用 pixel-wise 交叉熵作为训练目标：

这里 y 表示：通过 FCN 得到对像素 i 的未归一化的分类预测（unnormalized categorical score vector），也称为 logit，

表示像素 i 的真实标签，

为的 one-hot 编码。

然而，该优化目标函数有两个缺点：

1. 只对每个像素的预测独立地进行约束，而忽略了像素之间的关系 [9] ；
2. 由于使用了 softmax 操作，交叉熵的计算实际只依赖于 logits 之间的相对关系，却无法直接约束习得的像素特征（cannot directly supervise on the learned representations） [10] 。

尽管最近的一些结构化损失函数（如 pixel affinity loss， lovasz loss 等）意识到了缺点 1，但也只考虑了同一个图像内部的像素依赖关系，却忽略了不同图像间像素的语义一致性。而缺点 2，在语义分割领域则鲜有提及。

基于全监督、像素 - 像素对比学习的语义分割训练范式

图2：基于全监督、像素对比学习的语义分割模型。

本文提出的像素对比学习（pixel-wise contrastive learning）能够较好的解决交叉熵损失函数的两个缺点。在训练过程中，对于任意像素（锚点）i, 其正样本为其它与之同类的像素，负样本为其它与之不同类的像素。值得注意的是，对于锚点 i 的正负样本的选择不限于同一张图像。对像素 i，对比损失函数定义为：

这里

表示像素 i 的所有正样本像素的特征，

表示像素 i 的所有负样本像素的特征。由上式可以看出，通过像素 - 像素对比学习，研究者直接在语义分割的特征空间上，将同属于一个语义类别的像素拉近，同时迫使不同语义类别的像素互相远离，从而同时强调了交叉熵损失的两个缺点。

最终的语义分割损失函数定义为：

交叉熵损失促使分割模型学习具有判别力的特征、提升分类能力（强调性质 1），pixel-wise contrastive loss 通过探索像素之间的全局语义关系，从整体上约束语义分割特征空间（强调性质 2）。

下图对只使用交叉熵损失（左图）与上式混合损失函数（右图）学习到的分割特征进行了可视化，可以看出，通过引入 pixel-wise contrastive loss，同类像素特征更加紧致，而类与类之间可以更好的分离。这表明，通过结合一元（unary）交叉熵损失和二元（pari-wise）contrastive loss 的优势，分割网络能够学习到更好的特征表示。

图3：分割特征可视化。左图：交叉熵损失；右图：交叉熵 + 像素对比损失。

基于以上框架，研究者进一步探索了像素 - 区域对比学习及难例挖掘，并得出了若干有益结论，具体内容请移步原文。

实验结果

为验证所提方案的性能，研究者在三个标准的语义分割数据集（Cityscapes，PASCAL-Context，COCO-Stuff）上进行了实验。以目前主流的语义分割模型（HRNet，OCR）为基础，引入像素对比学习后的性能对比如下：

实验结果表明，在三个数据集上，本文提出的训练范式可以明显地提升现有语义分割算法的性能。此外，该算法在模型部署阶段，不会引入任何额外的计算开销，不改变分割网络结构，具有较高的灵活性。

进一步探讨

不同于当前主流算法只关注图像内部像素的局部上下文信息，本文提出跨图像、像素对比损失函数，来挖掘训练数据集中所有像素的全局关系，有效提升了语义分割的性能。这有助于我们重新思考当前主流的训练范式，不仅仅只关注训练样本自身的特性，同时从全局的角度，关注训练样本之间关系。

本文也带来了一些有益的启示，如：

1. 对比学习或度量学习依赖于正负样本的质量，更智能的采样策略能够帮助分割网络更快速有效的学习。
2. 从度量学习的角度，交叉熵损失为一元损失函数（unary loss）, 而对比损失为二元损失函数（pair-wise loss），探索高阶的度量损失函数有可能带来更大的提升。
3. 对比损失在计算中需要对正负样本采样，有可能借此更自然地实现训练中的类别再均衡（class rebalance）。
4. 本文方案在主流语义分割数据集上取得了有效的性能提升，并且有望在其它图像稠密预测任务中（如 2D 人体姿态估计，医疗图像分割等）发挥优势。

参考文献

[1] Jonathan Long, Evan Shelhamer, and Trevor Darrell. Fully convolutional networks for semantic segmentation. In CVPR, 2015.

[2] Tsung-Wei Ke, Jyh-Jing Hwang, Ziwei Liu, and Stella X Yu. Adaptive affinity fields for semantic segmentation. In ECCV, 2018

[3] Maxim Berman, Amal Rannen Triki, and Matthew B Blaschko. The lovasz-softmax loss: A tractable surrogate for the optimization of the intersection-over-union measure in neural networks. In CVPR, 2018.

[4] Weiyang Liu, Yandong Wen, Zhiding Yu, and Meng Yang. Large-margin softmax loss for convolutional neural networks. In ICML, 2016.

[5] Florian Schroff, Dmitry Kalenichenko, and James Philbin. Facenet: A unified embedding for face recognition and clustering. In CVPR, 2015.

[6] Alexey Dosovitskiy, Jost Tobias Springenberg, Martin Riedmiller, and Thomas Brox. Discriminative unsupervised feature learning with convolutional neural networks. In NeurIPS, 2014.

[7] Ting Chen, Simon Kornblith, Mohammad Norouzi, and Geoffrey Hinton. A simple framework for contrastive learning of visual representations. ICML, 2020.

[8] Michael Gutmann and Aapo Hyva ̈rinen. Noise-contrastive estimation: A new estimation principle for unnormalized statistical models. In AISTATS, 2010.

[9] Shuai Zhao, Yang Wang, Zheng Yang, and Deng Cai. Region mutual information loss for semantic segmentation. In NeurIPS, 2019.

[10] Tianyu Pang, Kun Xu, Yinpeng Dong, Chao Du, Ning Chen, and Jun Zhu. Rethinking softmax cross-entropy loss for adversarial robustness. In ICLR, 2020.

© THE END

转载请联系本公众号获得授权