MobileSAM来啦 | 比SAM小60倍，比FastSAM快4倍，速度和效果双赢

集智书童公众号

发布于 2023-09-04 10:43:33

1.1K0

发布于 2023-09-04 10:43:33

文章被收录于专栏：集智书童

Segment anything model （SAM）是一种快速引导的视觉基础模型，用于从其背景中剪切出感兴趣的目标。自Meta研究团队发布SA项目以来，SAM因其令人印象深刻的零样本传输性能和与其他模型兼容的高度通用性而备受关注，用于高级视觉应用，如具有细粒度控制的图像编辑。许多这样的用例需要在资源受限的边缘设备上运行，比如移动端应用程序。在这项工作中，作者的目标是通过用轻量化图像编码器取代复杂的图像编码器，使SAM对移动端友好。原始SAM文件中训练这种新SAM方式会导致性能不令人满意，尤其是当可用的训练来源有限时。作者发现，这主要是由图像编码器和Mask解码器的耦合优化引起的，因此作者提出了解耦蒸馏。具体地说，作者将原始SAM中的图像编码器ViT-H的知识蒸馏到一个轻量化的图像编码器中，该编码器可以自动与原始SAM中的Mask解码器兼容。训练可以在不到一天的时间内在单个GPU上完成，由此产生的轻量化SAM被称为MobileSAM，它比原始SAM小60多倍，但性能与原始SAM相当。就推理速度而言，MobileSAM每幅图像运行约10ms：图像编码器运行8ms，Mask解码器运行2ms。凭借卓越的性能和更高的通用性，作者的MobileSAM比并发的FastSAM小7倍，快4倍，更适合移动端应用程序。代码：https://github.com/ChaoningZhang/MobileSAM

1、简介

ChatGPT Zhang等人彻底改变了NLP领域，标志着生成人工智能（AIGC，又称人工智能生成内容）的突破。使这成为可能的是Brown等人、Radford等人的GPT系列模型，这些模型是Bommasani等人在网络规模的文本数据集上训练的基础模型。

继NLP中基础模型的成功之后，何、乔、张等人通过对比学习将图像编码器与文本编码器结合起来。最近，Meta Research团队发布了Kirillov等人的“Segment Anything”项目，其中提出了一个名为SAM的即时引导视觉基础，被认为是视觉的GPT时刻。

SAM由两个组件组成：

ViT-based图像编码器
Prompt-guided Mask解码器

它们按顺序工作（见图1）。

自SAM问世以来，由于多种原因引起了人们的极大关注。首先，它首次表明，愿景可以遵循NLP，走一条将基础模型与即时工程相结合的道路。其次，它是第一个执行无标签分割的方法，这是一项与标签预测并行的基本视觉任务。

此外，这项基本任务使SAM与其他模型兼容，以实现高级视觉应用，如文本引导的分割和具有细粒度控制的图像编辑。然而，许多这样的用例需要在资源受限的边缘设备上运行，比如移动端应用程序。如官方演示中所示，通过处理图像嵌入，SAM可以在资源受限的设备上工作，因为Mask解码器是轻量化的。使SAM流水线计算繁重的原因在于巨大的图像编码器。

在这项工作中，作者研究了如何获得适用于资源受限移动端设备的轻量化SAM，因此称为MobileSAM。

鉴于SAM中的默认图像编码器是基于ViT-H的，获得MobileSAM的一个简单方法是遵循Kirillov等人的管道，用较小的图像编码器重新训练新的SAM，比如用较小的ViT-L甚至更小的ViT-B替换ViT-H。

表3总结了图像编码器不同变体的SAM参数。如Kirillov等人所述，使用ViT-L或ViT-B作为图像编码器训练新的SAM需要128个GPU，持续多日。这种资源密集型的再训练可能是复制或改进其结果的一个不小的负担。这种优化难度主要来自于图像编码器和Mask解码器的耦合优化。

基于这种理解，作者提出将图像编码器和Mask解码器的优化解耦。具体来说，作者首先将知识从默认图像编码器ViT-H蒸馏到一个微小的ViT。之后，作者可以微调原始SAM中的Mask解码器，以更好地与蒸馏的图像编码器对齐。值得强调的是，对准优化是可选的，因为轻量化图像编码器是从默认图像编码器中蒸馏的，这保证了其与默认Mask解码器的固有对准。

通过将寻求新SAM管道的问题转化为解耦蒸馏，作者的方法具有简单有效的优点，同时可以以低成本（在不到一天的单个GPU上）进行复制。由此产生的MobileSAM将编码器参数减少了100倍，将总参数减少了60倍。令人惊讶的是，这样一个轻量化的MobileSAM的性能与最初的原始SAM不相上下，这为推动移动端应用程序的SAM迈出了重要的一步。

对于MobileSAM的推断，单个图像仅运行约10ms：在图像编码器上运行8ms，在Mask解码器上运行2ms。值得强调的是，作者的MobileSAM比同期的FastSAM赵等人小7倍，快4倍，同时实现了卓越的性能。

2、相关工作

2.1、SAM：泛化和通用功能性

自今年4月初出现以来，已经出现了许多从不同角度研究SAM的项目和论文。鉴于SAM声称可以分割任何东西，一系列工作报告了其在现实世界中的表现，包括医学图像，伪装物体和透明物体等。研究结果一致表明，SAM在一般设置中工作良好，但在上述具有挑战性的设置中则不然。

另一个重要的研究方向集中在增强SAM以提高其实用性。攻击SAM Zhang等人已经表明，SAM的输出Mask可以通过恶意生成的对抗性扰动被对抗性攻击轻易操纵。另一项工作Qiao等人进一步对SAM进行了全面的稳健性评估，从风格转移和常见腐蚀到局部遮挡和对抗性扰动。在Qiao等人中发现。SAM具有高鲁棒性，但不适用于对抗性扰动，这与Zhang等人的发现非常一致。

另一项工作侧重于展示SAM的多功能性。Grounded SAM IDEA Research是将Grounding DINO与SAM相结合的开创性工作，用于通过文本输入分割任何内容。具体来说，它依赖于Grounding DINO从文本生成边界框，然后生成的框可以用作分割Mask的prompt。

SAM预测没有标签的Mask和多项工作将SAM与CLIP-等其他模型相结合。除了目标分割，多项工作在其他领域也显示出了其多功能性，包括图像编辑、修复任务以及视频中的目标跟踪。除了2D视觉，SAM的研究还扩展到了3D目标重建，展示了其辅助从单个图像生成3D模型的能力。

2.2、ViT：轻巧高效

早期的移动视觉应用主要由轻量化神经网络提供支持，如MobileNet及其改进的变体。MobileNet的核心思想在于将普通卷积块分为深度卷积和点卷积，这大大减少了模式参数和计算时间。自ViT问世以来，许多工作都试图使其轻量化和高效。与原始ViT论文中的ViT-Huge（ViT-H）、ViT-Large（ViT-L）和ViT-Base（ViT-B）互补，Touvron等人引入了较小的ViT，并表示为Deit Small（Deit-S）和Deit Tiny（Deit-T）ViT-Small和ViT-Tiny。

MobileViT-Metha是将ViT与标准卷积相结合以提高其性能的开创性工作，其性能优于MobileNet v2 。主要动机是利用CNN的局部表示能力，随后进行了多项旨在提高模型速度的后续工作，包括EfficientFormer、EfficientViT 、NextViT和Tiny ViT。最近在轻量化和更快的ViT方面取得的进展是对提出的解耦蒸馏的补充，以使下一代SAM适用于资源受限的移动设备。

3、本文方法

3.1、背景和项目目标

1、Background on SAM

在这里，作者首先总结SAM的结构及其工作原理。SAM由一个基于ViT的图像编码器和一个 prompt-guided Mask解码器组成。图像编码器将图像作为输入并生成嵌入，然后将嵌入提供给Mask解码器。Mask解码器生成一个Mask，根据点（或框）等prompt从背景中剪切出任何目标。

此外，SAM允许为同一prompt生成多个Mask，以解决模糊性问题，这提供了宝贵的灵活性。考虑到这一点，这项工作保持了SAM的流水线，首先采用基于ViT的编码器来生成图像嵌入，然后采用 prompt-guided 解码器来生成所需的Mask。该管道是为“Segment Anything”而优化设计的，可用于“Segment 所有东西”的下游任务。

2、Project goal

该项目的目标是生成一个移动端友好型SAM（MobileSAM），以轻量化的方式实现令人满意的性能，并且比原始SAM快得多。原始SAM中的 prompt-guided Mask解码器的参数小于4M，因此被认为是轻量化的。

给定编码器处理的图像嵌入，如他们的公开中所示，SAM可以在资源受限的设备中工作，因为Mask解码器是轻量化的。然而，原始SAM中的默认图像编码器是基于ViT-H的，具有超过600M的参数，这是非常复杂的的，并使整个SAM管道与移动端设备不兼容。

因此，获得移动端友好SAM的关键在于用轻量化的图像编码器取代复杂的图像编码器，这也自动保持了原始SAM的所有功能和特性。

3.2、Proposed Method

1、耦合蒸馏

实现作者项目目标的一个简单方法是SAM中的官方管道，用较小的图像编码器重新训练新的SAM。如SAM所述，使用ViT-H图像编码器训练SAM需要在256个A100 GPU上花费68小时。用ViT-L或ViT-B取代ViT-H将所需的GPU减少到128个，这对于社区中的许多研究人员来说仍然是复制或改进他们的结果的一个不小的负担。

按照他们的方法，作者可以进一步采用更小的图像编码器，并用他们提供的11-T分割数据集重新训练新的SAM。请注意，所提供的数据集中的Mask是由预训练的SAM（使用ViT图像编码器）给出的。本质上，这种再训练过程是知识蒸馏，它将知识从基于ViT-H的SAM转移到具有较小图像编码器的SAM（见图2左）。

2、从半耦合到解耦的蒸馏

当执行从原始SAM到较小图像编码器的KD时，困难主要在于图像编码器和组合解码器的耦合优化。直观地说，图像编码器的优化取决于图像解码器的质量，反之亦然。当SAM中的两个模块都处于不良状态时，将它们训练到良好状态更具挑战性。受分治算法的启发，作者提出将KD任务分为两个子任务：

图像编码器的蒸馏
Mask解码器的微调

具体来说，作者首先通过将知识从ViT-H转移到较小的编码器来对图像编码器执行KD。由于原始SAM中的Mask解码器已经是轻量化的，作者计划保留其架构。这带来了易于使用的组合解码器的好处，用于微调而不是从头开始训练它。

为了缓解耦合蒸馏的优化问题，一种简单的方法是使用复制和冻结的Mask解码器优化图像编码器（见图2右侧）。冻结操作可以帮助防止Mask解码器的质量因较差的图像编码器而劣化。作者称这种蒸馏为半耦合，因为图像编码器的优化仍然没有与Mask解码器完全解耦。

根据经验，作者发现这种优化仍然具有挑战性，因为prompt的选择是随机的，这使得Mask解码器可变，从而增加了优化难度。因此，作者建议直接从原始SAM中的ViT-H中蒸馏小图像编码器，而不使用组合解码器，这被称为解耦蒸馏（见图3）。

对图像嵌入进行蒸馏的另一个优点是，作者可以采用简单的MSE损失，而不是使用Focal Loss的组合。Dice Loss用于进行Mask预测。

3、论Mask解码器微调的必要性

与半耦合蒸馏不同，上述解耦蒸馏产生了一个轻量化的图像编码器，该编码器可能与原始的冻结Mask解码器不太一致。

根据经验，作者发现这不是真的，因为从学生图像编码器生成的图像编码可以足够接近原始教师编码器的图像编码，这使得第二阶段对组合解码器的微调是可选的。预计在冻结的轻量化图像编码器上对Mask解码器进行微调或将它们联合微调可能会进一步提高性能。

4、初步评估

在此，作者进行了初步研究，以比较耦合蒸馏和解耦蒸馏。在这里，为了进行绩效评估，作者计算教师SAM和学生SAM在同一prompt点上生成的两个Mask之间的mIoU。

直观地，通过假设由ViT-H生成的Mask是GT，较高的mIoU指示较高的Mask预测性能。对于耦合蒸馏，作者采用原始SAM中提供的带有ViT-B的SAM。它在128个GPU（每个GPU 1个样本）上的SA-1B（11M图像）上进行了180k次迭代训练。

相反，在作者的解耦蒸馏设置中，作者在2个GPU上训练模型（每个GPU两个样本，以节省计算资源），在SA-1B数据集（11k）图像的0.1%样本上进行55k次迭代。

总的来说，解耦蒸馏比耦合蒸馏占用的计算资源不到1%，同时实现了耦合的mIoU 0.75比0.72的优越性能（在200个样本上平均）。由于ViT-B对移动端设备来说仍然是一个不平凡的负担，因此在下文中，作者基于作者提出的解耦蒸馏对TinyViT（具有5M参数）进行了实验。

4、实验

4.1、Lightweight Image Encoder

作者项目的目标是通过用移动端设备的轻量化图像编码器取代默认的ViT-H来获得高效的SAM。作为一个基于ViT的主干，ViT Tiny具有与Deit Tiny相似的参数，但性能更好。例如，在ImageNet-1K上，Deit Yiny实现了72.2%的准确率，而ViT Tiny实现了79.1%的准确率。

因此，作者采用ViT Tini作为概念验证，以证明作者提出的解耦蒸馏训练轻量化MobileSAM的有效性，该轻量MobileSAM可以比原始SAM快得多。所采用的轻量化图像编码器由四个阶段组成，这些阶段逐渐降低分辨率。第一级由具有反向残差的卷积块构建。而其余三级由Transformer块组成。

在模型开始时，有2个步长为2的卷积块用于对分辨率进行下采样。不同阶段之间的下采样操作由步长为2的卷积块来处理。作者将最后一次下采样卷积中的步长2设置为1，以使最终分辨率与原始SAM的ViT-H图像编码器的分辨率相匹配。

请注意，第2节中讨论的其他高效图像编码器也可以用作图像编码器。

4.2、Training and evaluation details

对于图像编码器上的解耦KD，作者使用SA-1B数据集Kirillov等人的1%来训练轻量化编码器。[2023]在单个GPU（RTX3090）上进行8个Epoch的训练。作者观察到，考虑到教师图像编码器比作者采用的学生图像编码器要重得多，因此在前向处理上花费了更多的计算（见上文）。

为了更快地进行蒸馏，作者遵循Wu等人的实践。预先保存图像嵌入，这样作者只需要运行一次正向过程。使用一个GPU，作者可以在不到一天的时间内获得作者的MobileSAM。用更多的GPU在更长的时间内训练作者的MobileSAM有望产生更好的性能。

执行Mask解码器微调的初步研究进一步提高了MobileSAM的性能，然而，为了简单起见，作者在本文的这个版本中省略了这一步骤。为了定量评估蒸馏的SAM，作者计算原始SAM预测的Mask和作者的MobileSAM之间的mIoU。

4.3、MobileSAM的性能与最初的SAM相当

对于主要结果，作者报告了具有两种类型prompt的预测Mask：点和框。作者不使用文本prompt报告结果，因为SAM的官方github项目没有为文本引导的Mask解码器提供预训练的模型。

以point为prompt的结果如图4所示，以box为prompt的效果如图5所示。作者观察到MobileSAM做出了与原始SAM类似的令人满意的Mask预测。

4.4、消融实验

在这里，作者对训练计算对SAM性能的影响进行了消融研究。表7中的结果表明，在相同的迭代次数下，增加Batch-size会提高模型性能。此外，在Batch-size的相同的情况下，通过增加训练Epoch，性能也受益于更多的更新迭代。

请注意，所有实验都是在单个GPU上进行的。作者预计，增加GPU的数量以允许更大的Batch-size或进一步增加迭代可以进一步提高性能。

4.5、MobileSAM在所有方面都优于FastSAM

请注意，SAM原始论文的标题是“Segment anything”，而不是“segment everything”。如SAM中所强调的，SAM执行可prompt分割的任务，该任务“在给定任何分割prompt的情况下返回有效的分割Mask”。

prompt的作用是指定要在图像中分割的内容。理论上，只要正确设置prompt，任何目标都可以被分割，因此，它被称为“Segment anything”。相比之下，“segment everything”本质上是目标建议生成，对此不需要prompt。在SAM中，选择“segment everything”（目标建议生成）作为下游任务之一，以演示其零样本传输性能。

总之，“Segment anything”解决了任何目标的可prompt分割的基础任务，而“segment everything”解决了为所有目标生成Mask建议的下游任务不一定需要prompt，FastSAM以无prompt的方式直接用YOLO v8生成Mask建议。为了实现可prompt分割，设计了一种映射算法来从提议Mask集中选择Mask。

值得强调的是，评估其泛化/鲁棒性或研究其多功能性的后续工作主要集中在任何而不是所有模式上，因为前者涉及基础任务。因此，与FastSAM的比较主要集中在“Segment anything”上，但为了完整性，作者也提供了关于“segment everything”的比较。

4.6、MobileSAM的速度更快、规模更小

FastSAM由基于YOLOv8的检测分支和基于YOLACT的分割分支组成，以执行无prompt的Mask建议生成。它有68M个参数，处理一张图像需要40ms。相比之下，MobileSAM的10M参数更少，这明显更小。就推理速度而言，在单个GPU上，处理图像需要40ms，而MobileSAM的仅需要10ms，这比FastSAM快4倍。

4.7、Segment anything模式下的mIoU比较

作者进一步比较了预测Mask与原始SAMMask之间的mIoU。请注意，FastSAM不能以单个点作为原始SAM来预测Mask。相反，它至少需要两个prompt点：一个用于前景，另一个用于背景。

表6中的结果显示，FastSAM的mIoU比MobileSAM的mIoU小得多，这表明FastSAM的Mask预测与原始SAM的Mask不同。此外，当两个prompt点之间的距离增加时，FastSAM的mIoU下降得非常快。这主要是由于当前台prompt点设置得离后台prompt点太近时，FastSAM经常无法预测目标。

4.8、Results for segment everything

“Segment everything”的结果如图6所示。

为了完整性，作者还报告了原始SAM的结果，它生成了一个令人满意的目标提案。作者有两个主要的观察结果。首先，作者的MobileSAM的结果与原始SAM的结果惊人地一致。

相比之下，FastSAM的结果往往不太令人满意。例如，FastSAM经常无法预测某些目标，例如第一张图像中的屋顶。此外，Mask方案有时很难解释（参见第一幅图像中舞台的Mask和第二幅图像中天空的Mask）。

其次，FastSAM通常生成具有非平滑边界的Mask，对此，作者建议读者放大以查看图6中的详细信息。例如，第三个图像中的支柱具有非平滑边界，而原始SAM和作者的MobileSAM没有这个问题。

5、总结

在这项工作中，作者的目标是通过用轻量化图像编码器取代复杂的图像编码器，使SAM对移动端友好。作者发现，原始SAM文件中训练这种新SAM的方式会导致不令人满意的性能，尤其是在训练来源有限的情况下。图像编码器和Mask解码器的耦合优化是原因，因此作者提出了解耦蒸馏，其中知识从原始SAM中的图像编码器ViT-H蒸馏到轻量化图像编码器。

作者展示了由此产生的轻量化图像编码器可以自动与原始SAM中的Mask解码器兼容。作者的MobileSAM比原始SAM小60多倍，但性能与原始SAM相当。

此外，作者与并发的FastSAM进行了比较，表明MobileSAM实现了卓越的性能。作者的MobileSAM也比并发的FastSAM快4倍，小7倍，使其更适合移动端应用程序。由于作者的MobileSAM保留了原始SAM的所有管道，只是取代了图像编码器，因此现有基于SAM的项目可以即插即用，以零成本从复杂的SAM转变为轻量化SAM。