论文:https://arxiv.org/abs/2005.10821
代码:https://github.com/nvidia/semanic-segmentation
论文核心:how to combine multi-scale predictions :层次注意力机制。
多尺度的推理是改善语义分割结果的常用方法。将多个图像尺度通过网络传递,然后将结果进行平均或最大池化而结合起来。在这项工作中,我们提出了一种基于注意力的方法来结合多尺度预测。我们表明,在特定规模下的预测能更好地解决特定的故障模式,而网络学会了在这种情况下支持那些规模,以便产生更好的预测。我们的注意力机制是分层的,这使得它的记忆效率比最近的其他方法提高了4倍。除了加快训练,这允许我们训练更大的作物尺寸,这导致更大的模型精度。我们在Cityscapes 和 Mapillary Vistas.两个数据集上展示了我们的方法的结果。对于城市景观,它有大量的弱标记图像,我们也利用自动标记来提高泛化。使用我们的方法,我们在Mapillary(61.1 IOU val)和城市景观(85.1 IOU测试)中都取得了新的最先进的结果。
Keywords: Semantic Segmentation · Attention · Auto-labelling
在语义分割任务中有一个权衡,即某些类型的预测在较低的推理分辨率下得到最好的处理,而其他任务在较高的推理分辨率下得到更好的处理。精细的细节,比如目标的边缘或者薄结构,通常用放大的图像尺寸比较好预测。而对大型结构的预测,需要更多的全局上下文信息,通常通过减小图像尺寸预测,结果较好。因为网络的感受与可以观察到更多有用的上下文信息。后者的情况我们成为类混淆,示意图如下:
使用多尺度推理是解决这种权衡的常见做法。预测是在一系列的尺度上完成的,并且结果使用平均池化或最大池化相结合。平均池化的结合方式通常可以改善结果,但它存在将最好的预测与较差的预测结合起来的问题。例如,如果对于给定的像素,最好的预测来自于2x的尺度,而更糟糕的预测来自于0.5x的尺度,那么平均这些预测就会结合起来,导致低于平均水平的输出。另一方面,最大池化只选择N个尺度中的一个用于给定像素,而最佳答案(最佳的组合结果)可能是跨不同尺度的预测的加权组合。
为了解决这个问题(如何结合不同尺度的预测以达到更好的分割效果),我们采用了注意力机制在像素级上去怎样去结合多尺度的预测结果,与Chen提出的相似。不同的是:我们提出了一个层次注意机制,通过网络学习预测相邻尺度之间的相对权重。
与其他方法相比,优点:计算量少;在推理阶段尺度选择灵活,不受限于训练时的尺度。
在Cityscape数据集上还是用的自动标注的方法,进一步增加分割精度,改善泛化能力。
(为了在城市景观中获得最先进的结果,我们还采用了粗糙图像的自动标记策略,以增加数据集中的方差,从而提高泛化。我们的策略受到了近期多项研究的启发,包括[2,3,4]。与典型的软标签策略相反,我们采用硬标签来管理标签存储大小,这有助于通过降低磁盘IO成本来提高训练吞吐量。)
本文工作贡献总结如下:
•一种有效的分层多尺度注意力机制,有助于消除类混淆,并让网络学习如何最好地结合来自多个推理尺度的预测
•基于硬阈值的自动标记策略,利用未标记的图像并提高IOU
•我们在城市景观Cityscapes(85.1 IOU)和Mapillary Vistas (61.1 IOU)方面取得了最先进的成果。
Multi-scale context methods(多尺度上下文方法):
PSPNet-spatial pyramid pooling module; deeplab- Atrous Spatial Pyramid Pooling (ASPP);
利用中间层特征而不是网络最后一层的特征,来创建多尺度上下文。
Relational context methods(关系上下文的方法):
在实践中,金字塔池化技术关注的是固定的方形上下文区域,因为池化和扩展通常以对称的方式使用。此外,这些技术往往是静态的,无法学习。然而,关系上下文方法是通过关注像素之间的关系来建立上下文,而不是局限于正方形区域。关系上下文方法的学习性质允许基于图像组合建立上下文。这样的技术可以为非方形语义区域建立更合适的上下文,例如一列长火车或一根细长的灯柱。
Multi-scale inference(多尺度推理)
两种最常用的结合多尺度预测结果的方法:平均池化和最大池化。
但是平均池化将不同尺度的预测结果赋予相同的权重,但结果可能不是最优的。(sub-optimal 次优的)。我们使用注意力机制来组合不同尺度的结果。
Chen同时在所有尺度计算注意力,将注意力加权后的结果再结合起来。
Yang使用网络不同层的特征再结合起来。
以上方法的不足之处:
然而,上述两种方法都有一个共同的特点,那就是网络和注意力头都是用一套固定的尺度来训练的。只有这些规模可以在运行时使用,否则网络必须重新训练。
我们提出了一种基于分层的注意机制,该机制不知道推理时间中尺度的数量。此外,我们还表明,我们提出的分层注意机制不仅提高了平均池的性能,而且允许我们诊断并可视化不同尺度对类和场景的重要性。
此外,我们的方法与其他注意力或金字塔池方法(如[22,25,26,9,27,10,28])是正交的,因为这些方法使用单尺度图像并执行注意力以更好地结合多层次特征来生成高分辨率预测。
Auto-labelling(自动标注):
“Cityscapes是关于城市街道场景的语义理解图片数据集。它主要包含来自50个不同城市的街道场景,拥有5000张在城市环境中驾驶场景的高质量像素级注释图像(其中 2975 for train,500 for val,1525 for test,共有19个类别);此外,它还有20000张粗糙标注的图像(gt coarse)。” https://blog.csdn.net/avideointerfaces/article/details/104139298
为了在Cityscapes数据集上达到SOTA(state of art),我们还采用了自动标注的策略。(受其他语义分割的半监督学习训练,伪标签的方法等启发。)我们为城市景观中的粗糙图像生成密集的标签。我们生成的标签很少有未标记的区域,因此我们能够利用粗图像的全部内容。
虽然大多数图像分类自动标记工作使用连续标签或软标签,但为了提高存储效率和训练速度,我们生成硬阈值标签。对于软标签,教师网络为图像的每个像素提供了N个类的连续概率,而对于硬标签,则使用阈值来为每个像素选择一个最高类。类似于[37,4],我们为粗糙的城市景观图像生成硬密集的标签。示例如图4所示。(自动标注的标签生成了更多细节。)
我们不执行对标签的迭代优化。相反,我们使用默认的粗糙和精细的标记提供的图像对我们的教师模型进行一次完整的训练。在这次联合训练之后,我们对粗糙的图像进行自动标记,然后将其替换为我们的教师训练数据(training recipe),从而获得最先进的测试结果。使用我们的伪生成硬标签和我们提出的分层注意力,我们能够获得最先进的城市景观结果。
受chen的工作启发(左边的图),每个尺度学习一个密集的掩模,这些多尺度的预测通过在掩模之间进行像素上的乘法(也就是乘以注意力),然后在不同尺度之间进行像素上的求和得到最终的结果。
我们不是学习每个固定尺度的所有注意力掩膜,而是学习相邻尺度之间的相对注意力掩膜。
(we learn a relative attention mask between adjacent scales.)(右边的图)
在训练网络时,我们只使用相邻的尺度对进行训练。如图2所示,给定一组来自单一(较低)尺度的图像特征,我们预测了两个图像尺度之间的密集关注度。在实践中,为了得到这对缩放后的图像,我们取一个输入图像并将其缩放2倍,这样我们就得到了一个1x的缩放后的输入和一个0.5x的缩放后的输入,尽管可以选择任意的缩放比例。需要注意的是,网络输入本身是原始训练图像的一个缩放版本,因为我们在训练时使用了图像缩放增强。这使得网络学会了预测一系列图像尺度下的相对注意力。
在进行推理时,我们可以分层地应用所学的注意,将N个尺度的预测组合在一起,形成一个计算链,如图所示,并由下面的等式描述。我们优先考虑低比例尺,并逐步向上扩展到更高的比例尺,我们的理念是,它们具有更多的全局背景,可以选择哪些地方的预测需要由更高尺度的预测来确定。(确定细节的优化)
我们定义r为缩放比例,r=0.5,下采样2倍。r=1为原图,r=2,上采样2倍。在训练过程中我们选择r=0.5, r=1这两个尺度。
两张图片(r=0.5, r=1)送入共享的网络,生成各自尺度的语义分数(L)和注意力掩膜(α)(semantic logits L and also an attention mask(α))
U为上采样操作,*和+分别是像素上的乘法和加法。
我们提出的这种策略有以下两个优点:
在推理时,我们现在可以灵活地选择尺度,通过我们提出的注意机制以分层的方式将0.5x和1.0x训练好的模型添加到0.25x或2.0x这样的新尺度(应该是指0.25和0.5、1和2,可以分别使用已训练好的层次注意力结构)。这不同于以前提出的仅限于使用模型训练中使用的相同比例的方法。
•与显式方法相比,这种分层结构可以提高我们的训练效率。使用显式方法,如果使用scale 0.5, 1.0, 2.0,相对于单尺度训练,培训成本为0.52 + 1.02 + 2.02 = 5.25。采用分层方法,训练成本仅为0.52 + 1.02 = 1.25。
本文的实验是基于HRNet-OCR+我们提出的层次注意力;为了验证我们提出的层次注意力的有效性,我们用deeplabv3+ + resnet50(backbone)进行了验证。
为了验证我们的多尺度注意力方法,我们训练的网络的基础主干结构为deeplabv3+ +resnet50作为骨干网。下表表明我们的层次注意力方法达到最好的IoU(51.6),在推理阶段引入新的尺度0.25x时,有明显提升。
与chen提出的显示方法不同,在增加0.25这个尺度时,我们的方法不需要重新训练网络。在推理阶段的灵活性是我们的方法的一个关键。We can train once but evaluate flexibly with a range of different scales.
从表格中看出,在使用平均池化这种方式组合不同尺度预测结果时,增加0.25这个尺度使IoU下降了0.7。但是在我们的基于相邻尺度注意力的方法中,增加0.25这个尺度,使IoU上升0.6。因为0.25的图片是很粗糙的,仅用平局池化的方法组合,粗糙程度会给边缘和精细的信息带来反面的影响。也说明了我们的组合方法,以一种合理的方式利用了0.25尺度图片的信息,没有在边缘这种地方利用上它。如图3 所示,在图3中可以观察到这方面的例子,在左边的图像中,只有很少的特征被0.5x的预测所关注,但是在2.0x尺度中出现了非常强烈的注意信号。相反地,对于右边非常大的区域,注意机制学习更多地利用较低的尺度(0.5x),很少利用错误的2.0x预测。
我们只采用了两个相邻尺度中较低的那个尺度特征送入attention head,我们也尝试了将两个尺度特征都送入attention head,在准确率方面没有明显变化,于是我们采用single-scale features.
在Cityscapes数据集里,有3500张精细标签和20000张粗标签。粗糙图像的标签质量非常低,包含大量未标记像素。如图4所示:
通过使用自动标注策略,我们改善了标签的质量,这反过来又帮助了提升模型的IoU。
另外我们使用了硬标签标注。常用于图像分类自动标注的软标签方法,需要教师网络为每一个像素提供一个概率分布,包含N个标签中每个标签的概率。非常占用内存:
20000 images * 2048w * 1024 h * 19 classes * 4B = 3.2TB.而使用硬标签,我们只为每个像素选出最高的概率的那一类。教师网络的预测超过阈值成为真实的标签,否则像素被标记为忽略类。在实践中,我们使用0.9的阈值。
训练细节:pytorch,8 GPUs per DGX node,分布数据并行训练,同步批标准化;SGD;
训练:momentum 0.9 and weight decay 5e−4;polynomial的学习率策略;在默认设置下,我们使用RMI (基于互信息的损失)作为主要损失函数,而对辅助损失函数使用交叉熵。
数据增强:高斯模糊,色彩增强,随机水平垂直翻转,随机缩放(0.5-2.0);将cityscapes裁剪为2048x1024,,将Mapillary裁剪为1856x1024。
5.1.1 在Cityscapes上的结果
两个数据集的训练都采用HRNet-OCR as the trunk along with our proposed multi-scale attention method。
在这项工作中,我们提出了一个分层的多尺度注意的语义分割方法。我们的方法在提高分割精度的同时也提高了内存和计算效率,这两个都是实际问题。训练效率限制了研究的速度,而GPU内存效率限制了作物网络的训练规模,这也限制了网络精度。我们的经验表明,使用我们提出的方法,城市景观和地图的都有了改善。