DFormer:用于通用图像分割的扩散引导transformer
摘要:本文介绍了一种通用的图像分割方法DFormer。所提出的DFormer将通用图像分割任务视为一个使用扩散模型的去噪过程。DFormer首先将不同级别的高斯噪声添加到地面真实掩码中,然后学习一个模型来预测从损坏的掩码中去除噪声的掩码。具体地说,我们将深度像素级特征和噪声掩码作为输入来生成掩码特征和注意掩码,并使用基于扩散的解码器逐步执行掩码预测。在推理中,我们的DFormer直接从一组随机生成的掩码中预测掩码和相应的类别。大量的实验表明,我们提出的贡献在不同的图像分割任务上的优点:全景分割、实例分割和语义分割。我们的DFormer在MS Coco val2017集合上的表现优于最新的基于扩散的全景分割方法Pix2Seq-D,收益为3.6%。此外,在ADE20K Val集合上,DFormer获得了良好的语义分割性能,比最近的基于扩散的方法高出2.2%。我们的源代码和模型将在https://github.com/cp3wan/DFormer.上公开
1 介绍
图像分割的目的是将图像中的像素分成不同的单位。分组有不同的表示法,例如语义类别或实例。因此,人们提出了多种图像分割任务,包括语义分割[1]、实例分割[2]、全视分割[3]等。语义分割将像素划分为不同的语义类别,而实例分割将像素划分为不同的实例。另一方面,泛视分割不仅将物体的像素分组到不同的实例中,而且还将物体分割成不同的语义类别,这可以看作是语义分割和实例分割的联合任务。
在过去的几年里,研究人员针对不同的图像分割任务提出了专门的架构。例如,全卷积网络的不同变体[1,4,5]被用于语义分割中的像素级分类。检测-然后分割[2,6,7]和端到端[8-10]架构用于实例分割,而分裂-然后合并管道[11,12]旨在结合语义和实例分割进行全视分割。虽然这些专门的方法在每个单独的分割任务中都取得了巨大的成功,但它们通常难以有效地推广到不同的图像分割任务中。
为了解决上述问题,很少有现有作品[13-15]试图开发一种通用架构,通过统一的框架执行不同的图像分割任务。启发
通过基于变压器的目标检测器,DETR[16],这些方法将图像分割视为端到端集合预测问题。例如,K-Net[13]采用了一组可学习的内核来动态分割实例和内容,而Mask2Former[17]引入了一个掩码-注意力掩码转换器来进行掩码分类和掩码预测。近年来,基于扩散模型的图像分割方法也得到了探索。例如,Chen等人[18]采用位扩散进行全光分割,Ji等人[19]提出将噪声和深度特征拼接进行语义分割。然而,基于扩散模型的图像分割方法仍然落后于现有的通用图像分割方法。在这项工作中,我们研究了设计一个有效的基于扩散模型的变压器方法的问题,该方法可以实现具有竞争力的通用图像分割性能。
我们提出了一个名为DFormer的扩散引导变压器框架,用于通用图像分割。我们的DFormer将图像分割视为噪声蒙版的生成任务。在训练过程中,我们将不同程度的高斯噪声加入到真值掩模中,得到噪声掩模。然后,利用二值阈值生成注意力蒙版,并对噪声蒙版和深度像素级特征进行聚合得到蒙版特征。然后,我们将它们馈送到变压器解码器,以预测具有掩码注意的每个掩码特征的接地真值掩码。在推理时,我们首先生成一组带噪声的掩模,然后使用基于扩散的解码器来预测掩模和相应的对象类别。
我们对三种不同的图像分割任务进行了实验:语义分割、实例分割和全景分割。我们的实验结果表明,在单一架构的三个分割任务上,所提出的贡献具有令人满意的性能(见图1)。使用骨干网ResNet50,我们的DFormer在MS COCO val2017集上实现了51.1%的PQ分数,从而优于最近基于扩散的全光分割方法Pix2Seq-D[18],绝对增益为3.6%(见图1)。在主干swing - t下,我们的DFormer在语义分割ADE20K值集上的mIoU得分达到48.3%,优于最近基于扩散的语义分割方法DDP[19]。
2 相关工作
语义分割。语义分割旨在将图像中的像素分组到不同的语义类别中(例如,人,汽车和道路)。随着深度学习的出现,语义分割已经取得了很好的进展。FCN[1]是最早用于语义分割的工作之一,它采用全卷积网络进行像素级分类。之后,研究人员开发了不同的FCN变体来改进语义分割。例如,一些研究侧重于采用编码器-解码器结构的局部上下文信息聚合[5,21,22]或空间金字塔结构[23,4]。此外,一些作品利用了非局部上下文聚合的注意机制[24-26],而其他作品则着眼于实时设计[27,28]。
最近,文献中提出了基于变压器的方法,如基于变压器的网络设计[29-31]和基于变压器的分段头[32,33]。
实例分割。实例分割对不同的对象实例(例如,人和车)进行分类和分割,并通过两阶段、单阶段和端到端方法来解决。两阶段方法[2,34,35]通常建立在两阶段目标检测器的基础上,首先生成一些候选区域建议,然后在裁剪的区域建议中预测目标掩码。
与两阶段方法相比,单阶段方法[36-38]直接使用全卷积网络进行掩码预测。最近,随着端到端目标检测器(如DETR[16])的成功,一些工作已经探索了端到端实例分割。queryst[39]为基于查询的实例分割设计了一个动态掩码头。SOLQ[40]为实例类、位置和掩码构建了统一的查询表示。SOTR[41]将双注意机制发展为实例分割的变压器。此外,FastInst[42]和SparseInst[43]侧重于高效设计。
展示全景的分割。Panoptic segmentation[3,11]将实例分割和语义分割结合为一个统一的任务。首先,一些研究[44,45,12,46,47]探讨了在两个子任务中使用不同分支,并设计融合模块将这些分支的结果结合起来进行全视分割。例如,Panoptic FPN[11]将两个不同的头网络连接到共享特征金字塔上,并使用后处理操作将两个头网络的结果组合在一起。这些方法虽然取得了可喜的结果,但相对复杂。之后,一些作品[48-51]探索将物与物的分割统一成一个单一的网络。例如,Panoptic FCN[48]首先为thing和stuff生成核,然后使用动态卷积进行预测。Panoptic SegFormer[50]引入了一种查询解耦策略,其中thing和stuff的预测基于不同的查询。
通用图像分割。通用图像分割旨在将语义、实例和全景图像分割任务统一到一个网络中。受视觉转换器在目标检测中的成功[16]的启发,研究人员[14,13]开始探索一种适用于不同图像分割任务的通用架构。K-Net[13]学习了一组可学习的内核来预测实例和语义类别。MaskFormer[14]和Mask2Former[17]将图像分割视为掩码分类问题,预测掩码嵌入和像素级嵌入,并使用它们的点积进行掩码预测。
感知任务的扩散模型。扩散模型[52]旨在通过参数化马尔可夫链从噪声中恢复样本,在图像生成任务上取得了令人满意的结果。最近,一些研究开始探索感知任务的扩散模型。Chen等人[53]介绍了用于目标检测的扩散模型,该模型通过渐进去噪过程从噪声盒中恢复目标。Baranchuk等人[54],Brempong等人[55],Kim等人[56],Wu等人[57]借助扩散模型探索标签高效分割。Amit等[58]提出将深度特征和估计的分割图相加进行逐步细化。Wolleb等[59]和Ji等[19]提出使用扩散模型来改进分割。Chen等人引入了用于离散数据生成的位扩散[60]和用于全光分割的扩展位扩散[18]。然而,我们注意到这些基于扩散模型的方法在不同的图像分割任务上仍然落后于通用的图像分割方法。在这项工作中,我们着手弥合基于扩散模型的感知方法和通用图像分割技术之间的性能差距。
3 方法
3.1 预备期
扩散模型
扩散模型[52]学习一系列的状态转换,从噪声中生成高质量的样本。在训练过程中,首先通过在样本数据中逐渐加入高斯噪声进行前向扩散过程,然后学习一个模型,从有噪声的样本数据中预测原始样本数据。具体来说,扩散过程在任意时间步长t产生有噪声的样本数据xt
式中,时间步长t在1 ~ t范围内均匀分布,¯αt是随t的增加从1到0的单调递减函数,当时间步长t = t时,¯αt接近于0,噪声数据xT接近于I。噪声I采用标准正态分布采样。然后,训练一个网络模型,从xt预测任意时间步长t的样本数据x0,其中去噪损失可以写成
式中,f表示学习到的模型。在推理过程中,样本数据x0通过一系列状态跃迁逐渐从噪声数据xT中恢复。xT→xT−∆→…→x0[52,61],其中反复使用学习到的模型。
通用图像分割。图像分割可以看作是掩模分类任务,它预测图像x中存在的东西和事物的掩模和对应的类别,可以写成M, C = f(x)。最近,基于变压器的图像分割取得了不错的效果,可以表示为M, C = ftrans(fpixel(x), Q),其中,fpixel表示用于像素级特征提取的像素级模块,ftrans表示变压器模块,Q表示可学习查询的固定大小。在本文中,我们的目标是通过扩散模型将图像分割重新定义为一个去噪过程。我们没有使用一组可学习的查询,而是直接从噪声掩模中分割图像为M, C = fd(fpixel(x), I),其中I表示噪声数据,fd表示基于扩散的解码器。
图2:用于通用图像分割的DFormer的总体架构。我们采用像素级模块,包括骨干和像素解码器,提取多尺度特征。在训练过程中,我们通过在地真掩模中加入高斯噪声来进行扩散过程,并使用变压器解码器从噪声掩模中预测掩模和相应的类别。
3.2体系结构
图2显示了我们的扩散引导变压器的整体架构,名为DFormer,它包含了用于通用图像分割的扩散模型。接下来[17],我们的DFormer包含一个像素级模块,并用提出的基于扩散的解码器取代标准解码器。像素级模块生成金字塔特征fip, i = 1,2,3和像素级嵌入Fpixel,而我们基于扩散的解码器以噪声掩码Mn,金字塔特征fip, i = 1,2,3和像素级嵌入Fpixel作为输入,并预测掩码嵌入Fmask。
最后,我们通过像素级嵌入Fpixel和掩码嵌入Fmask的点积生成掩码,并通过将掩码嵌入馈送到MLP层来预测掩码的类别。
进行像素级模块。
给定输入图像x∈RH×W ×3,像素级模块提取金字塔特征F ip∈RH /Si ×W /Si ×C, i = 1,2,3,4,其中Si表示特征映射到输入图像的步幅,C表示特征通道数。像素级模块包含主干和像素解码器。具体来说,我们首先使用骨干网络,如ResNet[62]或Swin Transformer[63],提取低分辨率Flow∈R H /64 × W /64 ×C的深度特征。然后,我们使用像素解码器逐步上采样深度特征,生成各种分辨率的金字塔特征F ip, i = 1,2,3,4。将前三个连续的金字塔特征fip, i = 1, 2, 3馈送到连续的变压器解码器层生成掩模嵌入,最后一个金字塔特征f4p作为像素级嵌入Fpixel。
基于扩散的解码器。
基于生成的金字塔特征F ip, i = 1,2,3和噪声掩码Mn,我们的基于扩散的解码器旨在预测掩码嵌入Fmask。基于扩散的解码器堆叠了L个变压器解码器层,如[17]所示,其中每个解码器层包含一个掩码注意层、一个自注意层和一个FFN层。每个解码器层的输入是注意力掩码Matt,金字塔特征之一fip和掩码特征Fm。对于第一解码器层,通过对噪声掩码进行阈值处理生成注意掩码Matt, Matt = Mn > 0, 掩码特征Fm是用噪声掩码Mn作为权值对像素级嵌入Fpixel进行聚合生成的,可以写成
其中,favg表示全局平均池化操作,fnorm表示归一化操作。除第一解码器层外,使用当前解码器层的输出掩码嵌入Fmask作为下一解码器层的输入掩码特征Fm,使用当前解码器层的预测掩码作为下一解码器层的输入注意掩码Matt。交替使用金字塔特征F ip, i = 1,2,3作为连续解码器层的输入图像特征,以提高训练效率。对于输出掩码嵌入Fmask,我们使用MLP层来预测掩码分类分数,并使用像素级嵌入的点积来生成预测掩码。
3.3训练与推理
在训练过程中,我们通过扩散过程从真值掩模中生成噪声掩模,并训练模型从噪声掩模中恢复真值掩模。在推理过程中,我们随机生成噪声掩模,并使用学习到的模型预测掩模和类别。
3.3.1 训练
算法1介绍了我们的DFormer的训练细节,包括像素级模块,掩码填充和编码,掩码损坏,基于扩散的解码器和损失。
掩码填充和编码。我们首先执行掩码填充,因为不同图像中的对象数量不同。我们将地面真相面具Mgt填充到固定数量的填充面具Mpad上。
具体来说,我们生成一些随机的二进制掩码Mrand,并将其与原始的groundtruth掩码连接为Mpad = cat(Mgt, Mrand),其中Mpad∈RN×H×W。
然后,我们探讨了在添加噪声之前对填充掩码进行编码的两种不同策略:基于二进制和基于随机洗牌的编码策略。二进制策略直接将填充掩码编码为二进制输入(0或1),基于随机洗牌的策略将填充掩码编码为0到1,其中目标掩码像素随机采样从0.5到1,非目标像素随机采样从0到0.5。
掩码变体。掩码损坏的目的是通过在编码掩码Menc上添加高斯噪声,在任意时间步长t上生成噪声掩码Mn。我们首先将编码蒙版的比例从[0,1]重新缩放到[-b,b],其中b是一个比例因子。然后,我们根据Eq. 1将高斯噪声添加到编码的掩模中。附加噪声的尺度1−¯αt取决于时间步长t,其中较大的t对应于噪声的大尺度。添加噪声后,我们用[-b, b]的范围箝位噪声掩模。
损失公式。使用N个填充的噪声掩模,我们可以预测N个掩模和相应的分类分数。然后,考虑掩码和分类匹配成本,我们通过基于匈牙利的算法在预测和基础事实之间进行二部匹配。之后,我们可以计算出总的训练损失为
其中Lcls是用于分类的交叉熵损失,Lce和Lmask是用于掩码预测的二元交叉熵损失和骰子损失。λ1、λ2、λ3表示损失权值,分别为2.0、5.0、5.0。
3.3.2 推理
算法2给出了我们的DFormer的推理细节。我们首先生成高斯分布的噪声掩模Mn,然后用多个采样步长预测实例掩模。在每个采样步骤中,我们首先根据前一个时间步长的噪声掩码预测实例掩码,然后在预测的实例掩码中加入噪声,生成当前时间步长的噪声掩码。
表1:MS-COCO val2017集合133个类别的全视分割对比。我们的DFormer比最近基于扩散的方法Pix2Seq-D[18]高出3.6%。
表2:MS-COCO val2017和80个类别的测试开发集的实例分割比较。我们的DFormer优于最近基于扩散的方法DiffusionInst[20],在val2017上的绝对增益为5.3%,在测试开发上的绝对增益为6.2%。
4 实验
4.1 数据集和指标
我们采用大规模数据集MS COCO[65]进行实例分割和全视分割,采用大规模数据集ADE20K[66]进行语义分割。
MS Coco:MS Coco[65]是常见视觉任务的最大数据集之一,例如对象检测、实例分割和全景分割。对于目标检测和实例分割,有80个目标类别。对于全景分割,有80个物体(物)类别和53个填充类别(即总共133个类别)。它包含Train 2017、val2017、test-dev集合,其中Train 2017集合有11.8k个图像,val2017集合有5k个图像,而test-dev集合有20k个图像。我们采用AP、AP50、AP75、APS、APM、APL作为实例分割的评价指标[65],采用PQ、PQTh、PQST作为全景分割的评价指标。
ADE20k。ADE20K[66]是最大的语义分割数据集之一,包含150个语义类别。数据集包含训练集、val集和测试开发集,其中训练集有20k个图像,val集有2k个图像,测试集有3k个图像。我们采用mIoU作为语义分割的评价指标。
4.2 实现细节
我们采用在ImageNet上预训练的深度模型ResNet50[62]或swing -t[63]作为主干。像素级解码器生成的特征映射F p1、2、3、4的步长分别为32、16、8、4特征通道数为256个。在基于扩散的解码器中,有L = 9级联解码器层。在训练过程中,对每个解码器层进行辅助损失监督,并将噪声掩码填充为100。在推理过程中,我们只保留最后一层译码器的预测掩码。对于不同的分割任务,我们采用后处理的方法得到预期的输出格式,如[14,17]。下面总结了不同图像分割任务的实现细节。
实例和全景分割。我们的方法在4个NVIDIA RTX3090 gpu上进行训练,批处理大小为12个。在训练过程中,我们采用AdamW作为优化器和步进学习率策略,总共有50个epoch。初始学习设置为1 × 10−4,权重衰减设置为0.05。学习率在总epoch的0.9和0.95的分数上降低了10倍。我们采用大规模抖动增强策略,如Mask2Former[17]。在推理过程中,我们采用标准的单尺度测试,将图像的短边缘调整为800像素,长边缘调整为1333像素。
语义分割。我们的方法在4个NVIDIA RTX3090 gpu上进行训练,批处理大小为16个。在训练过程中,我们采用AdamW作为优化器和多学习率调度,共进行了16万次迭代。初始学习设置为1 × 10−4,权重衰减设置为1 × 10−4。我们采用随机尺度抖动进行数据增强,如Mask2Former[17]。在推理过程中,我们将短边图像的大小调整为512像素。
4.3 先进水平对比
展示全景的分割。我们将我们的DFormer与表1中设置的MS COCO val2017的最先进的全光分割方法进行比较。所有的方法都使用ResNet-50作为主干。与最近基于扩散的方法Pix2Seq-D[18]相比,我们的DFormer在参数效率方面取得了显著的改进。DFormer在PQ上的绝对增益为3.6%,PQT h上的绝对增益为4.4%,PQSt上的绝对增益为2.5%,优于Pix2Seq-D。此外,我们的DFormer具有更少的参数和训练周期。与其他先进的方法相比,我们的DFormer实现了具有竞争力的性能。
实例分割。我们将我们的DFormer与表2中MS COCO val2017和测试开发集上的几种最先进的实例分割方法进行比较。所有的方法都使用ResNet-50作为主干。与最近基于扩散的方法DiffusionInst[20]相比,我们的DFormer在性能上有了显著的提高。就总体AP而言,我们的DFormer在val2017集上的绝对增益为5.3%,在测试开发集上的绝对增益为6.2%,优于DiffusionInst。与其他先进的方法相比,我们的DFormer实现了具有竞争力的性能。例如,我们的DFormer在测试开发集上比Mask Transfiner[70]高出3.9%。
语义分割。最后,我们将我们的DFormer与表3中ADE20K值集[66]上的几种最先进的方法进行比较。我们展示了使用骨干ResNet50和swing - t的结果。当使用骨干swing - t时,我们的DFormer优于最近基于扩散的方法DDP[19],绝对增益为2.2%。此外,我们的DFormer与文献中的现有方法相比表现良好。与K-Net相比[13],使用相同的主干,DFormer实现了2.5%的绝对增益。此外,在使用相同的swing - t骨干时,与最近的Mask2Former[17]相比,DFormer的AP值为48.3,性能也有所提高。
表3:在155个类别的ADE20k val集合上的语义分割对比。我们的方法也比最先进的方法有竞争力。
表4:消融研究。除了(c),我们在MS COCO val2017上进行了实验,使用ResNet50进行实例分割。(c)中,我们展示了不同模型在不同任务上的结果。
4.4 消融研究
解码器层。表4a给出了不同解码器层数的影响。当数字为3时,其mask AP得分为41.1,当数字为9时,其mask AP得分为42.6。因此,我们在最终模型中采用了9个解码器层。
推理期间的时间步长。表4b给出了推断过程中不同时间步长的影响。它几乎达到了最好的性能使用一步向前的过程。多步前进过程提供了性能的轻微改进。
不同的骨干负责不同的任务。表4c显示了不同模型对不同任务的影响。swwin - t在这三个分割任务上都比ResNet-50有更好的性能。
噪声掩模的数量。表4d显示了不同噪声掩模数量在推理过程中的影响。默认设置是使用100,与训练相同。我们观察到,较大数量的噪声掩模并不影响性能。
掩码编码和比例策略。表4e显示了不同掩码编码和缩放策略的影响。如前所述,我们介绍了两种不同的掩码编码策略,包括二进制和基于随机洗牌的策略。当掩码尺度b等于1时,基于随机洗牌的编码比二进制编码获得更好的性能。当掩码尺度b = 0.1时,随机洗牌策略略劣于二值洗牌策略。因此,我们采用尺度为0.1的二元策略。
定性的结果。图3给出了我们的DFormer在三个图像分割任务上的定性结果示例。我们的DFormer在具有挑战性的场景下为这些不同的任务提供了高质量的分割结果,例如分割遮挡对象(例如,第2行,第2列)和小规模对象(例如,第3行,第4列)。材料。
图3:三个图像分割任务的定性结果。DFormer在不同挑战场景下的三个任务上实现了相同架构的高质量分割。
5 结论
我们提出了一种用于通用图像分割的扩散引导变压器。我们的DFormer将通用图像分割视为基于扩散的掩模分类。在训练过程中,我们首先通过在地面真值掩模中加入高斯噪声进行扩散处理,然后学习一个变压器引导模型,从噪声掩模中预测地面真值掩模。在推理中,我们直接从高斯噪声中预测掩模和相应的类别。在全景分割、实例分割和语义分割任务上进行了实验,验证了DFormer的有效性。在所有三个图像分割任务中,Dformer都明显优于现有的基于扩散的方法。
局限性和未来方向。我们观察到,我们提出的扩散引导方法在小对象的实例分割上略低于最先进的方法。一个潜在的未来方向是探索利用扩散模型提高小物体的性能。
V 扩散过程的可视化
图4显示了在真值掩模中加入高斯噪声后的扩散过程的可视化。
总时间步长T等于1000。当时间步长t = t时,噪声掩模为高斯噪声。随着t的减小,噪声掩模变得更加清晰。当时间步长t = t时,噪声掩模为地真掩模。我们的DFormer使用单一架构在不同任务上呈现高质量的分割结果。
B 视频实例分割的扩展
我们通过添加跟踪分支将DFormer扩展到视频实例分割。表5显示了视频实例分割的结果。DFormer在视频实例分割方面也展示了具有竞争力的性能。
C 分割结果更加可视化
图5是我们的DFormer的全景分割结果,图6是我们的DFormer的实例分割结果,图7是我们的DFormer的语义分割结果。
图5:COCO2017 val集上的全视分割可视化。
图6:COCO2017 val集上的实例分割可视化。
图7:ADE20K值集上的语义分割可视化。