Masked-attention Mask Transformer for Universal Image Segmentation
图像分割是关于将不同语义的像素分组,例如,类别或实例成员关系,其中每个语义的选择定义了一个任务。虽然只是每个任务的语义不同,但目前的研究重点是为每个任务设计专门的体系结构。我们提出了一种新的架构Mask -attention Mask Transformer (Mask2Former),能够解决任何图像分割任务(全景、实例或语义)。它的关键组成部分包括Mask -attention,通过约束掩模区域内的交叉注意来提取局部特征。除了将研究工作量减少至少三倍之外,它在四个流行的数据集上的性能显著优于最好的专门架构。最值得注意的是,Mask2Former设置了一个新的最先进的全景分割(COCO上57.8 PQ),实例分割(COCO上50.1 AP)和语义分割(ADE20K上57.7 mIoU)。
本文提出通用的图像分割架构,它在不同的分割任务中优于专门的架构,同时仍然很容易在每个任务上训练。
在maskformer的基础上:
增加了masked attention机制
使用多尺度高分辨率特征来帮助模型分割小目标/区域
调整了decoder部分的self-attention和cross-attention的顺序
通过在少量随机采样点上计算mask loss,来加快训练速度
Mask2Former采用与MaskFormer相同的元架构,具有backbone、像素解码器和transformer解码器。
1、backbone从图像中提取低分辨率特征的主干。
2、像素解码器,从主干的输出中逐渐上采样低分辨率的特征,以生成高分辨率的逐像素嵌入。
3、Transformer解码器,它通过图像特征处理对象查询。
背景特征已经被证明对图像分割非常重要。然而,最近的研究表明,基于transformer的模型收敛缓慢是由于交叉注意层的全局上下文,交叉注意学习关注局部对象区域需要许多训练时间。我们假设局部特征足以更新查询特征,通过自我注意可以收集上下文信息。为此,我们提出了掩码注意,这是交叉注意的一种变体,对每个查询只关注预测掩码的前景区域。
标准交叉注意(带有残差路径)计算
masked attention模型的注意力矩阵计算:
特征位置(x, y)的注意掩码为 :
是前面第(l−1)个transformer解码器层调整掩码预测的二值化输出(阈值为0.5)。它被调整到相同的分辨率kl,M0为从X0得到的二进制掩码预测,即在向Transformer解码器输入查询特征之前的预测。
高分辨率特性提高了模型的性能,特别是对于小目标。然而,这需要计算。因此,我们提出了一种有效的多尺度策略,在控制计算量增加的同时引入高分辨率特征。我们不总是使用高分辨率特征图,而是使用一个特征金字塔,它包括低分辨率特征和高分辨率特征,并一次向一个Transformer decoder层提供一个多尺度特征的分辨率。
具体来说,我们使用pixel decoder生成的特征金字塔,分辨率分别为原始图像的1/32、1/16和1/8。对于每个分辨率,我们添加一个正弦位置嵌入,紧跟DETR,以及一个可学习尺度级嵌入,紧跟Deformable DETR。我们使用这些,从最低分辨率到最高分辨率的对应Transformer解码器层。我们重复这个3层Transformer解码器L次。因此,我们最终的Transformer解码器有3L层。更具体地说,前三层得到的特征图分辨率为H1 = H/32, H2 = H/16, H3 = H/8, W1 = W/32, W2 = W/16, W3 = W/8,其中H和W为原始图像分辨率。此模式以轮询方式对以下所有层重复。
标准的transformer解码器(左图) vs. 本文的transformer解码器(右图)
标准的Transformer解码器层由三个模块组成,依次处理查询特性:自注意模块、交叉注意模块和前馈网络(FFN)。此外,查询特性(X0)在被输入到Transformer解码器之前是零初始化的,并且与可学习的位置嵌入相关联。此外,dropout应用于残差连接和注意图。
为了优化Transformer解码器的设计,本文做了以下三点改进。
首先,我们交换一下self-attention和cross-attention(我们的新“masked attention”)顺序,使计算更有效:第一个自我注意层的查询特征还不依赖于图像特征,因此应用自我注意不会产生任何有意义的特征。
其次,我们使查询特征(X0)也可学习(我们仍然保持可学习的查询位置嵌入),可学习的查询特征在被用于Transformer解码器预测掩码(M0)之前是直接监督的。我们发现这些可学习的查询功能类似于区域提议网络RPN,并且能够生成mask proposal。
最后,我们发现dropout是不必要的,通常会降低性能。因此,我们在我们的解码器完全消除了dropout。
受 PointRend 和 Implicit PointRend 的启发,通过在 K个随机采样点而不是整个掩码上计算掩码损失来训练分割模型。在本文中设置K=12544,即112×112点。 本文在匹配损失和最终损失计算中使用采样点计算掩码损失:
1、在构建二分匹配成本矩阵的匹配损失中,对所有预测和真实掩码的相同的 K 点集进行均匀采样。
2、在预测之间的最终损失及其匹配的ground truth中,使用importance sampling为不同的预测和ground truth对不同的K点进行采样。
这种训练策略有效地将训练内存减少了 3 倍,从每张图像 18GB 到 6GB。