A Three-Stage Self-Training Framework for Semi-Supervised Semantic Segmentation

论文阅读笔记：

A Three-Stage Self-Training Framework for Semi-Supervised Semantic Segmentation

基本信息

\1.标题：A Three-Stage Self-Training Framework for Semi-Supervised Semantic Segmentation

\2.作者：\(Rihuan Ke^{*1} , Angelica Aviles-Rivero^{*1} , Saurabh Pandey^3 , Saikumar Reddy^4 , and\ Carola-Bibiane\ Schonlieb^1\)

\3.作者单位：\(^1Centre\ for\ Mathematical\ Sciences,\ University\ of\ Cambridge,\ Cambridge\ CB3\ 0WA,UK.^†\ ^2KritiKal Solutions^‡\)

\4.发表期刊/会议：CVPR

\5.发表时间：2020

\6.原文链接：https://arxiv.org/abs/2012.00827

Abstract

语义分割在社区中得到了广泛的研究，其中SOTA的技术是基于监督模型的。这些模型展现出了前所未有的性能，代价是需要大量高质量的分割mask。获取这样的标注是非常昂贵和耗时的，特别是在需要像素级标注的语义分割中。在这项工作中，我们提出了一个完整的解决方案，作为半监督语义分割的三阶段自训练框架。该技术的核心思想是提取伪标签统计信息，以减少预测的不确定性，同时在多任务方式下增强分割一致性。我们通过三个阶段的解决方案来实现这一目标。首先，我们训练一个分割网络来生成预测概率高度不确定的粗糙伪标签。其次，我们利用数据丰富的统计信息，利用多任务模型增强一致性，降低伪掩码的不确定性。我们将我们的方法与现有的半监督语义分割方法进行了比较，并通过大量的实验证明了其SOTA的性能。

1.Introduction

语义分割是计算机视觉中的一项基础任务，它的目标是从一组预定义的类中为图像中的每个像素分配一个标签。这一任务在文献中得到了广泛的探索，但尚未完全解决。目前最先进的模型是建立在深度网络的基础上的[30,5,34,53,54]。虽然这些技术已经报道了前所未有的结果，但它们依赖于非常高的标注数据体系。这是一个强烈的假设，因为标注是像素级的;这是昂贵的，耗时的，并且是人类固有的偏见。要解决缺乏大量且具有良好代表性的标签集的问题，可以更多地依赖于无标签的样本。

另一种方法是使用完全无标记的数据，即无监督学习。然而，由于样本和类之间缺乏对应关系，性能大大降低，因此这种范式在语义分割方面并不成功。另一种选择是使用弱监督技术[25,40]，然而，来自未标记样本的丰富信息没有得到充分利用，性能仍然有限。一个可行的选择是使用半监督学习，利用大量的标记数据和少量的注释。虽然semi-supervised learning(SSL)在社区里已被广泛开发[4],深度semi-supervised learning只是在过去的几年里有明显的进步,而且主要用于图像分类的任务如[24,32,44],这些理论近年来在语义分割领域得到广泛的应用[17,11,36,10]。

现有的SSL技术大致可以分为熵最小化[13]、生成式模型[23]、基于图的技术[56]、基于代理的技术(建立在伪标签/自训练之上)[46]、一致性正则化[24,32]和整合了每个原则中最佳原则的整体方法。对于语义分割，现有的技术使用生成模型和一致性正则化技术。尽管该任务报告的结果很有希望，但仍有很多改进的空间，特别是关于如何改进置信度预测。

为此，我们提出了一种新的语义分割框架，该框架假定数据的有标签样本非常少。我们的整体解决方案是半监督学习的三个阶段的自我训练技术。我们的框架的每个阶段都有一个目的:阶段1通过训练的分割网络生成初始伪掩模(在伪标签意义上)，阶段2是一个多任务模型，学习伪掩码统计信息(Task 2)和分割(Task 1)，用于生成更高质量的伪掩码，然后在第三阶段利用更新后的伪掩码和微量标注集对分割网络进行训练。我们的贡献如下:

• 我们提出了一种新的深度半监督语义分割的自训练框架，其中我们强调:

  • 一个端到端的优化模型被构建成一个三个阶段的解决方案。我们的模型使用多任务模型降低了伪掩模预测概率的不确定性，在利用数据统计信息的同时加强了分割的一致性。
  • 本文介绍了语义切分的一个新视角——整体原则。我们证明，虽然一致性正规化是重要的，但需要解释预测伪掩模的不确定性。我们表明，从两个来源学习信息可以增加确定性预测。

• 我们在一系列数值和可视化结果上验证了我们的技术，并将其与当前的SOTA深度半监督语义分割技术进行了比较。我们证明了我们设计良好的模型实现了它的目的，并且我们所提出的技术优于当前那些依赖非常有限的真值标签的SOTA语义分割模型。

2.Related Work

Long[30]的开创性工作表明，全卷积网络具有将学习到的表示转移到分割任务中的潜力，因此使用深度监督技术进行分割已经取得了重大进展，包括[2,5,34,38]。最近，复杂的机制被结合起来，以克服现有技术的性能限制，包括结构搜索[33,29,52,27]、注意力机制[53,50,51,16]和重新设计几种架构的原则[47,26,37,54,55]。

尽管这些技术显示了惊人的结果，一个主要的共有缺点是这些假设都是建立在用于足够多的标记数据的设置下。这推动了对标注数据依赖较少的技术的快速发展，如弱监督分割，如[1,25,40,19,42,18]。当前关注的焦点是深度半监督学习的最新发展，这也是本研究的重点。在本节中，我们依次回顾现有的技术。

Deep Semi-Supervised Learning. 自[4]领域的早期开发以来，半监督学习(SSL)已经得到了广泛的研究。然而，在过去的几年中，对这种范式的兴趣有了显著的增长，这尤其是因为SSL的基础理论与深网的强大功能相结合，深网报告了令人印象深刻的结果，很容易与完全监督的技术相竞争。这种性能增益主要报道在图像分类的背景下，其中已经发展了几种技术，例如[24,32,44]。然而，图像分类与语义分割之间存在显著差异，后者涉及到更密集、更复杂的预测。

图像分类的深层SSL可以使用不同的原则来实现，其中主要的成功是通过一致性正则化实现的[24,43,39,32,3]。该理论的核心思想是，在诱导扰动下的无标记样本(\(x_u∈D_u\)表示无标记集)δ不应改变性能输出，使\(f(x_u) = f(x_u + δ)\)。该原则通过将决策边界推到低密度区域来强制执行SSL的等价形式，从而放松了SSL的集群假设。虽然在图像分类中应用的原则可以以某种方式外推到语义分割，但任务之间的固有差距阻止语义分割技术达到类似的高性能，如分类中报告的[11]，因此，需要重新考虑用于语义分割的深层SSL的设计。

Deep Semi-Supervised Semantic Segmentation. 标注质量在技术性能中起着至关重要的作用。特别是在语义分割的任务中，标签的开销过大。例如，来自分割基准数据集Cityspaces[6]的一张分辨率为1024x2048的图像，涉及超过1M的像素标签，容易出现注释错误，需要考虑像素模糊的问题。SSL非常适合手头的任务，因为前面的方法依赖于一组很小的标签。用于语义分割的深层SSL最近只在少数工作中进行了探索。

早期的技术依赖于GANs[12]原理。

[41]作者提出通过生成gan型合成图像来扩大训练集，以丰富特征空间，加强无标记样本和有标记样本之间的关系。Hung等人[17]提出一种基于GAN的预测概率图与地面真值分割的区分技术。类似地，Mittal等人[31]提出了一个双分支解决方案，包括:i)为输入样本生成每像素类标签的GAN分支和ii)多标签Mean Teacher[43]分支，以消除假阳性预测。

最近，[11]的作者将CutMix[48]扩展到语义分割的上下文中。在这篇论文中，作者应用了强增强的原理，包括Cutout，从图像分类的发现。[36]的作者强调了基于特征、基于预测和随机扰动之间的一致性。Ke等人的[20]使用了缺陷概率图，并将双学生[21]扩展到像素级任务。在[10]中提出了一种基于伪标签的离线自训练方案，增强了有标记集和无标记集之间的一致性。

3.Proposed Technique

本节包含了我们提出的框架的三个核心部分:(i)如何生成初始伪掩码，(ii)多任务模型，以提高初始伪掩模的质量;(iii)如何在最终模型中传播高质量的伪掩模。整个工作流显示在图1中。

Our Deep SSL Setting. 我们的整体技术建立在深度半监督学习的三个阶段的自训练技术框架。在我们的设置中，我们假设我们有大量的未标记数据和少量的标注数据。形式上，我们用\(\{(x_i， y_i) | i∈L\}\)，其中\(x_i\)为图像，\(y_i\)为\(x_i\)对应的ground truth segmentation mask.未标记样本集表示为\(\{x_i| i∈U\}\)。我们考虑的学习问题，只有一小部分的图像被标记，\(|L|\ll|U|\)接下来，我们将详细讨论我们的技术。

3.1. Consistency Regularisation with Strong Augmentations

半监督学习的一个关键问题是如何对未标记数据\(x_i (i∈U)\)计算具有高确定性的标签。在深度学习的背景下，我们研究了先验知识，以便更好地利用丰富的未标记数据。

为了明确SSL分段设置中的一致性规则化，我们定义一个操作符为\(A:\mathbb{R}^{n\times m \times c}\rightarrow \mathbb{R}^{n\times m \times c}\)获取取图像x的随机扰动版本\(Ax\)。这种扰动的常见例子包括旋转、翻转、平移及其组合。我们还定义运算符为 \(B:\mathbb{R}^{n\times n \times c'}\rightarrow \mathbb{R}^{n\times n \times c'}\)它以与A相同的方式将分割掩码的像素映射到一个新的掩码。一致性损失定义为:

\[\mathcal{L}_{\text {con }}:=\sum_{i \in L \cup U} d_{c}\left(f_{\theta}\left(A x_{i}\right), B f_{\theta^{\prime}}\left(x_{i}\right)\right)
\]

其中\(f_\theta\)和\(f_{\theta'}\)分别为\(\theta\)和\(\theta'\)参数化的深度神经网络模型，\(d_c\)表示\(f_\theta(Ax_i)\)和\(Bf_{\theta'}(x_i)\)之间的距离。在这个工作中，A和B是随机算子，也就是说，它们对于不同的样本是不同的。

一致性正则化(1)的一个关键因素是扰动算子A和b。代价(1)的最小化使网络对数据样本邻居的预测变得平滑。因此，A算子的选择反映了邻居的大小，影响了分割结果。特别地，我们遵循使用强增强的原则，因为我们强制扰动是多样化的和自然的，以提高SSL性能[45]。然后结合RandAugment[7]和Cutout[8]实现随机算子A和B作为强增广器。

一致性损失(1)建立在两个网络模型之间的关系上。在文献中，常将\(f_θ\)网络和\(f_θ'\)网络分别称为学生模型和教师模型。教师模型的参数可以用几种方法确定。在\(Γ\)模型[24]中，θ'设为θ。在均值教师(MT)方案[43]中，θ'被计算为θ在训练过程中的指数移动平均(EMA)。在MT中，由于学生模型的收敛性，当学习过程继续进行[21]时，θ'越来越接近学生的权值θ。因此，随着训练步数的增加，教师模型\(f_{\theta '}\)缺乏有效的指导。

基于一致性正则化方法(例如[24,43])，模型输出的一致性依赖于训练期间教师和学生的耦合，基于这一事实，我们提出了一个通过多任务学习集成额外指导信号的框架。接下来，我们提出了一个三阶段的自训练方法，该方法包含伪掩模统计数据，以帮助学习更好的一致输出的未标记数据，从而提高不同阶段的伪掩模质量。

3.2. A Three-Stage Self-Training Framework

拟议的框架包括三个阶段(见图1)如下所述。

Stage 1: Generating Initial Pseudo-Masks. 在这一阶段，我们利用微小的标记图像集训练分割网络\(f_θ\)，生成具有高不确定性(即低质量伪掩模)的初始伪掩模。为此，我们使用如下定义的loss:

\[\mathcal{L}_{\mathrm{seg}}:=\sum_{i \in L} d_{s}\left(f_{\theta}\left(x_{i}\right), y_{i}\right)
\]

\(d_s\)是交叉熵损失。一旦网络被训练，伪分割掩模从模型\(\hat{y}_i:= arg max f_θ(x_i)\)的输出产生，其中arg max函数在像素级下操作。这一阶段的主要目的是生成初始伪掩模，由于真值标签数量有限，其最大预测概率具有高度不确定性。

Stage 2: Increasing Certainty for Pseudo-Masks. 在这一阶段，将标记的样本、未标记的图像和第一阶段生成的伪掩模\(\hat{y}_i\)一起使用，通过训练第二个模型来生成更高质量的伪掩模。由于掩码\(\hat{y}_i\)是不准确的，我们不使用它们来拟合分割网络，而是学习在多任务模型中复制它们。为此，我们考虑了两个任务，即分割任务(task 1)和辅助任务(task 2)。我们的任务2旨在从\(\hat{y}_i\)中提取伪掩码统计信息。

我们的两个任务在一个多任务网络中构建，如图1所示。第一个网络fθ是由N块(即蓝色块)组成的分割网络。第二个网络(Task 2的\(\hat{g}_{\hat{\theta}}\))由N个区块和与fθ共享的第一个N−2区块组成(见图1左侧)，\(\hat{\theta}\)为\(\hat{g}\)的共享权和自身权。Task 1采用了segmentation loss(2)。另外，为了提高分割精度，我们在这个任务中加入了一致性损失(1)。在训练[43]时，利用θ的指数移动平均(EMA)更新教师网络的权值\(\theta'\).对于Task 2和使用的伪掩码\(\hat{y}_i\)，损失定义为:

\[\mathcal{L}_{\mathrm{pl}}:=\sum_{i \in L \cup U} d_{c}\left(\hat{g}_{\hat{\theta}}\left(A x_{i}\right), B \hat{y}_{i}\right)
\]

我们强调\(\hat{g}_{\hat{\theta}}\)的输出与fθ的输出不相同，因为Task 1没有看到伪掩模。但是，Task 2中的伪掩码在优化过程中为Task 1产生额外的信号。总的来说，我们在这个阶段使用的损失是:

\[\mathcal{L}:=\mathcal{L}_{\mathrm{seg}}+\lambda_{1} \mathcal{L}_{\text {con }}+\lambda_{2} \mathcal{L}_{\mathrm{pl}}
\]

其中λ1和λ2是平衡任务的超参数。对多任务网络进行训练后，通过\(\tilde{y}_{i}:=\operatorname{argmax} f_{\theta}\left(x_{i}\right)\)计算出高质量的伪掩模。

Stage 3: Propagating Higher Quality Masks Information.最后，在这个阶段，使用阶段2的高质量伪掩模\(\tilde{y}_i\)来监督最终的模型。高品质pseudo-masks集成以第二阶段中类似的方式执行,除了我们用\(\tilde{g}_{\tilde{\theta}}\)更换网络\(\hat{g}_{\hat{\theta}}\),其中\(\tilde{g}_{\tilde{\theta}}\)有N个模块和前N-1个模块(见图1的右边)和fθ共享(网络\(\hat{g}_{\hat{\theta}}\),然而,只有N-2个模块和fθ相同)。这使得伪掩模信息更好地传播到fθ，这是因为\(\tilde{y}_i\)比阶段2中使用的\(\hat{y}_i\)更准确。这一阶段所用的loss为:

\[\mathcal{L}^{\prime}:=\mathcal{L}_{\mathrm{seg}}+\lambda_{1} \mathcal{L}_{\mathrm{con}}+\lambda_{2} \mathcal{L}_{\mathrm{pl}}^{\prime}
\]

其中：\(\mathcal{L}_{\mathrm{pl}}^{\prime}:=\sum_{i \in L \cup U} d_{c}\left(\tilde{g}_{\tilde{\theta}}\left(A x_{i}\right), B \tilde{y}_{i}\right)\)

算法1列出了我们的技术结合之前的各个阶段来解决语义分割问题的整个过程。