干货！全监督对比学习：语义分割新训练范式

点击蓝字

关注我们

AI TIME欢迎每一位AI爱好者的加入！

当前语义分割算法本质是通过深度神经网络，将图像像素映射到一个高维的特征空间。然而，当前算法只关注于局部上下文信息(单个图像内、像素之间的位置和语义依赖性)，却忽略了训练数据集的全局上下文信息(跨图像的、像素之间的语义相关性)，因而难以从整体的角度对习得的特征空间进行约束，进而限制了语义分割模型的性能。对此我们提出了一种全新的、全监督语义分割训练范式:像素对比学习(pixel-wise contrastive learning)，强调利用训练数据集中，跨图像的像素-像素对应关系(cross-image pixel-to-pixel relation)来学习一个结构化(well-structured)的特征空间，用来替代传统的、基于图像的(image-wise) 训练范式。该训练策略可直接应用于主流语义分割模型，带来显著的性能提升，并在模型推理阶段不引入额外计算开销。

本期AI TIME PhD直播间，我们邀请到苏黎世联邦理工学院博后研究院——王文冠，为我们带来报告分享《全监督对比学习：语义分割新训练范式》。

王文冠：

苏黎世联邦理工学院(ETH Zurich)博后研究员。主要研究方向为自动驾驶、语义分制、视频分析、以人为中心的场最解析和具身人工智能等。在国际权威期刊和会议(如TPAMI、TIP、TVCG、CVPR、ICCV、 ECCV、 AAAI、 Siggraph Asia)发表学术论文50多篇，谷歌学术引用7900+。曾获百度奖学金、ACM中国优博奖、中国人工智能学会优博奖、世界人工智能大会优秀青年论文奖，以及CVPR、 ICCV、 ECCV杰出审稿人，入选2020年Elsevier中国高被引学者榜单，获澳大利亚研究理事会青年职业奖。

01

 背   景 

语义分割任务是在图像上实现像素级的分类，自FCN（全卷积网络）提出后，一个发展的额主线是如何利用图像内的上下文信息。其中一类算法是通过采用某些特殊的神经网络架构，如空洞卷积（dilated convolution）、空间金字塔池化（SPP）、注意力机制（attention），来实现这一目的；另一个类算法是提出考虑图像内像素关系的、结构化训练损失函数（Structure-aware training objective），比如Pixel-affinity loss会对相邻像素之间分类预测不一致的情况进行惩罚，还有一些基于IoU的损失函数，比如lovász-softmax loss，这是常用的语义分割损失函数，交叉熵损失（cross-entropy loss）是逐像素计算的。

总而言之，现有方法大多只关注于局部上下文信息（单个图像内、像素之间的位置和语义依赖性），忽略了训练数据集的全局上下文信息（跨图像的、像素之间的语义相关性），因而难以从整体的角度对习得的特征空间进行约束，进而限制了语义分割模型的性能。所以我们第一个研究动机就是探讨如何利用跨图像的上下文信息来提高语义分割模型的性能。

另外，语义分割模型的本质上是利用深度神经网络将图像映射到高维的特征空间，之后将嵌入向量通过segmentation head得到不同类别的prediction score，然后进行分类预测。但这里忽略了一个本质问题——这个segmentation embedding space应该具备哪些性质？这就是我们的第二个研究动机。

我们认为一个理想的语义分割嵌入空间应具备以下性质：

1）单个像素的embedding具备良好的辨识度

2）整个空间结构合理（well-structured）—— 类内紧凑（属于同一个语义类别的像素embedding聚拢在一起）和类间分散（不属于同一个语义类别的像素embedding分开）

由此，本文提出了一个全新的语义分割训练范式。不同于以往的、image-wise的语义分割训练方法，我们的在语义分割训练中额外考虑了跨图像的语义相关性。通过考虑整个数据集的数据分布，约束整个语义分割嵌入空间，使得语义分割网络习得更好的像素embedding。

这里再简要介绍一下对比学习。对比学习在无监督的表征学习中得到了广泛运用。假设是一张没有标注的训练样本图像I的特征向量，v＋为图像I的一个正样本（positive sample）的特征，这个正样本往往是对I施以某种变换得到的（flipping、cropping操作等），v－为一个负样本的特征，训练集中其它非I的训练图像均被视为负样本（negative sample）。之后通过对比学习损失函数，如下面的InfoNCE loss进行无监督训练，目标是将正样本从海量负样本中识别出来：

02

 方   法 

在图像语义分割的全监督训练条件下，训练图像的每个像素的label已经给出，我们可以将正样本视为同属于相同语义类别的像素，将负样本视为不属于同一语义类别的像素，而不论它们是否来源于同一个训练图像，之后即可以利用度量学习或对比学习，对传统的cross-entropy loss加以补充，进而挖掘所有训练样本中像素的全局语义关系，进而获得一个高度结构化的分割特征空间。由此形成了一个基于像素对比学习的全新语义分割训练范式，pixel-wise contrastive learning，强调利用训练数据集的全局上下文信息，从整体的角度对习得的特征空间进行显式约束，使其在全局结构上具备良好的性质。

下面我们再从语义分割领域经典的Pixel-wise cross-entropy作为出发点，进一步探讨将对比学习引入语义分割训练的必要性。

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

这里y表示：通过FCN得到对像素i的未归一化的分类预测（unnormalized categorical score vector），也称为logit，表示像素i的真实标签,    为的one-hot编码。

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

1）它只对每个像素的预测独立地进行约束，而忽略了像素之间的关系；

2）由于使用了softmax操作，cross-entropy的计算实际只依赖于logits之间的相对关系，却无法直接约束习得的像素特征（cannot directly supervise on the learned representations）。尽管最近的一些结构化损失函数（如pixel affinity loss， lovasz loss等）意识到了缺点1，但只考虑了同一个图像内部的像素依赖关系，忽略了不同图像间像素的语义一致性。而缺点2，在语义分割领域，则鲜有提及。

下面的式子是我们提出的一个对比损失函数，其中和分别表示像素的正样本和负样本的pixel embedding集合。注意，正、负样本和像素不限于来自同一图像，也就是说在整个数据集中，只要与像素是同一类别标签的都属于正样本，否则都是负样本。因此克服了单图像交叉熵损失的缺陷，能够获取跨图像的上下文信息。

这种pixel2pixel的对比损失函数的目的是通过拉近同一类像素样本和分开不同类样本来优化pixel embedding space，学习到一个Structured Embedding Space。

由上式可以看出，通过像素-像素对比学习，我们直接在语义分割的特征空间上，将同属于一个语义类别的像素拉近，同时迫使不同语义类别的像素互相远离，克服了cross-entropy loss的两个缺点。总体的损失函数如下，我们同时考虑了传统的交叉熵损失和自己提出的对比损失函数。

为了进一步对对比损失函数进行分析，我们对其进行求梯度，得到下面的式子。表示正负样本与像素的匹配程度，匹配程度越高，值就越接近于1，匹配程度越低值越接近-1。从下式可以得出hard example mining(困难样本挖掘)对的计算非常重要，hard negative指的是与像素i相似度接近于1（趋向于1），但与i不属于同一类别的样本，hard positive指的是趋向于-1的样本。无论是hard negative还是hard positive都会对产生很大的影响，而主导训练优化的方向，因此我们需要更多地关注对比损失函数如何对正负样本进行采样。

为此，我们在算法中加入了hard example mining和hard anchor mining的机制。首先hard example mining机制使得在采样时加大hard positive和hard negative样本的权重，让它们有更大的概率被放入到中进行计算。hard anchor mining机制是作用于训练预测时各类别的预测分数都不太高（即未被正确分类的pixel）的pixel上，将这些pixel embedding以更大的概率作为anchor pixel放到进行计算。

另外，我们又提出了在两个不同的level上进行对比学习损失函数的计算，一个是pixel level的Pixel Rank，另一个是region level的Region Rank。Pixel Rank为每个类别维护一个队列，并存储从几个最新batch中采样的pixel embedding。Region Rank对每个图像同一类别的所有pixel embedding进行池化得到region level的embedding。

03

 实   验 

数据集：

CityScapes、CoCo、PASCAL、CamVid

Backbone：

ResNet、HRNet、Xception等等

实验结果：

下面展示了已有的语义分割算法加入本文的训练范式后的表现，使用本文的训练范式都能够有效提升模型的性能。

提

醒

论文链接：

https://arxiv.org/pdf/2101.11939.pdf

点击“阅读原文”，即可观看本场回放

整理：爱 国

审核：王文冠

实习推荐

往期精彩文章推荐

记得关注我们呀！每天都有新知识！

 关于AI TIME 

2019年，清华大学人工智能研究院院长张钹院士、唐杰教授和李涓子教授等人联合发起“AI TIME science debate”，希望用辩论的形式，探讨人工智能和人类未来之间的矛盾，探索人工智能领域的未来。

AI TIME是清华大学计算机系一群关注人工智能发展，并有思想情怀的青年学者创办的圈子。AI TIME旨在发扬科学思辨精神，邀请各界人士对人工智能理论、算法、场景、应用的本质问题进行探索，加强思想碰撞，链接全球AI学者、行业专家与爱好者，打造成为全球AI交流与知识分享的聚集地。

我知道你

在看

哦

~

点击 阅读原文 查看回放！