基于一致性的半监督语义分割方法

在本工作中，来自阿德莱德大学、乌鲁姆大学的研究者针对当前一致性学习出现的三个问题做了针对性的处理， 使得经典的 teacher-student 架构 （A.K.A Mean-Teacher） 在半监督图像切割任务上得到了显著的提升。

该研究已被计算机视觉顶会 CVPR 2022 大会接收，论文标题为《Perturbed and Strict Mean Teachers for Semi-supervised Semantic Segmentation》：

背景

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

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

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

一致性学习的介绍

简单来说， 一致性学习（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），进一步增强了我们模型的泛化性。

审核编辑：彭静