首篇!全面解读高效Segment Anything模型变体:各种图像分割加速策略和核心技术展示 精华
文章链接:https://arxiv.org/pdf/2410.04960
亮点直击
- 提供了一项系统的高效 SAM 变体全面回顾,旨在加速分割任务。介绍了一个结构良好的分类法,将这些方法根据其采用的加速策略进行分类。据我们所知,这是第一项专门关注该领域的调查。
- 对这些变体的效率和准确性进行了全面的评估和比较,旨在帮助研究人员选择最能满足其性能和应用需求的模型。
- 提出了几个未来研究的潜在方向,为读者提供启发,以推动该领域的持续发展。
Segment Anything Model (SAM) 是一款基础模型,专用于图像分割任务,以其在多样应用中的强大泛化能力而闻名。然而,它的出色性能伴随着显著的计算和资源需求,这使得在资源有限的环境(如移动设备)中进行部署变得具有挑战性。为了解决这个问题,已经提出了多种 SAM 变体,以提高效率而不牺牲准确性。
本综述首次全面回顾了这些高效的 SAM 变体。我们首先探讨推动该研究的动机,然后介绍 SAM 和模型加速中使用的核心技术。接下来,深入分析各种加速策略,并按方法进行分类。最后,提供了对这些方法的统一和全面的评估,评估它们在代表性基准上的效率和准确性,并对它们的整体性能进行明确比较。
1 引言
基础模型的出现已经彻底改变了人工智能(AI)领域。基础模型是大型神经网络,在海量数据上进行预训练,具有强大的表示能力和出色的泛化能力,能够在各种任务中表现出色。在自然语言处理(NLP)领域,最近流行的研究趋势是大型语言模型(LLMs),如OpenAI的GPT系列、Google的PaLM系列和Meta的LLaMA系列等,这些模型取得了显著的发展。同时,视觉转换器(ViT)的成功首次将转换器架构引入计算机视觉(CV)领域,为视觉基础模型(VFM)开辟了新纪元。像CLIP、LLaVA和Video-ChatGPT等视觉-语言基础模型旨在对齐视觉和语言模态,在众多下游视觉任务中表现出色。
最近,Meta提出了一种新型的基础模型,用于通用图像分割,即 Segment Anything Model (SAM) 。SAM 完全在他们提出的 SA-1B 数据集上进行训练,该数据集包含超过十亿个mask和一千一百万张图像,旨在通过任何提示(例如点、框、mask和文本)实现有效的分割。SAM 能够很好地泛化到广泛的下游任务(如边缘检测、目标提议和实例分割),只需提供适当的提示。SAM 的出现引起了研究界的广泛关注,并催生了大量相关工作,探索 SAM 在各种场景中的泛化能力,包括不同的图像分割任务、视频分析任务和 3D 视觉任务。随着 SAM 的巨大成功,升级版的 Segment Anything Model 2 (SAM 2)
尽管 SAM 在广泛应用中取得了成功,原始的 Segment Anything Model (SAM),特别是 SAM-H,由于运行速度慢和计算成本高面临显著的局限性。在资源受限或实时环境(如边缘设备和移动应用)中部署 SAM 时,这些挑战尤为明显。随着对在实际资源受限场景中部署机器学习模型的需求不断增加,SAM 当前的设计在广泛使用中显得效率低下。这导致了对更轻量级、高效变体的迫切需求,这些变体可以在满足模型强大分割能力的同时应对这些限制。优化 SAM 以提高效率的挑战进一步加剧,因为在实时应用、移动平台和嵌入式系统中对计算资源的需求越来越高。随着研究界努力克服这些障碍,全面了解在提高 SAM 效率方面的最新进展变得至关重要。因此,开展一项详细的调查,以研究提高 SAM 效率和扩展其在多样环境中的适用性的努力,是及时且必要的。
随着与 SAM 相关的研究不断增加,已有若干调查从不同角度提供了概述。然而,这些现有的调查主要集中在 SAM 的下游应用上,并存在几项局限性:
- 没有一项调查关注到提高 SAM 效率的新兴领域,而这一领域正在获得重要关注,对实际应用的部署至关重要。
- 除了一项外,这些调查缺乏结构化的分类法,以便更清晰地组织和参考。
- 大多数先前的调查主要集中在收集和描述基于 SAM 的方法,但缺乏对这些方法的系统性评估或比较。
为了解决这些不足,本文开展了本次调查,旨在全面回顾高效 Segment Anything 模型的发展,并公平地评估和比较它们。
本调查的其余部分组织如下。在第2节中,我们首先介绍原始 SAM 的背景,然后回顾可以应用于提高 SAM 效率的高效视觉表示骨干和模型压缩技术。在第3节中,我们根据目标和技术对现有方法进行分类,详细回顾每个类别。我们还讨论了进一步加速 SAM 的几个潜在研究方向。在第4节中,我们对这些模型在效率、准确性和相应的权衡方面进行公平评估。最后,在第5节中,我们简要总结了本次调查。
2 初步概述
2.1 Segment Anything Model
Segment Anything Model (SAM) 是图像分割领域中的一个强大基础模型,旨在通过基本的可提示分割任务,利用提示工程来统一各种分割任务。该项目的一个显著贡献是 SA-1B 数据集,该数据集包含来自 1100 万张经过许可和隐私保护的图像生成的超过 10 亿个masks。在如此丰富和高质量的数据上进行训练,SAM 展现出强大的鲁棒性和泛化能力。SAM 的巨大潜力迅速引起了研究人员对探索其在广泛现实应用中的能力以及改进其架构以实现更高效或更准确分割的兴趣。
最近,Segment Anything Model 2 (SAM 2) 被提出作为继任者,专注于高效的可提示视觉分割 (PVS),适用于图像和视频。为了使 SAM 2 能够在视频中分割任何内容,研究人员在 SAM 的原始架构中引入了流式记忆机制。SAM 2 采用两阶段训练:
- 在 SA-1B 数据集上进行可提示分割任务的预训练;
- 在混合数据上进行可提示分割任务的训练,处理图像和视频。
与 SAM 类似,研究人员开发了一个数据引擎,以创建一个用于视频分割的大规模数据集,命名为 SA-V 数据集。该数据集最终收集了 50,900 个视频中的 642.6K 个mask(视频中的对象分割)。在本次调查中,将 SAM 2 视为高效的 SAM 变体,并将其纳入评估和比较中。
2.1.1 模型
SAM 由三个主要组成部分构成:图像编码器、提示解码器和mask解码器,如下图 1(a) 所示。图像编码器是一个经过微小修改的 MAE预训练视觉Transformer (ViT)。它将预处理后的图像作为输入,并为每张图像输出一个图像嵌入。提示解码器用于嵌入提示:点、框、mask和文本。这两个嵌入然后输入到轻量级mask解码器,该解码器基于两个修改过的Transformer解码器块和一些预测头,生成有效的mask。
基于 SAM 的架构,SAM 2 额外引入了流式记忆机制。具体来说,这个机制包括一个记忆编码器、一个记忆库和一个记忆注意模块。SAM 2 的结构如上图 1(b) 所示。通过记忆机制,SAM 2 能够逐帧处理视频。记忆编码器生成当前帧的预测记忆并将其发送到记忆库。记忆库存储最近的预测、提示帧的特征图以及目标对象的高级语义信息(即对象指针)。记忆注意机制使来自图像编码器的图像嵌入与来自记忆库的信息充分交互,从而得到精细化的嵌入。除了记忆机制,SAM 2 还采用 MAE 预训练的 Hiera作为图像编码器,这比 ViT-H 更高效,预期可以实现更快的速度。
2.1.2 任务
可提示分割任务被提出作为 SAM 的基本任务,其目标是根据任何给定的提示(例如一个点、一个框、一个mask或文本)返回有效的mask。这个任务不仅是 SAM 训练过程中的目标,也是使 SAM 能够解决各种下游任务的基础。另一个重要任务是全mask生成,它会对图片中的所有对象进行分割。通过使用点网格对 SAM 进行提示,并根据这些稠密的提示预测mask来实现。全mask生成也是数据引擎最后阶段中的关键步骤,旨在增强 SA-1B 中mask的多样性。如下图 2 所示,可提示分割任务称为 Segment Anything (SegAny),而全mask生成任务则称为 Segment Everything (SegEvery)。这两个任务总结了 SAM 的分割能力,并为提升 SAM 效率指明了两个研究方向。在本次调查中,我们遵循这两个定义,探讨基于 SAM 的高效变体在 SegAny 和 SegEvery 任务中的表现。
2.1.3 应用
由于 SAM 及其继任者 SAM 2 在大量零样本下游任务中展示了强大的泛化能力,研究界正深入探索它们在更多场景中的应用。
SAM 的一个主要应用领域是医学图像分割 。根据[157]的说法,该领域的工作可以分为两类。一类旨在测试 SAM 在 CT 图像、MRI 图像、病理图像等中的分割性能。另一类则专注于通过微调、自动提示或框架修改来提高 SAM 在这些任务中的适应性。此外,像[35, 57]这样的研究试图提高医学 SAM 方法的效率。SAM 还被应用于不同现实场景中的对象检测,包括土木基础设施缺陷评估中的裂缝检测和陨石坑检测,农业中的作物病虫害检测,异常检测和遥感。此外,Segment Anything 已被改编为 Edit Everything、Inpaint Anything和 Caption Anything来处理图像编辑任务。
除了图像分割任务,SAM 还广泛应用于各种视频任务。大量研究集中于两个基本任务:视频对象分割 (VOS)和视频对象跟踪 (VOT)。研究人员还探索了 SAM 在生成相关任务中的应用,例如视频超分辨率和视频数据集标注生成。除此之外,SAM 还作为视频编辑任务中的一个中间工具被进一步利用。除了 2D 任务,SAM 还扩展到 3D 视觉领域。研究[100]将 SAM 应用于 3D 重建,而研究[28]则将其应用于 3D 点云分割。最近的工作[134]旨在实现在线设置下的任何 3D 物体的实时分割。对于最近提出的 SAM 2,已经有一些研究探索其在图像和视频任务中的应用。一个流行的趋势是将 SAM 2 应用于医学图像和视频任务。像[27, 99, 135]这样的研究评估了 SAM 2 在 2D 和 3D 模态下医学图像中的表现,而一些其他研究[75, 140]则测试了其在外科视频分割任务中的表现。研究人员还在寻找更好地将 SAM 2 适应医学任务的策略。此外,SAM 2 还被应用于一些特定的图像分割任务,如数字病理语义分割、网格部件分割和太阳能电池板分割。此外,一些研究[70, 113]利用 SAM 2 在具有挑战性的规模化视频对象分割 (LSVOS) 任务中取得了良好的结果。
2.1.4 局限性
尽管 SAM 在各种任务中表现出色,但在实际应用中仍面临两个关键挑战:
- 它通常难以为精细结构预测完整的mask,从而导致边界不精确;
- 它不是实时的,资源消耗仍然很高,特别是在使用像 ViT-H 这样的重型图像编码器时。
为了解决这些问题,研究如[49, 54]旨在通过利用高分辨率图像来改善mask质量,而其他研究[131, 147, 161, 164]则专注于创建更高效的架构,以减少 SAM 的时间和资源消耗。之前的调查[78, 100, 158]已探讨了提高 SAM 质量结果的最新进展。在本次调查中,我们专注于提高 SAM 效率的努力。
2.2 高效骨干网络
SAM 的低效主要源于其重量级的图像编码器。SAM 图像编码器的尺寸详见表 1,并在后面小节中提供了 SAM 总参数的进一步估算,比如在 SAM-H 中,ViT-H 图像编码器包含大约 632M 个参数,而模型的总大小为 641M,这意味着图像编码器占模型参数的大部分。
即使在最小的变体 SAM-B 中,图像编码器仍占总参数的 90% 以上。因此,快速有效地加速 SAM 的一种方法是用更高效的骨干网络替换大型图像编码器。这些高效的骨干网络可以包括纯卷积神经网络 (CNN)、高效视觉Transformer架构以及最近的Transformer替代模型[87]。
2.2.1 高效视觉Transformer
提高视觉Transformer效率的努力通常可以分为两种方法:
- 设计更高效的架构;
- 重构注意力机制。
从结构角度减少计算成本,MobileViT是一种开创性的混合架构,创造性地将卷积神经网络(CNN)模块(MobileNetV2 模块)与Transformer模块集成到一个模型中。随后,像 [10, 62, 127] 这样的后续工作基本上遵循这个思路,构建高效的混合结构视觉Transformer,这些Transformer已被广泛用来替代 SAM 的重型图像编码器。在 [103, 161] 中,TinyViT 作为高效骨干网络,而在 [162] 和 [159] 中,EfficientFormerV2 和 EfficientViT分别替代了 SAM 的原始图像编码器。另一个具有影响力的视觉Transformer设计,MetaFormer,将注意力机制抽象为一个更广泛的概念,称为 token mixer,可以使用各种 token mixer 实现与Transformer相媲美的性能。最简单的变种 PoolFormer,通过使用池化操作作为 token mixer,而不引入额外的可学习参数,已被用作为 Lite-SAM开发轻量级图像编码器的基础架构。
研究人员在优化注意力机制方面也取得了显著进展。观察到注意力机制中的 softmax 操作显著增加了整体计算成本。在 EfficientViT 中,提出了一种新颖的 ReLU 线性注意力机制,以更高效地实现全局感受场。这种高效的骨干网络在中被进一步采用以加速 SAM。注意力机制的改进还在硬件层面上进行。FlashAttention通过切片、内核融合和重计算等技术显著降低了计算成本,并被应用于 SAM 加速工作,以减少内存需求并提高计算效率。
2.2.2 替代Transformer模型
虽然Transformer目前在语言和视觉领域占据主导地位,但一些新提出的模型在效率和性能方面显示出了超越Transformer的潜力。
接受权重关键值(RWKV)模型结合了递归神经网络(RNN)和Transformer的优势,随着序列长度的增加,实现了线性时间复杂度。RWKV 非常适合处理长序列处理的挑战。
为了促进全局信息交互,RWKV 用更高效的 WKV 操作符和输出门控机制替代了传统的具有二次复杂度的注意力机制。这些机制的公式如下:
其中 r, k, v 分别表示接收的tokens、键和值的移位tokens,W表示权重。RWKV 也已扩展到视觉任务。视觉 RWKV (VRWKV) 模型在效率上表现出与视觉Transformer (ViT) 相当的性能。为了将 RWKV 从一维序列适应到二维图像,引入了 Q-shift tokens,以在四个方向上融合邻域信息。在 [145] 中,一种基于 RWKV 的 SAM 变体通过采用混合 MobileNetV2 块 和 VRWKV 块的高效主干,取得了卓越的效率表现。
2.3 模型压缩
模型压缩涵盖一系列旨在减少模型大小和计算复杂度的技术,使其在计算资源有限的实际应用中部署大型模型成为必要。模型压缩和加速的四种主要方法是知识蒸馏、量化、剪枝和低秩分解。
2.3.1 知识蒸馏
知识蒸馏 (KD)最初被引入作为在资源受限环境中部署大型复杂神经网络的解决方案。其核心概念是将来自较大、训练良好的模型(教师模型)的知识和表征能力转移到较小、更高效的模型(学生模型)。
在将 KD 应用于加速 SAM 时,目标是从原始的较大 SAM 中提取知识,并将其传授给更高效的 SAM 类模型。鉴于 SAM 的编码器-解码器架构,KD 通常可以分为两种方法:蒸馏整个 SAM 模型或仅蒸馏图像编码器。大多数工作,如 [84, 103, 117, 147, 159],专注于仅蒸馏高效主干,同时保留原始 SAM 的提示编码器和掩膜解码器。然而,其他方法,如 [162, 164],旨在通过监督编码器和解码器的输出来蒸馏整个模型。
2.3.2 量化
量化是将模型的高精度权重/激活值 X(例如 32 位浮点数)转换为低精度格式(例如 16 位浮点数、8 位整数)的过程。广泛使用的一种量化函数是均匀对称量化,其公式如下:
2.3.3 剪枝
模型剪枝通过消除冗余的权重或连接来减少模型的大小和复杂性,同时尽量保持精度。剪枝方法通常分为两种类型:结构化剪枝和非结构化剪枝 。结构化剪枝根据特定标准以组的方式删除参数,系统性地针对通道、层或块等子结构。相比之下,非结构化剪枝则专注于单个权重,通常导致稀疏和碎片化的网络。然而,由于剩余网络结构的不规则性,非结构化剪枝可能在一般硬件上并未带来有效的加速。在 [18] 中,应用结构化剪枝来减轻 SAM 的重量,通过移除大量冗余权重显著减少了模型的大小,同时保留了大部分 SAM 的能力。
2.3.4 低秩分解
3 SAM 的高效变体
本节回顾了自 SAM 受到关注以来,为开发轻量级高效的 SAM 类模型所做的努力。这些工作旨在降低模型的高计算成本并实现高效性能,同时保持 SAM 强大的分割能力和泛化能力。正如前面所述,SAM 解决了两个主要任务,包括 Segment Anything (SegAny) 和 Segment Everything (SegEvery)。因此,我们分别讨论旨在改进每个任务的研究:首先,关注加速 SegAny 任务,然后涵盖加速 SegEvery 任务的努力。值得注意的是,一些方法适用于这两个任务,将单独讨论这些贡献。此外,根据所采用的技术将所有模型分类,并在下图 3 中呈现方法论的分类。最后,概述了该领域未来研究的四个潜在方向。
3.1 加速 SegAny 任务
正如前面所分析的,SegAny 任务的主要瓶颈在于 SAM 的庞大架构。一种直接的解决方案是用更高效的主干替换编码器。或者,采用保持与 SAM 相同分割能力的不同架构也是一种方法。遵循这些策略的工作要么涉及完全从头开始训练轻量级模型,要么使用合适的监督进行知识蒸馏训练模型。此外,一些研究探索量化、剪枝或局部优化等方法,以直接压缩 SAM,而不替换编码器或构建新架构。
3.1.1 从头开始训练
本小节聚焦于完全从头开始训练 SAM 变体的工作。根据它们的架构,这些模型可以分为两种类型:与 SAM 架构不同的架构和类似 SAM 的架构。我们将按照这种分类详细探讨每种类型。
FastSAM 是第一个不依赖于 SAM 原始编码器-解码器架构的 SAM 变体之一。为了实现更快的分割,它将 SegAny 任务分为两个子任务:全实例分割和基于提示的选择。由于实例分割已被许多基于 CNN 的方法有效解决,FastSAM 提供了比基于Transformer的 SAM 更高的效率。对于实例分割,FastSAM 采用 YOLOv8-Seg 模型,并使用 YOLACT 方法来提升性能。FastSAM 可以使用点、框或文本作为提示,可靠地预测感兴趣的物体。除了加速 SegAny 任务外,FastSAM 还在 SegEvery 任务中表现出色,因为这可以与全实例分割高效地实现。然而,作为 SAM 的早期高效变体,FastSAM 仍然存在一些局限性,例如对较小物体生成低质量mask和生成边界不够平滑的mask。尽管存在这些缺点,FastSAM 通过在该领域引入基于 CNN 的架构标志着显著的进展。FastSAM 的架构如下图 4 所示。
基于 FastSAM 中 CNN 成功应用的基础,Varadarajan 等人引入了 SqueezeSAM,进一步用 U-Net 结构替换了 SAM 的基于Transformer的架构。U-Net 由用于特征提取的编码器和用于信息恢复的解码器组成。SqueezeSAM 保留了一般的 U-Net 架构,但在 U-Net 的最低尺度处加入了两个Transformer层,以在速度和准确性之间取得平衡。此外,SqueezeSAM 还具有几个微观级别的优化,例如将输出通道限制在 256,使用 BatchNorm代替 LayerNorm以提高效率,并在编码器和解码器之间引入跳跃连接。
SqueezeSAM 的一个独特挑战在于处理提示。与 SAM 在解码阶段使用提示tokens不同,SqueezeSAM 采用了早期融合策略,将编码的提示作为额外输入通道,添加到输入编码器之前。该模型使用 SA-1B 数据集从头开始训练,数据增强技术用于解决低质量数据问题。SqueezeSAM 主要设计用于摄影应用的部署,其中需要高效的交互式分割。如图 5 所示,其工作流程包括生成显著物体的初始mask,然后通过用户点击进行精细分割。
EfficientSAM并没有引入全新的网络,而是保留了 SAM 的原始架构,但替换了图像编码器。他们使用 ViT-tiny 或 ViT-small 作为轻量级编码器,并利用基于 SAM 的 Masked Image (SAMI) 预训练策略从头开始重新训练。SAMI 是从 Masked AutoEncoder (MAE) 框架适配而来,该框架最初用于预训练 SAM 的原始图像编码器。SAMI 遵循编码器-解码器管道:编码器从未遮罩的图块中生成潜在特征表示,而解码器则重构遮罩图块的缺失嵌入。这个过程由重构损失监督,比较 SAM 的 ViT-H 编码器生成的嵌入与 SAMI 管道生成的嵌入。预训练完成后,轻量级编码器从 SAMI 管道中提取并与 SAM 的其余组件集成,形成 EfficientSAM。最后一步是对整个模型在 SA-1B 数据集上进行微调,以进一步对齐和细化。SAMI 是一种通用的预训练方法,可以应用于训练任何用于 SAM 变体的主干。SAMI 和 EfficientSAM 的整体结构如下图 6 所示。
Xu 等人提出了 RAP-SAM,该模型旨在实现实时的通用分割,包括全景分割 (PS)、视频实例分割 (VIS) 和交互式分割(相当于 SegAny 任务)。RAP-SAM 保留了 SAM 的基本编码器-解码器架构,但结合了更高效的组件以增强性能。对于编码器,RAP-SAM 结合了特征金字塔网络 (FPN) 和可变形卷积 来从图像和视频中提取特征,同时使用提示编码器嵌入视觉提示。在解码器中,RAP-SAM 采用三阶段管道,利用新颖的基于池化的动态卷积来细化masktokens。每个阶段生成的tokens以及来自编码器的特征图作为输入。首先,这些输入通过动态卷积处理,然后利用多头自注意力 (MHSA) 和前馈网络 (FFN) 进行细化。在解码器之后,额外引入两个提示适配器,以增强视觉提示和分割tokens之间的交互。最终mask是通过计算更新后的tokens和更新后的提示之间的内积生成的。RAP-SAM 的架构如下图 7 所示。
最近,Ravi 等人引入了 Segment Anything Model 2 (SAM 2),这是对原始 SAM 的扩展。SAM 2 的目标是在图像和视频中提供高质量、实时的可提示分割。在图像分割任务中,SAM 2 报告显示其准确性更高,并且效率比原始 SAM 提高了 6 倍。这一显著进展主要归功于其高效的图像编码器 Hiera,这是一种分层 ViT,通过去除冗余组件并利用 MAE 框架进行训练,从 MViTv2 中简化而来。Hiera 是一种精简的、纯Transformer基础的架构,在图像和视频任务中运行速度更快,准确性更高。
3.1.2 基于知识蒸馏的方法
从前面图 3 所示的分类法中,我们观察到许多方法利用知识蒸馏,因为这种方法通常比完全模型训练需要更少的时间和资源。在本节中,我们回顾采用高效主干作为图像编码器,同时利用知识蒸馏进行训练的 SAM 变体。我们根据编码器类型将这些模型分为三类:具有 (i) 轻量级 ViT 编码器、(ii) 纯 CNN 编码器和 (iii) 修改过的注意力编码器的模型。我们将依次介绍每个类别。
(i) 轻量级 ViT 编码器Zhang 等人首次尝试用更高效的 TinyViT 替换 SAM 的重型 ViT 编码器,形成名为 MobileSAM 的集成模型。如 [55] 所强调,训练 SAM 从头开始需要多天和 128 个 GPU。MobileSAM 将这一复杂性归因于同时优化编码器和解码器的挑战。为了解决这个问题,他们提出了一种仅编码器蒸馏策略,如下图 8 所示,旨在将 ViT-H 的视觉表示能力转移到 TinyViT。所使用的损失函数是两个编码器的输出图像嵌入之间的简单均方误差 (MSE)。进一步微调提示编码器或mask解码器是可选的,并且可以提高准确性。
与 MobileSAM 类似,后来的 ESAM利用 EfficientFormerV2作为其主干,旨在在 CPU 环境中改善性能,特别是在资源有限的医疗设备上。鉴于专家模型在医疗应用中通常优于 SAM,ESAM 采用了一种新颖的知识蒸馏 (KD) 策略,称为整体知识蒸馏 (HKD),以将知识从专家模型转移到 ESAM。HKD 包括两个组成部分:特征图蒸馏和输出mask蒸馏。对于特征图蒸馏,结合了三种不同的聚焦方法 来指导学习过程。对于输出mask蒸馏,ESAM 使用教师mask和学生mask之间的均方误差 (MSE) 损失,辅以教师mask和真实mask之间的二元交叉熵 (BCE) 损失。为了进一步对齐专家模型和 ESAM 之间的特征图,提出了一种教师引导模块 (TGM),如下图 9 所示。
Shu 等人对 MobileSAM 进行了分析,发现仅编码器蒸馏可能导致显著的性能下降。为了解决这个问题,他们提出了更有效的全阶段知识蒸馏策略,称为 Hard Mining Full-Stage Knowledge Distillation,如下图 10 所示。
TinySAM 是一个新的 SAM 变体,其结构与 MobileSAM 相同,通过这种改进的 KD 策略进行训练。具体而言,该策略不仅监督图像嵌入,还监督输出tokens和输出mask,均使用 L1 损失。为了进一步增强蒸馏过程,他们引入了困难mask加权策略,该策略对更难预测的mask分配更大的权重,从而提高学习效率。因子 H 的计算如下:
(ii) 基于 CNN 的编码器
来自 NVIDIA 的研究人员基于 MobileSAM 引入了一种新的 SAM 变体,称为 NanoSAM。它旨在在 NVIDIA Jetson Orin 平台上实现实时性能,使用 NVIDIA TensorRT。NanoSAM 将基于 ViT 的编码器替换为纯卷积网络,具体为 ResNet18,同时保留了 MobileSAM 的其他组件。NanoSAM 从 MobileSAM 中进行蒸馏,两个模型都使用 TensorRT 进行重新训练以优化性能。MobileSAM 的图像编码器使用 FP32 精度进行优化,而 NanoSAM 的图像编码器则使用 FP16 精度以实现更快的执行。在 Jetson Orin Nano 和 Jetson AGX Orin 上的推理延迟结果表明,NanoSAM 比 MobileSAM 快 5 倍,且准确性损失最小。
Wang 等人 开发了一种高效的 SAM 变体 RepViT-SAM,使用他们新提出的基于 CNN 的主干 RepViT作为图像编码器。RepViT 背后的核心思想是将高效视觉Transformer (ViTs) 的有效设计原则集成到轻量级 CNN 中。这些设计原则应用于三个层次:块级、宏观和微观。在块级,RepViT 分离了tokens混合器和通道混合器,减少了扩展比,并增加了块的宽度。在宏观设计中,采用早期卷积作为输入干,深化下采样层 ,采用更简单的分类器,并调整各个阶段的块比例。在微观层面,仅使用 3x3 卷积,并且仅在奇数块中应用挤压与激励层。RepViT-SAM 使用知识蒸馏进行训练,遵循 [147] 中的相同流程,与 MobileSAM 相比,推理速度提高了 10 倍。
与 RepViT-SAM 的开发同时,Zhou 等人观察到,MobileSAM在边缘设备(如手机)上部署时仍难以实现实时性能。为了解决这个问题,他们引入了 EdgeSAM,该模型将基于 Transformer 的编码器替换为更轻量高效的纯 CNN 基础的 RepViT,旨在提高在资源有限设备上的性能。与 [162] 中的方法类似,Zhou 等人认为仅使用编码器蒸馏是不够的,因为它与任务无关,不能完全捕捉模型的任务特定需求。为了克服这一问题,他们提出了“循环中的提示蒸馏”方法,为输出mask增加额外的监督。“循环中的提示”是指一种动态采样策略,迭代地从教师和学生预测mask的非重叠区域中采样新提示。经过若干次迭代,累计损失被反向传播以更新编码器和解码器。为了进一步提高输出质量,EdgeSAM 提供了一个可选模块,从特定数据集中嵌入粒度先验。EdgeSAM 的蒸馏整体框架如下图 11 所示。
(iii) 修改过的注意力编码器
Zhang 等人引入了 EfficientViT-SAM,利用 EfficientViT作为图像编码器。EfficientViT 的主要优点是使用了 ReLU 线性注意力机制,该机制促进了全局信息交互,同时提高了硬件效率。通过消除不利于硬件的 softmax 操作并用 ReLU 替代,注意力计算重新公式化如下:
Shen 等人引入了 FastSAM3D,这是一个专门为 3D 体积医学图像设计的高效分割模型。该工作的关键贡献是开发了 3D 稀疏闪存注意力机制。这种新颖的注意力方法结合了 3D 膨胀注意力 扩展感受野)的优点和 FlashAttention以加速计算。FastSAM3D 使用修改过的 ViT-Tiny 作为图像编码器,从 ViT-Base 编码器蒸馏而来,确保在不妥协性能的情况下实现效率。作者实现了一种层级渐进蒸馏策略,以迭代地对齐两个编码器之间的特征图。
Yuan 等人的最新研究 RWKV-SAM,在加速 SAM 方面代表了一个重要进展,引入了流行的线性注意力模型作为高效主干。在他们的研究中,他们比较了基于 RWKV 和 Mamba 的架构,选择基于 RWKV 的方法构建 SAM 的轻量版本。该主干是混合设计,前两层由来自 [97] 的移动卷积块组成,最后一层使用视觉 RWKV 块构建。有关 RWKV 的更多细节见第 2.2.2 节。此外,SAM 类架构中还纳入了一个细化模块,通过融合每个阶段生成的不同层次的特征来增强mask质量。RWKV-SAM 的整体架构如下图 12 所示。该模型使用“蒸馏-微调”策略进行训练,首先将 SAM-H 中的知识蒸馏到主干,然后对整个模型进行微调。RWKV-SAM 显示出显著的效率提升,同时保持与 SAM 相当的分割性能。
3.1.3 基于量化的方法
3.1.4 基于剪枝的方法
Chen 等人首先开发了一种有效的剪枝策略,以减少 SAM 的规模和复杂度,得到的模型称为 SlimSAM。如第 2.3.3 节所述,剪枝算法旨在以结构化或个别的方式删除冗余参数。在应用于 SAM 的重型编码器时,初始步骤涉及估计权重和激活值的重要性,以确定哪些应被剪枝。评估重要性的核心思想是评估给定参数有无的损失差异。SlimSAM 引入了扰动泰勒重要性方法,该方法使用一阶泰勒展开来近似参数的重要性,并引入高斯噪声N以防止梯度变为零。这个过程被公式化如下:
一旦参数的重要性被估计,便采用一种称为交替剪枝的策略来执行结构化剪枝和后对齐。基于 ViT 的编码器首先被分为两个子结构:嵌入层和瓶颈层。该策略在修剪嵌入/瓶颈层以减小模型规模和对齐瓶颈/嵌入层以保持模型质量之间交替进行,确保效率与性能的平衡。该过程的工作流程如下图 13 所示。
3.1.5 代码重构
由 PyTorch 团队开发的 Segment Anything Fast 模型(SAMfast)是 SAM 的重写版本,利用了纯原生 PyTorch 的优化。报告显示,SAMfast 比原始实现快 8 倍,同时保持几乎相同的准确性。这一改进是通过系统地识别瓶颈并应用针对性的优化实现的。最初,团队发现了导致同步阻塞的长函数调用,从而重写了相应的代码。另一个显著的瓶颈是耗时的矩阵乘法,通过使用 bfloat16 精度来减轻这一问题。进行这些调整后,团队利用 torch.compile
融合更小的操作,并采用 PyTorch 的缩放点积注意力(SDPA)来加速 GPU 上的注意力计算。此外,通过集成使用 Triton 构建的新内核,进一步减少了 GPU 的内存使用。当 SAM 使用批量预测方法时,不同大小的输入张量被统一为 NestedTensors,从而显著提高了吞吐量。尽管进行了这些优化,矩阵乘法仍然是一个关键瓶颈。为了解决这个问题,团队实现了 int8 量化,并使用半结构稀疏性来近似矩阵乘法。有关逐步优化过程的更多细节,建议查看官方博客以获取更多信息。
3.2 加速 SegEvery 任务
如第 3.1 节所述,SegAny 任务的主要效率瓶颈在于重型图像编码器。任何具有更轻量架构的 SAM 变体在本质上都能比原始 SAM 更快地进行分割。然而,正如 Zhang 等人分析的那样,SegEvery 任务的主要挑战源于密集网格采样策略。该策略首先基于点网格预测大量mask,然后选择有效mask,这在计算上是昂贵的。因此,设计一种更高效的采样策略以减少预测mask的数量已成为加速 SegEvery 任务的核心方法。另一种潜在解决方案是将 SegEvery 任务转换为另一个成熟的任务,例如全实例分割,正如 FastSAM 所做的那样。在这一部分中,我们将回顾专门提出优化采样策略以加速 SegEvery 任务的相关工作。
基于 SAM 的结构,Zhang 等人 提出了一个面向对象的提示采样策略,以增强 SegEvery 任务的效率。这个项目名为 MobileSAMv2,独立于他们之前的工作,后者专注于加速 SegAny 任务。在 MobileSAMv2 中,研究人员使用训练在 SA-1B 小子集上的 YOLOv8 模型进行对象发现。该模型生成大量对应于潜在对象的边界框。通过非极大值抑制(NMS)过滤高度重叠的框,剩余的框作为框提示使用。通过使用这些过滤后的框作为提示,MobileSAMv2 消除了过滤预测mask的需要——这是一个耗时更长的过程。最大提示数设置为 320,新策略被报告为比传统的 32*32 网格采样策略快 16 倍。此外,MobileSAMv2 可以与 MobileSAM 集成,以创建一个统一模型,在 SegAny 和 SegEvery 任务中实现高效率。
Shu 等人观察到,使用密集点网格(例如,3232,6464)往往会生成大量冗余mask,这些mask在后处理过程中被过滤掉,这一操作会带来显著的时间成本。实际上,网格中只需少数点即可生成置信度高的mask。为了解决这一低效问题,他们提出了一种分层高效采样策略,逐步选择用于mask生成的最佳点。该策略涉及两轮提示生成。在第一轮中,使用稀疏网格,只包括每边默认点的四分之一(约 1/4)。基于这些点生成mask,经过过滤后,仅保留高置信度mask作为最终预测。在第二轮中,应用更密集的网格,遵循默认配置。然而,位于已被高置信度mask覆盖区域的点被排除,从而显著减少点的数量。两轮的预测结果随后被融合以生成最终输出。这一分层策略的流程如下图 14 所示。通过采用这种两轮方法,采样过程变得更加节省时间和细致,从而在 SegEvery 任务中显著加速,同时性能下降最小。
与前述所有工作不同,Fu et al.提出了一个专门为 SegEvery 任务设计的端到端训练管道,旨在开发一个能够更高效地进行全局分割的 SAM 变体。他们的模型名为 Lite-SAM,保留了原始 SAM 的整体架构,但用更轻量化的解决方案替换了重型图像编码器。Lite-SAM 的架构概述如下图 15 所示。
Lite-SAM 采用了一种称为 Lite-ViT 的 CNN-Transformer 混合结构,由 2、2、6 和 2 个 Lite-ViT 模块构成。Lite-ViT 的关键创新在于多尺度池化模块 (MSPM),该模块作为传统注意力机制的替代方案。MSPM 从 PoolFormer 模块 适配而来,利用通道层归一化并将池化操作扩展到多个尺度。正如前面所讨论的,SAM 的另一个主要瓶颈在于耗时的网格采样策略。为了解决这个问题,Lite-SAM 引入了一种自动提示提议网络(AutoPPN)以提高采样效率。AutoPPN 以编码器生成的特征图作为输入,直接预测点和框提示。为了确保高质量的提示,Lite-SAM 使用了比 CNN 更强大的基于 MSPM 的网络,并结合距离变换来估计点提示的置信度。虽然 Lite-SAM 主要设计用于加速 SegEvery 任务,但由于其轻量化的图像编码器,它在 SegAny 任务中也显示出更高的效率。
3.3 未来研究方向
通过对高效 SAM 变体的全面回顾,我们概述了加速 SAM 的当前进展。然而,仍然存在进一步探索和创新的机会。以下是几个潜在的未来研究方向,提供初步见解,希望能够激励读者为该领域的持续发展做出贡献。
3.3.1 探索先进的架构
尽管当前的 SAM 变体通过采用高效架构和模型压缩技术展示了效率提升,但仍有很大的改进潜力。探索 Transformer 替代模型,如 Mamba、RetNet、KAN和 TTT,为设计更轻量化和高效的结构提供了令人兴奋的机会。这些模型可能在计算效率方面具有优势,同时不牺牲分割准确性。除了替代模型外,进一步提升图像编码器和mask解码器中注意力机制的效率也是至关重要的。诸如线性注意力、低秩分解或结合卷积和基于注意力的设计的混合架构等方法应进一步研究。解决计算和内存使用中的瓶颈将增强 SAM 在不同硬件环境中的部署能力。
3.3.2 利用稀疏性和加速技术
在深度神经网络中观察到的固有稀疏性,即只有一部分参数对模型输出产生显著贡献,提供了提高 SAM 效率的有希望的途径。剪枝、量化和结构稀疏等技术可以进一步降低 SAM 的计算需求。尽管在 [18] 中的初步稀疏化努力已显示出成功,但未来的研究可以集中在理解 SAM 架构中稀疏性的分布和动态特性上。这包括研究可以在不影响性能的情况下被剪枝或稀疏化的 SAM 最优层或组件。此外,稀疏注意力机制、推理过程中的动态剪枝以及低精度训练等技术应被探索,以在大型部署中平衡准确性和效率。通过将这些与先进的知识蒸馏技术结合,可能实现更紧凑、高效的 SAM 变体。
3.3.3 硬件特定优化
对 SAM 进行针对特定硬件平台的优化,包括 GPU、TPU、专业 AI 加速器(例如 NVIDIA 的 TensorRT 或 Google 的 Edge TPU)以及边缘设备,可以显著提升性能和效率。硬件感知模型优化技术,如操作符融合、量化感知训练和自定义 CUDA 内核,可以在现代硬件平台上最大化吞吐量并降低延迟。在边缘设备的上下文中,由于面临存储、计算能力和能量供应的极端限制,这些优化对实时应用(如无人机或物联网设备上的分割)至关重要。未来的研究可以探索分层云-边缘架构,将计算密集型任务卸载到云端,同时在边缘设备上运行轻量级模型。此外,利用专用 AI 硬件(如现场可编程门阵列 (FPGA))或使用硬件感知的神经架构搜索 (NAS) 和混合精度量化等技术,可以进一步优化 SAM,以实现低延迟和资源受限环境中的有效运行,确保模型在不同硬件平台上高效运作。
3.3.4 视频和多模态数据的高效分割
视频和多模态任务处理复杂、动态的环境,正迅速在众多实际应用中获得相关性。尽管一些初步努力,例如 SAM 2用于视频分割和 MM-SAM用于多模态任务,已经扩展了 SAM 的适用性,但效率仍然是一个迫切问题。视频数据包含时间冗余,而多模态数据通常表现出模态之间的相关性。通过时间聚合和跨模态特征共享等技术利用这些固有的冗余,可能显著降低计算成本。未来的工作可以集中在利用时空注意力、为时间数据设计高效的内存机制和早期融合技术,以减少模态特定计算的数量,从而优化 SAM 的运行复杂性。开发能够动态适应不同帧或模态冗余水平的框架,可以进一步推动实际应用中的计算效率。
4. 评估
在本节中,我们系统比较了前述 SAM 变体的效率和准确性。参考这些工作的实验,我们选择大多数工作所执行的任务,并在其常用数据集上进行评估,并使用相应的度量标准。我们的评估在单个 24GB RTX 3090 GPU 和 14 个 vCPU 的 Intel(R) Xeon(R) Gold 6330 处理器 @ 2.00GHz 上进行。以下子节提供更多细节:第 4.1 节介绍了用于评估的数据集和度量标准;第 4.2 和 4.3 节分别报告了效率和准确性的定量结果。
4.1 数据集和度量标准
研究者们选择 COCO 2017和 LVIS v1作为评估数据集。COCO 是一个大规模数据集,旨在用于目标检测、分割和标注,包含 330K 张图像和 150 万个目标实例,覆盖 80 个目标类别。LVIS 针对大词汇实例分割进行了定制,特征超过 200 万个高质量的分割掩膜,涵盖 1200 多个类别,在 164K 张图像中。为了进行评估,我们使用两个数据集的验证集,COCO 包含 5000 张图像中的 36,781 个实例,LVIS 包含 19,809 张图像中的 244,707 个实例。为了评估效率,我们首先测试几个软指标,如参数数量 (#Params)、浮点运算量 (FLOPs)、乘加运算量 (MACs) 和内存使用情况。我们进一步计算有效错误率 (EER),这是一种更全面的评估,如 [86] 中所述。EER 定义为:
其中 N 是度量的数量, 和 分别表示被测试模型和参考模型的第 i 个度量。在我们的评估中,参考模型设定为 SAM-H。除了这些度量之外,还报告模型的运行时间和吞吐量。为了评估准确性,使用均值交并比(mIoU)来评估 SegAny 任务,并使用均值平均精度(AP)来评估实例分割任务。
4.2 效率比较
首先报告 SAM 及其变体的效率结果。以 SAM 官方示例中的图像 groceries.jpg 作为输入,利用一个边界框作为提示,使用工具 calflops 评估模型的参数数量、FLOPs 和 MACs。我们还计算了 ERR 以进行全面比较。结果如表 2 所示。在高效变体中,我们观察到 EdgeSAM 的参数数量、FLOPs、MACs 及其相应的 EER 是最低的,而 EfficientViT-SAM-XL1 的这些数值是最高的,其 EER 比 SAM-B 高出 3%。与最重的 SAM-H 相比,所有变体在模型大小和计算量上都表现出明显的降低。
还测量了模型在 SegAny 和 SegEvery 模式下的推理时间,使用来自 COCO 验证集的 100 张图像作为评估数据。在 SegAny 任务中,每张图像都提示 50 个固定的边界框。
通过曲线图(如下图16所示)报告每10个边界框的累计时间。基于此,计算处理一个图像并使用一个边界框提示所需的平均时间,并将其作为 SegAny 任务的推理时间报告。评估在 CPU 和 GPU 环境下进行,并同时记录 GPU 内存使用情况。此外,还测试了每个变体在 COCO 验证集上的吞吐量,使用真实的边界框。
结果总结在表3中。我们的研究发现,EfficientViT-SAM-L0 在 SegAny 任务中实现了最短的推理时间,与最重的模型 SAM-H 相比,在 GPU 上加速接近 30 倍,而在 CPU 上几乎加速 50 倍。EdgeSAM 也表现出令人印象深刻的性能,其 CPU 延迟为 259 毫秒,而 NanoSAM 在 GPU 上的延迟为 20 毫秒,均接近最佳结果。在 COCO 数据集的吞吐量测试中,NanoSAM 以每秒处理 27.9 张图像领先。其他两个变体,EfficientSAM-Ti 和 EfficientViT-SAM-L0,也展示了强大的吞吐量,各自超过每秒 20 张图像。
对于 SegEvery 任务,报告使用不同点网格大小(1616、3232、6464)或专用采样策略生成图像所有mask所需的平均时间。结果在表4中呈现。使用默认的 3232 网格,SAMfast-H 展示了最高的效率,延迟为 848 毫秒——比 SAM-H 快两倍以上。EfficientViT-SAM-L0 在 1616 和 6464 网格上表现最佳,分别实现了 258 毫秒和 3938 毫秒的延迟。有趣的是,我们观察到在使用较低网格密度时,EfficientSAM-S 比 SAM-H 更慢,16x16 网格的延迟为 1100 毫秒,32*32 网格的延迟为 2290 毫秒。采用替代采样策略的模型在效率上显示出显著改善,特别是 FastSAM 记录了 196 毫秒的延迟,而 MobileSAMv2 的延迟为 173 毫秒。
4.3. 精度比较
本小节报告了 SAM 及其变体在 SegAny 任务(使用点/框提示)和实例分割任务上的精度结果。我循了 [159] 中的评估框架,并通过引入其他评估模块进行了统一的变体评估。
为了评估 SegAny 任务,采用两种类型的点作为提示:
- 1)真实边界框的中心点
- 2)从真实mask中均匀随机采样的点,遵循 [131] 中的设置。
在 COCO 和 LVIS 上对变体进行了评估,并在下表 5 中报告了平均交并比(mIoU)。当使用中心点提示时,SAM2-B+ 和 EfficientViT-SAM-XL1 在 COCO 上达到了 54.3% 的最高 mIoU,超过了 SAM-H 的 53.6% mIoU,而 SAMfast-H 也以 53.6% 的 mIoU 在 LVIS 上展示了变体中的最佳性能。在随机点提示的设置下,EfficientViT-SAM-XL1 的表现优于 SAM-H,特别是在使用 3 个点提示时,分别增加了 2.7% 和 0.7%。从数据集的角度来看,我们观察到 LVIS 的结果普遍低于 COCO,尤其是对于 FastSAM 和 EfficientSAM-Ti,其在 LVIS 上的准确率降至 30% 以下。
此外,还评估了使用两种类型框提示的 SegAny 任务的精度:
- 1)真实的边界框
- 2)对应真实mask的最紧边界框,灵感来自于 [131, 159] 中的实验。
在 COCO 和 LVIS 上报告了 mIoU 的结果,如表 6 所示。我们观察到 EfficientViT-SAM-XL1 在每种设置中都展示了最高的准确性,分别超过 SAM-H 1.5%、1.1%、1.9% 和 0.6%。SAMfast-H 和 EfficientViT-SAM-L0 在框提示的分割任务中也展现了接近 SAM-H 的性能。
对于实例分割任务,采用了 ViTDet、YOLOv8、GrounddingDINO、Detic 和 H-Deformable-DETR 与 Swin-L 作为物体检测器,帮助生成潜在物体的边界框,参考文献 [117, 159, 164]。评估了所有物体的平均精度(AP),以及小型、中型和大型物体的 AP。结果报告在下面表 7、8 和 9 中。与之前的结果类似,我们发现 EfficientViT-SAM-XL1 在 COCO 数据集上总是以最高的 AP 表现(除 H-Deformable-DETR 外)。在将 ViTDet 作为检测器并在 LVIS 数据集上测试的设置下,SAMfast-H 以 44.5% 的 AP 超过了所有其他变体。
根据第 4.2 节和第 4.3 节的结果,进一步绘制了吞吐量-mIoU 散点图,以观察变体的效率-精度权衡。具体而言,选择了在 COCO 数据集上使用真实边界框作为提示进行评估的吞吐量和 mIoU。结果如下图 17 所示。
5. 结论
在本次调查中,主要讨论和评估了关注于高效分割任意物体和分割所有物体的显著工作,这些方法旨在减少资源消耗和降低延迟。对于高效的 SegAny 任务,大多数工作采用替换图像编码器或整个架构为轻量级替代方案的方法,随后进行从头训练或通过知识蒸馏的方式进行训练。其他工作则旨在利用量化、剪枝或局部优化等技术压缩原始模型。对于高效的 SegEvery 任务,采用有效且高效的采样策略生成提示是至关重要的。
在详细回顾这些方法后,还概述了四个可能的未来研究方向,这些方向可能推动该领域的新趋势。此外,我们在一致的环境中评估了这些模型的效率、准确性及其相应的权衡,提供了公正而有价值的比较。我们的分析表明,一些变体在特定场景中已经超过了原始的 SAM,我们相信它们的成功将激励该领域的进一步探索和创新。
本文转自AI生成未来 ,作者:AI生成未来