点击下方卡片,关注「AI视界引擎」公众号
( 添加时备注:方向+学校/公司+昵称/姓名 )
低分辨率图像分割在机器人技术、增强现实以及大规模场景理解等实际应用中至关重要,由于计算限制,这类应用中通常无法获得高分辨率数据。
为应对这一挑战,作者提出了MaskAttn-UNet,这是一种通过引入 Mask 注意力机制来增强传统UNet架构的新颖分割框架。该模型选择性地强调重要区域的同时抑制无关背景,从而提高在杂乱和复杂场景中的分割准确性。
与传统的UNet变体不同,MaskAttn-UNet能够有效平衡局部特征提取与更广泛的上下文 Aware ,使其特别适用于低分辨率输入。作者在三个基准数据集上对作者的方法进行评估,输入图像重新缩放到
,并在语义分割、实例分割和全景分割任务中展示了竞争力。
结果显示,MaskAttn-UNet在显著降低计算成本的情况下达到了与基于Transformer的方法相当的准确度,使其成为一个资源受限场景下高效且可扩展的低分辨率分割解决方案。
- Introduction
在复杂场景中实现准确的多类分割对于自动驾驶、机器人技术和增强现实等应用至关重要。例如,在自动驾驶车辆中,精确的像素级车辆和行人的标注对安全导航至关重要;而在工业机器人中,详细的工具和障碍物分割能够实现可靠的操控和有效的避障[31, 60]。然而,许多实际的视觉系统——从低成本的监控摄像头到无人驾驶航空器(UAV)和移动机器人——由于传感器限制和硬件局限性,通常在低分辨率图像下运行[1, 30, 63]。这种图像细节的减少给分割算法带来了重大挑战[8, 23, 36, 41],这些算法仍需要在这样的受限条件下准确地区分多个目标类别。因此,迫切需要能够在这些受限条件下保持鲁棒性和精度的分割方法。
类似于U-Net [52]的编码器-解码器架构通过其多尺度结构在提取局部特征和细微细节方面表现有效。然而,当单张图像中存在多个目标或类别时,它们往往难以捕捉长距离依赖关系 [15, 68],导致复杂场景中的模糊性 [51, 58]。相比之下,基于 Transformer 的视觉模型通过Self-Attention机制引入全局上下文,从而能够建模像素或区域之间的长距离关系 [55, 66]。这种全局表示需要大量的内存和计算资源,因为Self-Attention的二次复杂性可能导致这类模型不适合嵌入式或实时系统 [14, 22]。此外,由于视觉 Transformer 缺乏CNN固有的归纳偏置(尤其是局部偏置),完全依赖注意力机制的模型可能会忽略区分小型或重叠目标所需的微细化特征 [6, 50, 65]。这些局限性突显了需要一种平衡局部特征精度、全局上下文捕捉能力和计算效率的分割方法的重要性。
在本文中,作者提出了MaskAttn-UNet,这是一种创新性的U-Net框架的扩展,通过集成一种新颖的 Mask 注意力模块来解决上述挑战。MaskAttn-UNet架构保留了U-Net通过 Shortcut 捕获精细局部细节的优势,而 Mask 注意力模块则有选择地强调特征图中的重要区域,以注入更广泛的上下文信息。通过专注于相关区域(而不是全局关注所有像素),作者的方法可以更有效地捕捉长距离依赖性,并缓解通常与 Transformer 相关的内存负担。作者特别设计该网络用于低分辨率输入(如
的图像),这虽然大大减少了计算需求,但仍使模型能够学习丰富的表示。这种设计选择反映了现实世界中有限图像分辨率的实际应用场景,确保MaskAttn-UNet适用于资源受限的方案。
作者在语义分割、实例分割和全景分割的标准基准上评估了MaskAttn-UNet。尽管处理相对低分辨率的输入,作者提出的模型在平均交并比(mIoU)、全景质量(PQ)和平均精度(AP)等方面的表现与当前最先进的方法相当。值得一提的是,MaskAttn-UNet 在推理过程中保持了适度的内存占用,这使得它相比于许多具有相似准确性的全Transformer基模型,部署起来更为实际。这些结果表明,作者的混合方法有效结合了卷积归纳偏差和目标导向的自注意力的优点,从而在多样且复杂的场景中实现了稳健的多类别分割。
作者的贡献总结如下:
- 作者提出了一种新的MaskAttn-UNet,这是一种集成新颖的 Mask 注意力模块的自注意力U-Net变种,能够有效捕捉局部细节和长程依赖关系。
- 作者设计了适用于低分辨率分割的架构,使用
尺寸的输入,从而降低计算需求同时保持稳健的性能。·作者在多个数据集上验证了该方法,结果显示与基于Transformer的方法相比,在更低的内存消耗下实现了分割指标的改进。
这些贡献突显了结合卷积归纳偏置与目标注意力机制以在实际图表中实现准确高效的分割的价值。在后续部分中,作者将讨论启发作者方法的相关工作,包括U-Net扩展、视觉Transformer以及基于 Mask 的分割方法。
- Related Work
2.1. U-Net
U-Net [52] 引入了一种编码器-解码器架构,已成为图像分割领域的基石。其设计包括一个压缩路径,通过连续的卷积和池化操作提取多尺度特征,以及一个扩张路径,利用上采样层恢复空间分辨率。编码器和解码器相应层之间的跳接连接使网络能够合并深层语义信息与高分辨率的空间细节。该结构在医学影像分割等应用中表现出色,尤其是在需要精确定位的情况下 [5, 27, 52, 61]。
尽管U-Net在许多任务上取得了成功,但标准卷积层固有的固定感受野限制了它捕捉长距离依赖的能力[24, 44, 64]。在包含多个相互作用物体或重叠结构的场景中,这可能导致将不同区域错误分类或合并[42, 53, 62]。
为了应对这些局限性,已经提出了一些扩展。例如,Attention U-Net [43]引入了注意力门控机制以改进 Short-Cut 连接,使得模型能够有选择地强调相关特征。类似地,Residual U-Net [69]通过引入残差连接来促进更深网络的训练。虽然这些修改有助于改善梯度 Stream 和局部特征提取,但它们仍未彻底解决在整个图像上聚合全局上下文的挑战。
2.2. Swin Transformers
基于Transformer的模型由于能够通过自注意力机制建模长距离依赖关系而成为了图像分割的强大替代方案[33, 54, 66]。Swin Transformer [37] 在这一领域取得了显著进展,采用了分层结构。该模型将输入图像划分为不重叠的Patch,并在局部窗口内计算自注意力。一个关键创新在于逐层之间引入了移位窗口,这使得窗口间交互成为可能,并有效扩展了感受野范围,同时避免了全局自注意力带来的高计算成本。这种分层设计便于多尺度特征学习,使Swin能够在处理高分辨率图像时平衡局部细节和全局语境。
然而,这些好处会带来显著的计算成本。随着深度配置或输入尺寸增大,内存需求和处理需求增加,这在实时推理或资源受限的硬件上部署时可能是不可行的[29, 40, 59]。因此,分割性能与效率之间的权衡仍然是基于Transformer的方法的一个活跃研究挑战。
2.3. Mask2Former
Mask2Former [17] 在 MaskFormer [16] 所确立的框架基础上,通过将分割重新公式化为集合预测问题而进行改进。在 MaskFormer 中,分割是通过为每个目标实例分配一个独特的 Mask 来实现的,从而统一处理“物体”和“场景”两种类型。Mask2Former 通过引入动态注意力 Mask 对相关图像区域进行自适应聚焦,进一步完善了这一方法。这种机制使得模型能够有效地区分重叠的目标,并生成高质量的实例级和全景分割输出。
Mask2Former中的动态 Mask 生成有助于捕捉全局上下文信息,同时保持界定物体边界所需的灵活性。然而,仅仅依赖注意力机制可能会影响对细粒度局部细节的保留,而这些细节对于准确界定边界至关重要,尤其是在处理小物体或复杂纹理的场景中[4, 18, 70]。尽管存在这些挑战,Mask2Former在标准分割基准测试中仍表现出稳健的性能。其设计突显了基于 Mask 的注意力机制在实现局部精确性和全局上下文之间的平衡方面的潜力,尽管计算复杂度和训练数据需求仍然是需要进一步改进的地方。
- Methods
在本节中,作者详细描述了所提出的一种分割方法。该方法处理输入图像以生成像素级分类 Mask ,其中每个像素都被分配一个语义标签。为了扩展该方法的功能,作者还引入了实例分割和全景分割分支,这些分支利用共享特征表示和专门的损失函数。作者首先概述了模型的整体架构,然后解释了训练目标和优化过程。
3.1. Architecture Overview
MaskAttn-UNet网络采用编码器和解码器结构,并在多个尺度上集成了 Mask 注意力模块。编码器通过一系列逐步降低空间分辨率的卷积块提取分层特征。在每个尺度上,特征通过 Mask 注意力模块进行细化,该模块生成可学习的二进制 Mask 以抑制不相关信息区域并强调显著结构。 Shortcut 将相应的编码器和解码器层连接起来,有助于解码器恢复分割输出中的高分辨率细节。解码器逐渐上采样并融合特征(通过 Shortcut 进行增强),以生成最终预测。
图1提供了该架构的总体概述。在图1(a)中,展示了整体的U-Net风格处理Pipeline。图1(b)说明了 Mask 注意力模块的内部结构,该模块应用可学习的注意力 Mask 以增强特征表示。图1(c)详细展示了多尺度的编码器-解码器设计,包括每个分辨率 Level 上卷积层、 Shortcut 和 Mask 注意力块的排列方式。
3.2. Mask Attention Module
每个 Mask 注意力模块受到了多头自注意力机制的启发,并且额外包含了一个可学习的 Mask ,该 Mask 调节了注意力权重。给定来自编码器或解码器的输入特征图 (X),作者首先将其 Reshape 为 (X' \in \mathbb{R}^{B \times H \times W \times C}),其中 (B) 是批量大小,(H \times W) 是空间维度,而 (C) 是通道数。然后,作者应用多头 Mask 自注意力(在作者的实现中使用了四个头)。注意力权重通过缩放点积注意力机制计算,同时还增加了一个 Mask 矩阵 (M)。
其中,
,
,
,且
是 Query 向量和键向量的维度。在此,
是一个可学习的(或动态计算得到的) Mask ,用于抑制注意力矩阵中不包含有用信息区域的贡献。直觉上,
使注意力偏向关注相关空间位置。
对于给定的Head,注意力操作的输出然后在所有Head之间进行组合(如同多头注意力机制),并通过残差连接添加到原始输入中。令
表示 Mask 多头注意力操作(合并Head之后)的结果。作者通过一个包含GELU非线性的两层 FFN (FFN)处理
,并在最后添加残差
:
其中,
是全连接层的权重矩阵,
是偏置项。这生成了 Mask 自注意力模块的最终输出
。通过结合 Mask 自注意力和残差全连接层,可以在保留通过 Short-Cut 传递的原始信息的同时,增强特征表示并整合全局上下文。
3.3. Segmentation Loss
作者使用一个综合损失函数来优化网络,该函数结合了语义分割损失和实例 Level 对比损失。平衡这些目标使模型能够学习像素 Level 的类别区分以及实例特定的可分性。
语义分割损失。在语义分割中,每个像素属于
个类别中的一个,作者使用标准的交叉熵损失。令
表示像素
的真实类别标签,
表示像素
属于类别
的预测概率。损失函数为:
其中,
是克罗内克δ函数(当其参数为真时值为1,否则为0)。这种按像素的交叉熵鼓励每个像素正确的类别预测。
实例对比损失。对于实例分割(以及全景分割中的实例部分),作者采用一种对比嵌入损失来鼓励属于同一目标实例的像素具有类似的特征嵌入,同时将不同实例的像素嵌入推开。令
表示网络生成的像素
的嵌入向量。对于给定的像素
,定义其正像素索引集合
为与
属于同一 GT 实例的所有像素索引,定义其负像素索引集合
为与
属于不同实例的所有像素索引。首先,作者计算
考虑的所有配对的归一化因子
。
其中
是一个温度参数,用于控制对比分布的尖锐度。对于正样本对
和
(即来自同一实例的两个像素),每对的对比损失为:
这是一种惩罚机制,如果嵌入向量 (e_{ij}) 和 (e_{kl}) 在分子中没有显著更接近彼此,相比于所有其他配对在分母中的情况。像素 ((i,j)) 的实例对比损失通过对其所有正伙伴 ((k,l) \in P_{ij}) 计算 (l_{ij,kl}) 的平均值来计算,然后对所有像素进行平均。
其中 ( N ) 是考虑的总像素数(出于效率考虑,这可以是一个所有像素对的采样子集)。实际上,( L_{\mathrm{IC}} ) 促使同一实例的嵌入在特征空间中聚类,而不同实例的嵌入则保持分离。最终的分割损失是由这两部分按权重相加得到的:
其中,
控制语义分割损失和实例对比损失之间的平衡。在作者的实验中,作者调整
确保 neither 一项主导,从而使模型能够学习准确的像素级分类和良好的实例嵌入。关于选择
的详细信息,请参见附录 B.1。
- Experiments
作者在三个常用的分割基准上评估了MaskAttn-UNet,报告了其语义分割、实例分割和全景分割的结果。然后,作者与最先进的方法在全景分割上进行对比,并通过不同训练数据集的比例来检验其性能。这些实验支持了作者设计的 Mask 注意力模块的选择,并展示了MaskAttn-UNet在不同数据集上的泛化能力。
数据集。作者采用了三种广泛使用的多任务图像分割基准数据集。COCO [35] 是一个大规模数据集,包含80个“thing”物体类别和多个“stuff”(背景)类别,支持语义分割、实例分割和全景分割任务。ADE20K [71] 包含150个语义类别(100个“thing”和50个“stuff”),用于语义分割、实例分割和全景分割任务。Cityscapes [19] 专注于城市街道场景,包含19个类别(8个“thing”和11个“stuff”),常用于自动驾驶场景中的语义和全景分割任务。所有图像均调整至
像素以减少计算开销,并模拟某些实际应用场景中遇到的低分辨率条件。
评估指标。作者为每项任务使用了标准指标。对于语义分割,作者报告了平均交并比(mIoU)[21],用于衡量预测区域与ground-truth区域之间的类别平均重叠程度。对于实例分割,作者使用不同交并比阈值下的平均精度(AP)[35],以评估单个目标实例的检测和分割效果(较高的AP表示更好的精确召回贸易)。
对于全景分割,作者报告了全景质量(PQ)指标[34],该指标综合了“thing”和“stuff”类的识别质量和分割质量。除了
和在
阈值下的AP之外,这些指标为每个数据集上的分割性能提供了全面的评估。
4.1. Implementation Details
作者的实现基于一个具有四层下采样编码阶段和四层上采样解码阶段的U-Net Backbone 网络,并通过 Shortcut 恢复空间细节。编码器逐步增加特征通道数,从64增加到128,再到256(包含两个256通道的块),进入瓶 Neck 分以压缩并细化全局上下文。解码器则对称地上采样,减少通道维度,并从相应的编码器层级合并特征图,以在输出中重建精细的空间细节。
在每一层编码器和解码器中,作者引入了一个MaskAttn模块,具体参见第3.2节。每个模块包含一个可学习的二值 Mask ,可以选择性地抑制不相关的激活。这使得网络能够将注意力集中在重要的区域,例如物体边界和显著结构上,即使是在低分辨率特征图上也是如此。通过在整个网络中集成这些模块,MaskAttn-UNet保留了卷积的地方归纳偏置,同时获得了在每个尺度上捕捉长程依赖性的能力(有关更详细的分析,请参见附录B.3)。
4.2. Training Settings
损失函数。作者通过在第3.3节中定义的损失函数分别训练模型的所有分支(语义、实例、全景)。对于语义分割,作者应用交叉熵损失(式3),它提供了强大的像素级分类信号。对于实例和全景分割,作者包括实例对比损失(式6),以学习不同的目标嵌入。实际上,作者还为每个预测实例添加了标准分类损失项(用于预测其语义类别),确保每个实例嵌入与特定类别相关联。通过平衡
和
的贡献,在式7中选择了
,模型能够同时学习类别的区分和细粒度的实例分离。
训练设置。作者在每个数据集上训练模型共1000个epochs。所有图像均统一调整为
大小。对于语义分割实验,作者使用配备了
NVIDIA V100 GPU的系统,批量大小设为8。对于更耗内存的实例和全景分割实验,作者则使用配备了
NVIDIA A100 GPU的系统,批量大小设为14。作者采用Adam W优化器进行模型优化,初始学习率为
,权重衰减率为
。作者使用数据增强技术,包括随机尺度抖动和水平翻转,以提供输入多样性而不过度复杂化训练分布。这些训练配置旨在平衡每项任务的内存使用率和吞吐量,从而在不同的分割目标下实现一致的收敛。
4.3. Main Results
语义分割。作者使用COCO panoptic-val2017、ADE20K val和Cityscapes val数据集评估了作者的语义分割模型,所有标签均针对语义分割进行了专门处理。如表1所示,MaskAttn-UNet在COCO上的mIoU为
,在ADE20K上的mIoU为
,在Cityscapes上的mIoU为
。这些结果表明,尽管输入分辨率较低,该网络能够有效地融合全局上下文信息与局部空间细节。图2(a)和图2(b)左侧两列中的分割图显示,目标边界得到了很好的保留,并且纹理复杂的区域被清晰地分割开。特别是对于包含重叠物体或精细结构细节的区域,该机制也表现出了很好的处理能力。这表明 Mask 注意力机制成功地抑制了背景噪声,并增强了学习特征的区分性。
在不同照明和对比度条件下(图2(b)中的第一张图片),网络能够维持较高的准确性,确保既能准确标注大面积的均匀区域,也能精确标注细小复杂的细节。 Mask 注意力模块的设计似乎能够有选择地增强重要特征,同时减少来自不informative区域的干扰,从而在不同的类别和挑战性场景下获得更一致的预测结果。
实例分割。实例分割性能分别在COCO val2017、ADE20K val和Cityscapes val数据集上进行了评估,并对所有标签进行了细化以适应任务需求。表2总结了MaskAttn-UNet的表现。在COCO数据集上,模型在IoU阈值为30时的平均精度(AP)为35.0%,而在IoU阈值为100时下降到30.2%。在ADE20K数据集中,对于IoU阈值分别为30、50、70和100的情况,对应的AP值分别为33.8%、33.2%、30.5%和30.5%;而在Cityscapes数据集上,对应AP值分别为38.9%、36.6%、36.2%和35.5%。这些结果表明,MaskAttn-UNet能够在重叠或密集排列的目标场景中可靠地隔离并划分出独立的目标实例。
调整IoU阈值的效果揭示了网络处理不同实例复杂性的能力。较低的IoU阈值主要捕捉最突出的目标,导致更高的AP值;而较高的IoU阈值则扩展到更小或部分遮挡的实例,这会导致AP略有下降。图2(a)和图2(b)右两列中的实例分割输出可视化表明, Mask 注意力模块有助于在多个尺度上细化特征图,从而改善相邻目标的区分和细微结构细节的识别。例如,图2(b)中第三张图片的分割结果显示,MaskAttn-UNet减少了重叠目标和复杂边界区域的误检测。网络在拥挤区域保留了目标轮廓,保持了个体实例的一致区分。这些结果表明,所提出的设计架构能够很好地适应各种真实世界条件,增强了其在实际实例分割应用中的有效性。
全景分割。在COCO panoptic-val2017、ADE20K验证集和Cityscapes验证集上评估了全景分割性能。表3展示了结果,其中MaskAttn-UNet在"Cuff’区域上达到了平均交并比(mIoU)为45.3%,在"thing"实例上的平均精度(AP)为31.5%,并在COCO数据集上的全景质量(PQ)为35.7%。在ADE20K上,模型达到了mIoU为45.9%、AP为30.7%以及PQ为33.6%。类似地,在Cityscapes上,它记录了mIoU为70.1%、AP为35.5%以及PQ为58.3%。这些结果表明,MaskAttn-UNet能够在前景目标和背景区域上提供均衡的分割性能。
Cityscapes上较高的mloU值表明网络有效地利用了城市场景的结构特性,实现了一致的背景分割。同时,COCO和ADE20K上稳定的AP和PQ值证明了其处理更复杂多变环境的能力。稳健的背景分割与精确的目标轮廓划定相结合,有助于实现较高的整体PQ值,进一步证实了网络具备全面的场景理解能力。
与 Baseline 模型的比较。为了验证作者方法的稳健性,作者将模型在COCO、ADE20K和Cityscapes三个数据集上与几个具备相当或略高复杂度的最新模型进行了基准测试(见表4)。例如,相较于仅含有32M参数且计算量为4G FLOPs的U-Net,MaskAttn-UNet在COCO数据集上的mIoU提升了超过10%,PQ提升了15%,AP提升了近20%,而参数量仅增加了约15M。尽管U-Net较低的FLOPs适用于简单的场景,但在存在重叠物体和复杂纹理的环境中,它表现不佳。相比之下,MaskAttnUnet拥有11G FLOPs和46M参数,在捕捉细粒度的空间细节及处理复杂场景方面表现更为出色。
作者还评估了基于DETR的模型,包括DETRR50(86G FLOPs,41M参数)和DETR-R101(152G FLOPs,60M参数)。基于DETR的模型同样需要增加计算量,仅能获得适度的性能提升,这限制了它们在资源受限或实时场景中的应用。
此外,Mask2Former-R50 需要高达
FLOPs 的计算量,而采用 ResNet101 后端的 Mask2Former-R101 则需要更多的计算资源。值得注意的是,Mask2Former-R101 在 ADE20K 数据集上的 PQ 值仅比 MaskAttn-UNet 高出
,在 Cityscapes 数据集上的 mIoU 值高出
。然而,这些微小的改进却伴随着巨大的成本:Mask2Former-R101 比 MaskAttn-UNet 多使用了
参数(大约增加了
),这表明单纯增加计算资源并不一定会显著提高分割质量。总体而言,MaskAttn-UNet 达到了精确度和效率之间的良好平衡,因此在计算资源受限的应用场景中是一个实用的选择。
少样本训练。在许多实际场景中,收集大量的标注数据既昂贵又耗时。为了评估训练数据量对分割性能的影响,作者在COCO panoptic-val2017数据集上评估了MaskAttn-UNet模型,并使用全数据集的不同比例(
、
、
、
和
)进行模型训练。结果显示如图3所示的趋势:
10% 训练数据:模型实现了36.7% 的mIoU、25.3% 的PQ 和22.6% 的AP。尽管这些指标并不理想,但它们表明网络在数据稀缺条件下仍能够提取有意义的特征。25% 训练数据:性能提升至37.3% 的mIoU、27.6% 的PQ 和25.1% 的AP,这显示出适度增加训练数据会导致显著提高分割精度。50% 训练数据:模型达到40.1% 的mIoU、33.3% 的PQ 和29.4% 的AP,表明使用完整数据集的一半足以学习稳健的特征表示。75% 训练数据:指标进一步上升至43.4% 的mIoU、35.1% 的PQ 和30.1% 的AP,这确认了额外的数据继续为模型带来益处,同时保持高效的训练动态。100% 训练数据:利用整个数据集,模型达到了45.3% 的mIoU、35.7% 的PQ 和31.5% 的AP,突显了其充分利用可用标注以获得最佳分割结果的能力。
这些趋势突显了MaskAttn-UNet 强大的数据效率,使其在标注数据有限的情况下成为一个实用的选择。值得注意的是,随着可用数据量的增加,该模型的性能显著提升,这与描述神经网络性能随训练集规模变化的传统神经网络扩展定律相一致[2, 32]。这种数据效率在医疗影像计算及其他难以获取大规模标注数据的领域中尤为重要[48, 56]。因此,MaskAttn-UNet 能够在减少训练数据的情况下有效运行,使其成为数据受限环境中的一个可行解决方案。
- Discussion
作者提出了MaskAttn-UNet,这是一种通过将 Mask 注意力模块整合到传统的U-Net架构中而在分割模型方面的重要进展,有效增强了局部和全局特征提取能力。这种混合方法利用了卷积网络在建模局部上下文方面的优势以及 Mask 注意力机制处理长范围依赖的能力。
作者在COCO、ADE20K和Cityscapes等数据集上的实证评估表明,MaskAttn-UNet在计算资源使用方面显著少于基于Transformer的架构(如Mask2former)的同时,持续优于标准的U-Net模型。这些发现突显了选择性注意力机制在低分辨率分割任务中的潜力,填补了卷积效率与Transformer模型的全局上下文感知之间的差距。
参考
[0]. MaskAttn-UNet: A Mask Attention-Driven Framework for Universal Low-Resolution Image Segmentation .
点击上方卡片,关注「AI视界引擎」公众号