https://arxiv.org/abs/2005.10821
https://github.com/NVIDIA/semantic-segmentation
1.5k星
21 May 2020
最近开源了：https://paperswithcode.com/paper/hierarchical-multi-scale-attention-for#code
摘要：多尺度推理常用于改善语义分割结果。将图像缩放到多个不同尺度，再输入给同一个网络，再将结果取平均或最大池化。本文提出一个基于注意力的方法来组合多尺度预测。我们表明，某些尺度下的预测更能处理特定的故障情况，而网络学会了在这种特定情况下更倾向于这些尺度，以便输出更好的结果。我们的注意力机制是分层的，这使得它的内存效率比其他最近的方法快4倍。这样除了能够训练地更快以外，还能让我们使用更大的裁剪尺寸训练，从而使得模型准确度更高。我们在Cityscapes和Mapillary Vistas两个数据集上验证我们的方法。对于具有大量弱标注图像的Cityscapes数据集，我们还利用自动标注改进泛化能力。我们的方法在Cityscapes数据集（85.1 IOU test）和Mapillary Vistas数据集（61.1 IOU val）上都取得最先进水平。
关键词：语义分割、注意力、自动标注
1.引言
语义分割的任务是给图像中所有像素打上多个类别中的一个的标签。该任务有个折中，就是某些情况下的预测在低分辨率下推理得更好，而另一些情况下的预测在高分辨率下推理得更好。在放大的图像尺寸下，精细的细节，例如物体边缘或者细长的结构，会推理得更好。同时，缩小的图像尺寸下，对需要更多的全局上下文信息的大物体会推理得更好，因为网络的感受野可以观察到更多的必要的上下文信息。我们将后一种问题称为类别混淆(class confusion)。两种情况的例子如图1所示。
多尺度推理是这一问题的常见解决办法。图像缩放到多个不同尺度，再输入给同一个网络，再将结果取平均或最大池化。取平均来组合多个尺度的预测通常可提高结果，但存在将最佳预测和最差预测结合的问题。煮个栗子，对于某像素，最佳预测来自2×尺度，其0.5×尺度的预测要差很多，取平均的话会将这两者结合起来，得到一个次优的输出。另一方面，最大池化是选择该像素在多个尺度预测中的一个，而最优结果可能是多个尺度预测的加权组合。
为了解决这个问题，我们提出一种使用注意力机制来将某像素的多个尺度预测结合在一起的办法，类似于Chen[1]。我们提出一种层次化的注意力机制，学习预测相邻尺度之间的关系权重，由于它是层次化的，所以我们只需要增加一个额外的尺度的训练pipeline，而其它的方法，例如[1]，则需要在训练阶段显示地增加每一个额外的推理尺度。例如，当多尺度推理为{0.5，1.0，2.0}时，其它的注意力方法要求网络在所有这些尺度上训练，导致额外的倍的训练成本。我们的方法只需要再增加0.5×尺度的训练，导致增加倍的训练成本。此外，我们提出的层次化机制还可使推理时具有选择额外的尺度的灵活性，而以往的方法在推理的时候只能使用训练时候的尺度。为了在Cityscapes上取得最先进水平，我们还使用了一种对粗标注图像的自动标注策略，以增加数据集中的方差，以提高泛化能力。我们的自标注方法是受一些最近的工作的启发，包括[2,3,4]。和典型的软标签不同，我们使用硬标注，这样标签存储内存就更小，可以降低磁盘IO成本，提高训练速度。
1.1 贡献
一种高效地层次化注意力机制，使网络可以学习如何最好地组合多个尺度的推理，能解决类别混淆和细节的问题。
一种基于硬-阈值的自动标注策略，利用未标记图像来提升IOU。
在Cityscapes数据集（85.1 IOU test）和Mapillary Vistas数据集（61.1 IOU val）上取得最先进水平。
2 相关工作
多尺度上下文 最先进的语义分割网络使用低输出步长的主干网络，这可以更好地分割细节，但是感受野也较小。小的感受野会导致网络难以预测大物体。金字塔池化可通过组合多个尺度的上下文来解决感受野较小的问题。PSPNet[5]使用一个空间金字塔池化模块，使用主干网络的最后一层的特征，用一些列卷积和池化来组合多个尺度的特征。DeepLab[6]使用Atrous Spatial Pyramid Pooling(ASPP)，用不同扩张率的扩张卷积，以建立比PSPNet更加稠密的特征。最近ZigZagNet[7]和ACNet[8]使用中间层特征，而不仅仅是最后层特征，来创建多尺度上下文。
关系上下文 实际上金字塔池化关注的是固定的方形上下文区域，因为池化和扩张通常都是对称的。此外这些方法是固定的，不是可学习的。关系上下文通过关注像素之间的关系来构建上下文，而不是固定的方向区域。关系上下文的可学习特性使得上下文的构建可基于图像的各个分区。该方法能够对非正方形的物体区域构建更合适的上下文，例如长火车和细高的灯柱。OCRNet[9]、DANET[10]、CFNet[11]、OCNet[12]和其他相关工作[13,14,15,16,17,18,19,20]利用这些关系来构建更好的上下文。
多尺度推理 一些关系上下文和多尺度上下文方法都使用了多尺度推理来达到最好的结果，例如[21,22,23,9]。多尺度推理有两种常见方法：取平均和最大池化，取平均要更常见些。然而，取平均对不同尺度的输出是等权重的，这样的结果可能是次优的。为了解决这个问题，[1,24]使用注意力组合多个尺度。[1]使用网络的最后层特征来训练一个建模各个尺度间关系的注意力头。[24]组合不同层的特征来构建更好的上下文信息。然而，这两种方法有共同的确定，网络和注意力头是用一组固定尺度训练的，推理时只能使用这些尺度，否则网络要重新训练。我们提出一种层次化的注意力机制，它对推理时用多少个尺度是无关的。我们方法不仅比取平均的性能更好，而且，更进一步地，我们的方法可用于可视化、诊断不同尺度分别对不同物体类别和场景的重要性。此外，我们的方法和其它的注意力方法和金字塔池化方法是正交的，例如[22,25,26,9,27,10,28]，因为这些方法使用单个尺度，并使用注意力来更好地组合不同层级的特征，以生成高分辨率的预测。
自动标注 很多在Cityscapes上的语义分割工作都使用了大约2w张的粗标注图像以达到最先进水平[12,29]。然而，由于标签是粗标注的，图像中相当一部分是未标记的。为了在Cityscapes上取得最先进水平，我们提出一种自动标注策略，我们的灵感是来自于NoisyStudent[2]，以及其他语义分割半监督自训练方法[30,31,32,33,34]，以及其他基于伪标注方法[4,35,36,3]。我们为Cityscapes的粗标注图像生成稠密标注。我们生成的标签只有很少的未标注区域，这样我们就能够充分利用粗标注数据的全部图像内容。
虽然大多数图像分类的自标注方法使用软标签，但为了提高存储效率和训练速度，我们用的是被阈值过的硬标签。软标签在语义分割中是指，教师网络为图像的每个像素上都提供分别属于每个类别的概率。而硬标签是用阈值从各个类别的概率中选出最大的。类似于[37,4]，我们为Cityscapes粗标注图像生成稠密的硬标签。图4是一些例子。不同于NoisyStudent[2]，我们对生成的伪标签不做迭代的调精。相反，我们使用默认的粗标注数据和精细标注数据，对教师模型进行一次完整的训练。在联合训练之后，对粗标注图像进行自动标注，在将这些自动标注数据替换我们教师模型用的粗标注数据，以获得最好的测试结果。使用我们的伪硬标注和层次化注意力，我们能获得在Cityscapes上的最先进水平。
3.层次多尺度注意力
我们的注意机制在概念上与[1]（Attention to scale）非常相似，在[1]中，在每个尺度下都学习，得到dense mask，并通过在对mask执行像素级乘法，然后在不同尺度之间进行像素级求和，来组合这些多尺度预测，以获得最终结果，见图2。我们称陈的方法为显式方法。使用我们的分层方法，我们学习相邻尺度之间的相对注意mask，而不是学习每个尺度的所有注意mask。在训练网络时，我们只使用相邻的尺度对(adjacent scale pairs)进行训练。如图2所示，给定一组来自单个（较低）尺度的图像特征，我们预测两个图像尺度之间的像素级的密集的相对注意力权重(dense pixel-wise the relative attention)。实际上，为了获得一对缩放图像，我们获取一个输入图像并将其缩小2倍，这样我们就剩下1倍缩放输入和0.5倍缩放输入，尽管可以选择任何缩小比率。需要注意的是，网络输入本身是原始训练图像的重新缩放版本(re-scaled version)，因为我们在训练时使用图像缩放增强。这允许网络学习预测一个范围内(a range of)的图像尺度的相对注意力。在进行推理时，我们可以分层地应用所学的注意力，将N个预测尺度组合在一起，形成一个计算链，如图所示，并由下面的等式描述。我们优先考虑较低的尺度，并努力向更高的尺度发展，我们的想法是，它们有更多的全局上下文，可以选择哪些区域是需要通过更高尺度的预测来refine的。
更正式地，在训练期间，给定的输入图像缩放因子，r=0.5表示下采样两倍，r=2.0表示上采样2倍，r=1表示不进行操作。对于我们的训练，我们选择r=0.5和r=1.0。然后，r=1和r=0.5的两个图像通过一个共同的网络块(the shared network trunk)，该网络块输出各个尺度下的semantic logits和注意力掩码，该掩码用来组合不同尺度下的logits。因此，如果训练和推理都使用2个尺度的话，令表示双线性上采样，∗ 和+分别为逐像素乘法和加法，公式可形式化为：
使用我们提出的方法有两个优势：
• 在推理时，我们现在可以灵活地选择不同的尺度，因此，通过我们提出的分层式注意力机制链，对于在0.5x和1.0x尺度下训练的模型，我们在推理时可以添加新的尺度例如如0.25x或2.0x。这不同于先前提出的方法，先前的方法训练时使用哪些尺度，在推理时就得用哪些尺度。
• 与显式方法相比，这种层次结构使我们能够提高训练效率。使用显式方法，如果使用0.5、1.0、2.0尺度，相对于单量表训练，训练成本为0.5^2+1.0^2+2.0^2=5.25。采用我们的分层方法，训练成本仅为0.5^2+1.02=1.25。
3.1 架构
主干：对于本节中的消融研究，我们使用ResNet-50[38]（配置为输出步幅8）作为我们网络的trunk。为了获得最先进的结果，我们使用了更大、功能更强大的主干HRNet OCR[9]。
语义头：语义预测由一个专用的全卷积头执行，该卷积头包括（3x3 conv）→ （BN）→ （ReLU）→ （3x3卷积）→ （BN）→ （ReLU）→ （1x1 conv）。最后的卷积输出num_classes个通道。
注意头:
注意力预测是使用一个单独的头部来完成的，该头部在结构上与语义头部相同，除了最终的卷积输出，它输出一个通道。当使用ResNet-50作为主干时，语义头和注意头会收到ResNet-50最后阶段的特征。当使用HRNet OCR时，语义和注意头会收到OCR块中的特征。使用HRNet OCR的时候，还存在一个辅助语义头，在OCR之前直接从HRNet主干获取其特征。该头部包括（1x1 conv）→ （BN）→ （ReLU）→ （1x1 conv）。在将注意力应用于语义logits之后，使用双线性上采样将预测上采样到目标图像大小。
3.2 分析
为了评估我们的多尺度注意力方法的有效性，我们使用DeepLab V3+架构和ResNet50主干训练网络。在表1中，我们表明，与基线平均法（49.4）或显式法（51.4）相比，我们的分层注意法具有更好的准确性（51.6）。我们还观察到，当添加0.25倍尺度时，我们的方法的结果明显更好。与显式方法不同，当使用额外的0.25倍尺度时，我们的方法不需要重新训练网络。这种推理时的灵活性是我们方法的一个主要优点。我们可以训练一次，但可以灵活地用不同的尺度进行评估。
此外，我们还观察到，对于基线平均多尺度方法，简单地添加0.25倍尺度对精度有害，因为它会导致IOU减少0.7，而对于我们的方法，添加额外的0.25倍尺度会将精度再提高0.6个IOU。使用基线平均法，0.25倍的预测是如此粗糙，以至于当与到其他尺度平均时，我们观察到诸如车道标记、人洞、电话亭、路灯、交通灯和交通标志（前后）、自行车架等类别下降1.5 IOU。预测的粗糙度损害了边缘和细节。然而，在我们提出的注意方法中，添加0.25倍的比例将我们的结果提高了0.6，因为我们的网络能够以最合适的方式应用0.25倍的预测，避免在边缘使用它。这方面的例子可以在图3中观察到，其中对于左侧图像中的fine posts，0.5x预测很少关注posts，但2.0x尺度中存在非常强的注意信号。相反，对于右侧非常大的区域，注意机制学习到最大程度地利用较低的尺度（0.5倍），很少利用错误的2.0倍预测。
3.2.1单尺度与双尺度特征
虽然我们确定的架构只从两个相邻图像尺度中较低的一个尺度的特征输入给注意力头（见图2），但我们也尝试用两个相邻尺度的特征训练注意力头。我们没有观察到准确度上的显著差异，所以我们选择了单个尺度的特征。
4.Cityscapes上的自动标注
受最近的在图像分类上任务上的自动标注工作[2,39]的启发，我们在Cityscapes数据集上使用自动标注来提升有效数据集的大小和标注的质量。Cityscapes有2w张粗标注图像和3500张精细标注图像。粗标注图像的标注非常一般，有大量像素是无标签的，如图4所示。通过我们的自动标注方法，我们可以提高标签质量，这反过来又有助于提升模型的IOU。
图像分类中常用的是软标签，或者称为连续标签。在语义分割中，软标签是为每一像素标上所属各个类别的概率，这样很占磁盘空间，训练速度会很慢，标签会占大约3.2TB的磁盘：20000×2048×1024×19×4B=3.2TB。训练期间读取这些标签会大大降低训练速度。
因此，我们用硬标签，即对每个像素，选择教师网络输出的最大类别概率。我们对教师网络输出的概率标签进行阈值化。超过阈值的概率视为标签，否则该像素被视为忽略类。实践中，我们使用0.9的阈值。
5 结果
5.1 实施协议
在本节中，我们将详细描述我们的实现协议。
训练详情
我们的模型在Nvidia DGX服务器上使用Pyrotch[40]进行训练，每个节点包含8个GPU，具有混合精度、分布式数据并行训练和同步批处理规范化。我们的优化器使用随机梯度下降（SGD），每GPU的批量大小为1，动量为0.9，重量衰减为5e−4.在训练中。我们采用“多项式”学习率策略[41]。在默认设置下，我们使用RMI[42]作为主要损失函数，并使用交叉熵作为辅助损失函数。对于城市景观，我们使用2.0的多指数，初始学习率为0.01，并在2个DGX节点上训练175个历元。对于Mapillary，我们使用1.0的多边形指数、0.02的初始学习率，并在4个DGX节点上训练200个历元。如[29]中所述，我们在数据加载器中使用类均匀采样，从每个类中平均采样，这有助于在数据分布不均匀时改进结果。
数据增强：
在训练过程中，我们对输入图像采用高斯模糊、颜色增强、随机水平翻转和随机缩放（0.5x-2.0x）来增强数据集。城市景观的裁剪尺寸为2048x1024，Mapillary Vistas为1856x1024。
5.1.1城市景观结果
Cityscapes[43]是一个大型数据集，在5000幅高分辨率图像上标记了19个语义类。对于城市景观，我们使用HRNet OCR作为主干，并提出了多尺度注意方法。我们使用RMI作为主分割头的损失，但对于辅助分割头，我们使用交叉熵，因为我们发现使用RMI损失会导致深入训练的训练精度降低。我们的最佳效果是首先在较大的Mapillary Vistas数据集上进行预训练，然后在城市景观上进行训练。对于Mapillary预训练任务，我们不会使用注意力机制进行训练。除了自动标记的粗略图像外，我们还使用train+val图像实现了最先进的城市景观配方。我们从train+val集合中用50%的概率采样，另外我们从自动标记的图像池中采样。在推断时，我们使用scales={0.5,1.0,2.0}和图像翻转。
我们对城市景观验证集进行了消融研究，如表2所示。多尺度注意力产生的IOU比基线HRNet OCR架构平均池高0.5%。自动标签比基线提高了1.1%的IOU。将这两种技术结合在一起可获得1.4%IOU的总增益。
最后，在表3中，我们展示了我们的方法与Cityscapes测试集中其他性能最好的方法相比的结果。我们的方法得分为85.1，这是所有方法中报告的最佳城市景观测试分数，比之前的最佳分数高0.6 IOU。此外，我们的方法在除三个类别外的所有类别中都拥有最高的每班分数。图5显示了一些结果。
5.1.2 Mapillary Vistas远景的结果
Mapillary Vistas[45]是一个大型数据集，包含25000个高分辨率图像，注释为66个对象类别。对于Mapillary，我们使用HRNet OCR作为主干，并提出了多尺度注意方法。因为Mapillary Vistas图像可以具有非常高且不同的分辨率，所以我们调整图像的大小，使长边为2177，如[23]中所述。我们使用在ImageNet分类上训练的HRNet的权重初始化模型的HRNet部分。由于Mapillary中66个类需要更大的内存，我们将裁剪大小减少到1856 x 1024。在表4中，我们展示了我们的方法在Mapillary验证集上的结果。我们单个模型达到61.1，比下一个最接近的方法Panopoptic Deeplab[23]高2.4，后者使用模型集成实现58.7。
6 结论
在这项工作中，我们提出了一种分层的多尺度注意语义分割方法。我们的方法在提高分割精度的同时，也提高了内存和计算效率，这两者都是实际问题。训练效率限制了研究的速度，而GPU内存效率限制了裁剪网络的训练尺度，这也限制了网络的准确性。我们的经验表明，使用我们提出的方法可以持续改善城市景观和Mapillary Vistas。
致谢：我们要感谢Sanja Fidler、Kevin Shih、Tommi Koivisto和Timo Roman的有益讨论。