CVPR 2022 | PS-MT：半监督语义分割需要更稳定的一致性训练！

作者 | 机器学习社区  编辑 | 汽车人

原文链接：https://zhuanlan.zhihu.com/p/553558981

点击下方卡片，关注“自动驾驶之心”公众号

ADAS巨卷干货，即可获取

点击进入→自动驾驶之心技术交流群

后台回复【ECCV2022】获取ECCV2022所有自动驾驶方向论文！

后台回复【领域综述】获取自动驾驶全栈近80篇综述论文！

后台回复【数据集下载】获取计算机视觉近30种数据集！

在这里和大家分享一波我们CVPR 2022中稿的工作"Perturbed and Strict Mean Teachers for Semi-supervised Semantic Segmentation". 在本工作中，我们针对当前一致性学习(consistency learning) 出现的三个问题做了针对性的处理, 使得经典的teacher-student架构 (A.K.A Mean-Teacher) 在半监督图像分割任务上得到了显著的提升。

作者单位：澳洲机器学习研究院, 乌鲁姆大学

文章地址：https://arxiv.org/abs/2111.12903

代码地址：https://github.com/yyliu01/PS-MT

背景

语义分割是一项重要的像素级别分类任务。但是由于其非常依赖于数据的特性(data hungary), 模型的整体性能会因为数据集的大小而产生大幅度变化。同时, 相比于图像级别的标注, 针对图像切割的像素级标注会多花费十几倍的时间。因此, 在近些年来半监督图像切割得到了越来越多的关注。

半监督分割的任务依赖于一部分像素级标记图像和无标签图像 (通常来说无标签图像个数大于等于有标签个数)，其中两种类型的图像都遵从相同的数据分布。该任务的挑战之处在于如何从未标记的图像中提取额外且有用的训练信号，以使模型的训练能够加强自身的泛化能力。

在当前领域内有两个比较火热的研究方向, 分别是自监督训练(self-training) 和 一致性学习 (consistency learning)。我们的项目主要基于后者来进行。

一致性学习的介绍

简单来说, 一致性学习过程可以分为3步来描述: 1). 用不做数据增强的”简单”图像来给像素区域打上伪标签, 2). 用数据增强(或扰动)之后的”复杂”图片进行2次预测, 和 3). 用伪标签的结果来惩罚增强之后的结果。

可是, 为什么要进行这3步呢? 先用简单图像打标签, 复杂图像学习的意义在哪?

从细节来说, 如上图所示, 假设我们有一个像素的分类问题 (在此简化为2分类, 左下的三角和右上的圆圈) 。我们假设中间虚线为真实分布, 蓝色曲线为模型的判别边界。

在这个栗子中, 假设这个像素的标签是圆圈, 并且由1).得到的伪标签结果是正确的 (y_tilde=Circ.)。在2).中如果像素的增强或扰动可以让预测成三角类, 那么随着3)步骤的惩罚, 模型的判别边界会(顺着红色箭头)挪向真实分布。由此, 模型的泛化能力得到加强。

由此得出, 在1).中使用”简单”的样本更容易确保伪标签的正确性, 在2).时使用增强后的”复杂”样本来确保预测掉在边界的另一端来增强泛化能力。可是在实践中,

1). 没有经受过增强的样本也很可能被判断错 (hard samples), 导致模型在学习过程中打的伪标签正确性下降。

2). 随着训练的进行, 一般的图像增强将不能让模型做出错误判断。这时, 一致性学习的效率会大幅度下降。

3). 被广泛实用的半监督loss例如MSE, 在切割任务里不能给到足够的力量来有效的推动判别边界。而Cross-entropy很容易让模型过拟合错误标签, 造成认知偏差 (confirmation bias)。

针对这三个问题, 我们提出了:

1). 新的基于一致性的半监督语义分割 MT 模型。通过新引入的teacher模型提高未标记训练图像的分割精度。同时, 用置信加权 CE 损失 (Conf-CE) 代替 MT 的 MSE 损失，从而实现更强的收敛性和整体上更好的训练准确性。

2).一种结合输入、特征和网络扰动结合的数据增强方式，以提高模型的泛化能力。

3). 一种新型的特征扰动，称为 T-VAT。它基于Teacher模型的预测结果生成具有挑战性的对抗性噪声进一步加强了student模型的学习效率.

方法介绍

1). Dual-Teacher Architecture

我们的方法基于Mean-Teacher, 其中student的模型基于反向传播做正常训练。在每个iteration结束后, student模型内的参数以expotional moving average (EMA)的方式转移给teacher模型。

在我们的方法中, 我们使用了两个Teacher模型。在做伪标签时, 我们用两个teacher预测的结果做一个ensemble来进一步增强伪标签的稳定性。我们在每一个epoch的训练内只更新其中一个teacher模型的参数, 来增加两个teacher之间的diversity。

由于双teacher模型并没有参加到反向传播的运算中, 在每个iteration内他们只会消耗很小的运算成本来更新参数。

2). Semi-supervised Loss

在训练中, teacher模型的输出经过softmax后的置信度代表着它对对应伪标签的信心。置信度越高, 说明这个伪标签潜在的准确率可能会更高。在我们的模型中, 我们首先对同一张图两个teacher的预测取平均值。然后通过最后的confidence作为权重, 对student模型的输出做一个基于cross-entropy惩罚。同时, 我们会舍弃掉置信度过低的像素标签, 因为他们是噪音的可能性会更大。

3). Teacher-based Virtual Adversarial Training (T-VAT)

Virtual Adversarial Training (VAT) 是半监督学习中常用的添加扰动的方式, 它以部分反向传播的方式来寻找能最大化预测和伪标签距离的噪音。

在我们的模型中, dual-teacher的预测比学生的更加准确, 并且(由于EMA的更新方式使其)更加稳定。我们使用teacher模型替代student来寻找扰动性最强的对抗性噪音, 进而让student的预测出错的可能性加大, 最后达到增强一致性学习效率的目的。

4). 训练流程

i). supervised part: 我们用strong-augmentation后的图片通过cross-entropy来训练student 模型。

ii). unsupervised part: 我们首先喂给dual-teacher模型们一个weak-augmentation的图片, 并且用他们ensemble的结果生成标签。之后我们用strong-augmentation后的图片喂给student模型。在通过encoder之后, 我们用dual-teachers来通过T-VAT寻找具有最强扰动性的噪音并且注入到(student encoded之后的)特征图里, 并让其decoder来做最终预测。

iii). 我们通过dual-teachers的结果用conf-ce惩罚student的预测

iv). 基于student模型的内部参数, 以EMA的方式更新一个teacher模型。

实验

1). Compare with SOTAs.

Pascal VOC12 Dataset:

训练log可视化链接:

https://wandb.ai/pyedog1976/PS-MT(VOC12)?workspace=user-pyedog1976

该数据集包含超过 13,000 张图像和 21个类别。它提供了 1,464 张高质量标签的图像用于训练，1,449图像用于验证，1,456图像用于测试。我们follow以往的工作, 使了10582张低质量标签来做扩展学习, 并且使用了和相同的label id。

Low-quality Experiments

该实验从整个数据集中随机sample不同ratio的样本来当作训练集 (其中包含高质量和低质量两种标签), 旨在测试模型在有不同数量的标签时所展示的泛化能力。

在此实验中, 我们使用了DeeplabV3+当作架构, 并且用ResNet50 和 ResNet101得到了所有ratio的SOTA。

High-quality Experiments

该实验从数据集提供的高质量标签内随机挑取不同ratio的标签, 来测试模型在极少标签下的泛化能力。我们的模型在不同的架构下(e.g., Deeplabv3+ and PSPNet)都取得了最好的结果。

Cityscapes Dataset:

训练log可视化链接:

https://wandb.ai/pyedog1976/PS-MT(City)?workspace=user-pyedog1976

Cityscapes是城市驾驶场景数据集，其中包含 2,975 张训练图像、500张验证图像和 1,525 张测试图像。数据集中的每张图像的分辨率为 2,048 ×1,024，总共有 19个类别。

在2021年之前, 大多数方法用712x712 作为训练的resolution, 并且拿Cross-entropy当作supervised 的loss function。在最近, 越来越多的方式倾向于用大resolution (800x800)当作输入, OHEM当作supervised loss function。为了公平的对比之前的工作, 我们分别对两种setting做了单独的训练并且都拿到了SOTA的结果。

2). Ablation Learnings.

我们使用VOC数据集中1/8的ratio来进行消融实验。原本的MT我们依照之前的工作使用了MSE的loss方式。可以看到, conf-CE带来了接近3个点的巨大提升。在这之后, T-VAT (teacher-based virtual adversarial training)使student模型的一致性学习更有效率, 它对两个架构带来了接近1%的提升。最后, dual-teacher的架构给两个backbone分别带来了 0.83% 和0.84%的提升。

同时我们对比了多种针对feature的扰动的方法, 依次分别为不使用perturbation, 使用uniform sample的噪音, 使用原本的VAT和我们提出的T-VAT。T-VAT依然带来了最好的结果。

3). Improvements over Supervised Baseline.

我们的方法相较于相同架构但只使用label part的数据集的结果有了巨大提升。以Pascal VOC12为例, 在 1/16 的比率中 (即662 张标记图像), 我们的方法分别 (在ResNet50 和ResNet101中) 超过了基于全监督训练的结果 6.01% 和 5.97%。在其他ratio上，我们的方法也显示出一致的改进。

总结

在本文中，我们提出了一种新的基于一致性的半监督语义分割方法。在我们的贡献中，我们引入了一个新的 MT 模型，它基于多个teacher和一个student模型，它显示了对促进一致性学习的未标记图像更准确的预测，使我们能够使用比原始 MT 的 MSE 更严格的基于置信度的 CE 来增强一致性学习的效率。这种更准确的预测还使我们能够使用网络、特征和输入图像扰动的具有挑战性的组合，从而显示出更好的泛化性。

此外，我们提出了一种新的对抗性特征扰动 (T-VAT)，进一步增强了我们模型的泛化性。

【自动驾驶之心】全栈技术交流群

自动驾驶之心是首个自动驾驶开发者社区，聚焦目标检测、语义分割、全景分割、实例分割、关键点检测、车道线、目标跟踪、3D感知、多传感器融合、SLAM、高精地图、规划控制、AI模型部署落地等方向；

加入我们：自动驾驶之心技术交流群汇总！

自动驾驶之心【知识星球】

想要了解更多自动驾驶感知（分类、检测、分割、关键点、车道线、3D感知、多传感器融合、目标跟踪）、自动驾驶定位建图（SLAM、高精地图）、自动驾驶规划控制、领域技术方案、AI模型部署落地实战、行业动态、岗位发布，欢迎扫描下方二维码，加入自动驾驶之心知识星球（三天内无条件退款），日常分享论文+代码，这里汇聚行业和学术界大佬，前沿技术方向尽在掌握中，期待交流！