【读点论文】FASTER SEGMENT ANYTHING: TOWARDS LIGHTWEIGHT SAM FOR MOBILE APPLICATIONS。先换为小图像编码器，一种解耦的知识蒸馏方法

FASTER SEGMENT ANYTHING: TOWARDS LIGHTWEIGHT SAM FOR MOBILE APPLICATIONS

ABSTRACT

• Segment Anything Model（SAM）因其令人印象深刻的 zero-shot transfer 性能和对许多视觉应用（如具有细粒度控制的图像编辑）的高度通用性而吸引了大量关注。许多这样的应用程序需要在资源受限的边缘设备上运行，比如手机。在这项工作中，我们的目标是通过用轻量级图像编码器取代重量级图像编码器，使SAM对移动友好。原始SAM文件中训练这种新SAM的天真方式会导致不令人满意的性能，尤其是当可用的训练来源有限时。我们发现，这主要是由图像编码器和掩模解码器的耦合优化引起的，因此我们提出了解耦蒸馏。
• 具体地说，我们将知识从重图像编码器（原始SAM中的ViT-H）提取到轻量级图像编码器，该编码器可以自动与原始SAM的掩码解码器兼容。训练可以在不到一天的时间内在单个GPU上完成，由此产生的轻量级SAM被称为MobileSAM，它比原始SAM小60多倍，但性能与原始SAM相当。就推理速度而言，使用单个GPU，MobileSAM每个图像运行约10ms：图像编码器运行8ms，掩码解码器运行4ms。凭借卓越的性能，我们的MobileSAM比并发的FastSAM快约5倍，小7倍，更适合移动应用程序。此外，我们还证明了MobileSAM在CPU上可以相对平稳地运行。我们项目的代码在MobileSAM上提供，演示显示MobileSAM可以在CPU上相对平稳地运行。https://github.com/ChaoningZhang/MobileSAM
• 论文地址：[2306.14289] Faster Segment Anything: Towards Lightweight SAM for Mobile Applications (arxiv.org)
• 主要研究利用知识蒸馏的技术，将sam的大模型迁移到一个小模型上，可以应用到移动设备上，模型大小相比原生sam小了60倍。考虑到sam的image encoder参数太多，不适合移动设备，论文的目的也就是用一个轻量级的image encoder替换原先比较重的image encoder模块。

Introduction

• ChatGPT彻底改变了NLP领域，标志着生成人工智能（AIGC，又称人工智能生成内容）的突破。使这成为可能的是GPT系列模型，这些模型是[On the opportunities and risks of foundation models]在网络规模的文本数据集上训练的基础模型。继NLP中的基础模型取得成功后，多项工作通过对比学习将图像编码器与文本编码器结合起来。最近，Meta Research团队发布了“Segment Anything”项目，其中提出了一个名为SAM的即时引导视觉基础，被认为是视觉的GPT时刻。SAM由两个组件组成：基于ViT的图像编码器和提示引导的掩码解码器，它们按顺序工作（见下图）。

  • 

  • The overview of Segment Anything Model.

• 自SAM问世以来，由于多种原因引起了人们的极大关注。首先，它首次证明了视觉可以遵循NLP来追求一条将基础模型与即时工程相结合的道路。其次，它首次执行了无标签分割，这是一项与标签预测并行的基本视觉任务。而且这个基本任务使得SAM与其他模型兼容以实现高级视觉应用，如文本引导的分割和具有细粒度控制的图像编辑。然而，许多此类用例需要在资源受限的边缘设备上运行，如移动应用程序。如官方演示所示，通过处理图像嵌入，SAM可以在资源受限的设备上工作，因为掩码解码器是轻量级的。使SAM流水线计算繁重的原因在于巨大的图像编码器。在这项工作中，我们研究了如何获得适用于资源受限移动设备的轻量级SAM，因此称为MobileSAM。

• 鉴于SAM中的默认图像编码器基于ViT-H，获得MobileSAM的一个简单方法是遵循[Segment anything]中的官方管道，用较小的图像编码器重新训练新的SAM，比如用较小的ViT-L甚至更小的ViT-B替换ViT-H。在下表中总结了具有图像编码器的不同变体的SAM的参数。

  • 

  • Parameters SAM with different image encoders.

• 使用ViT-L或ViT-B作为图像编码器训练新的SAM需要128个GPU，持续多日。这种资源密集型的再训练可能是复制或提高其结果的一个不小的负担。这种优化难度主要来自于图像编码器和掩模解码器的耦合优化。基于这一理解，我们提出将图像编码器和掩模解码器的优化解耦。具体来说，我们首先将知识从默认图像编码器ViT-H提取到一个微小的ViT。之后，我们可以微调原始SAM中的掩码解码器，以更好地与提取的图像编码器对齐。值得强调的是，对齐优化是可选的，因为轻量级图像编码器是从默认图像编码器中提取的，这保证了其与默认掩码解码器的固有对齐。

• 通过将寻求新SAM管道的问题转化为解耦蒸馏，我们的方法具有简单有效的优点，同时可以以低成本（在不到一天的单个GPU上）进行复制。由此产生的MobileSAM将编码器参数减少了100倍，将总参数减少了60倍。令人惊讶的是，这样一个轻量级的MobileSAM的性能与最初的重量级SAM不相上下，这是推动移动应用程序SAM的重要一步。对于MobileSAM的推断，单个图像仅运行约10ms：在图像编码器上运行8ms，在掩码解码器上运行4ms。值得强调的是，我们的MobileSAM比同期的FastSAM快约5倍，小约7倍，同时实现了卓越的性能。

Related work

• SAM：通用性和多功能性。自今年4月初出现以来，已经出现了许多从不同角度研究SAM的项目和论文。鉴于SAM声称可以分割任何东西，一系列工作已经报道了它在现实世界中的表现，包括医学图像、camouflaged 物体和 transparent 物体。研究结果一致表明，SAM在一般设置中效果良好，但在上述具有挑战性的设置中效果不佳。另一个重要的研究方向集中在增强SAM以提高其实用性上。Attack SAM 已经表明，SAM的输出掩码可以通过恶意生成的对抗性扰动被对抗性攻击轻易操纵。另一项工作进一步对SAM进行了全面的稳健性评估，从风格转移和常见破坏到局部遮挡和对抗性扰动。在研究中发现，SAM具有很高的鲁棒性，但不适用于对抗性扰动，这与Zhang等人中的发现非常一致。
• 另一个工作重点是展示SAM的多功能性。Grounded SAM IDEA Research是将Grounding DINO与SAM相结合的开创性工作，用于通过文本输入分割任何内容。具体来说，它依赖于Grounding DINO从文本生成边界框，然后生成的框可以用作分割遮罩的提示。SAM预测没有标签的掩码，多项工作将SAM与CLIP等其他模型相结合，对任何事物进行语义分割。除了对象分割，多项工作在其他领域也显示出了其多功能性，包括图像编辑，以及修复任务，视频中的对象跟踪。除了2D视觉，SAM的研究还扩展到了3D对象重建，展示了其辅助从单个图像生成3D模型的能力。关于SAM的完整调查，建议读者参考Zhang等人[A survey on segment anything model (sam): Vision foundation model meets prompt engineering]。
• ViT：重量轻，效率高。早期的移动视觉应用主要由轻量级的细胞神经网络提供支持，如MobileNet及其改进的变体。MobileNet的核心思想在于将普通卷积块分为深度卷积和点卷积，这大大减少了模式参数和计算时间。自从ViT，许多工作都试图使其轻量级和高效。与原始ViT论文中的ViT-Huge（ViT-H）、ViT-Large（ViT-L）和ViT-Base（ViT-B）互补，Touvron等人【Training data-efficient image transformers & distillation through attention】中引入了较小的ViT，并表示为Deit Small（Deit-S）和Deit Tiny（Deit-T）ViT-Small和ViT-Tiny。MobileViT 是将ViT与标准卷积相结合以提高其性能的开创性工作，其性能优于MobileNet v2 。主要动机是利用CNN的局部表示能力，在这一实践之后，进行了多项旨在提高模型速度的后续工作，包括EfficientFormer 、EfficientViT 、Next-ViT 和Tiny-ViT。最近在轻量级和更快的ViT方面取得的进展是对我们提出的去耦蒸馏的补充，旨在使下一代SAM适用于资源受限的移动设备。

Mobile-Friendly SAM

Background and Project Goal

• SAM的背景。在这里，我们首先总结SAM的结构及其工作原理。SAM由一个基于ViT的图像编码器和一个提示引导的掩码解码器组成。图像编码器将图像作为输入并生成嵌入，然后将嵌入提供给掩码解码器。mask 解码器生成一个蒙版，根据点（或框）等提示从背景中剪切出任何对象。此外，SAM允许为同一提示生成多个掩码，以解决歧义问题，这提供了宝贵的灵活性。考虑到这一点，这项工作保持了SAM的流水线，首先采用基于ViT的编码器来生成图像嵌入，然后采用提示引导解码器来生成所需的掩码。该管道是为“segment anything”而优化设计的，可用于“segment anything”的下游任务。
• 项目目标。该项目的目标是生成一个移动友好型SAM（MobileSAM），以轻量级的方式实现令人满意的性能，并且比原始SAM快得多。原始SAM中的提示引导掩码解码器的参数小于4M，因此被认为是轻量级的。给定编码器处理的图像嵌入，如他们的公开演示中所示，SAM可以在资源受限的设备中工作，因为掩码解码器是轻量级的。然而，原始SAM中的默认图像编码器是基于ViT-H的，具有超过600M的参数，这是非常重的，并且使得整个SAM管道与移动设备不兼容。因此，获得移动友好SAM的关键在于用轻量级的图像编码器取代重量级的图像编码器，这也自动保持了原始SAM的所有功能和特性。在下文中，我们详细介绍了我们提出的实现该项目目标的方法。

Proposed Method

• 耦合蒸馏方法，困难主要在于image encoder 和prompt-guided mask docoder的耦合优化。因为他们互相影响，当这两个模块都比较差时，想要将它们训练到好非常困难。作者将其两个子任务：蒸馏image encoder和微调mask docoder，对于mask docoder参数，采用复制和冻结的方法，可以避免mask decoder受image encoder的影响，作者称这种蒸馏方法为半耦合蒸馏方法。然而作者发现这种方法依然有问题，因为prompt是随机选择的，导致mask解码也是可变的，从而增加了优化难度，最后作者提出了解耦蒸馏的方法，直接蒸馏image embedding，该方法还有一个好处训练时时可以采用简单的MSE loss，而不需要sam原来的Focal Loss和Dice Loss。由于蒸馏的image encoder和原始的image encoder的结果非常相似，feautre的大小一致，这使得mask decoder的微调部分是可选的，论文中提到加上微调（冻结vit，微调mask decoder或者一起微调）的效果会更好。

• Coupled distillation. 实现我们项目目标的一个简单方法是遵循SAM中的官方管道，用较小的图像编码器重新训练新的SAM。如SAM所述，使用ViT-H图像编码器训练SAM需要在256个A100 GPU上花费68小时。用ViT-L或ViT-B取代ViT-H将所需的GPU减少到128个，这对社区中的许多研究人员来说仍然是复制或改进他们的结果的一个不小的负担。按照他们的方法，我们可以进一步采用更小的图像编码器，并用他们提供的分割数据集11-T重新训练新的SAM。注意，提供的数据集中的掩码是由预训练的SAM（使用ViT图像编码器）给出的。本质上，这种再训练过程是知识提炼，它将知识从基于ViT-H的SAM转移到具有较小图像编码器的SAM（见下图左）。

  • 

  • SAM的耦合知识蒸馏。左图表示完全耦合蒸馏，右图表示半耦合蒸馏。

• 从半耦合蒸馏到解耦蒸馏。当执行从原始SAM到较小图像编码器的KD时，困难主要在于图像编码器和组合解码器的耦合优化。直观地说，图像编码器的优化取决于图像解码器的质量，反之亦然。当SAM中的两个模块都处于不良状态时，将它们训练到良好状态更具挑战性。受分治算法的启发，我们提出将KD任务分为两个子任务：图像编码器的提取和掩模解码器的微调。具体来说，我们首先通过将知识从ViT-H转移到较小的编码器来对图像编码器执行KD。由于原始SAM中的掩码解码器已经是轻量级的，我们计划保留其架构。这带来了易于使用的组合解码器的好处，用于微调，而不是从头开始训练它。为了缓解耦合蒸馏的优化问题，一种简单的方法是使用复制和冻结的掩码解码器优化图像编码器（见上图右侧）。

• 冻结操作可以帮助防止掩模解码器的质量因较差的图像编码器而劣化。我们称这种蒸馏为半耦合，因为图像编码器的优化仍然没有与掩模解码器完全解耦。根据经验，我们发现这种优化仍然具有挑战性，因为提示的选择是随机的，这使得掩码解码器可变，从而增加了优化难度。因此，我们建议直接从原始SAM中的ViT-H中提取小图像编码器，而不使用组合解码器，这被称为解耦提取（见下图）。对图像嵌入进行提取的另一个优点是，我们可以采用简单的MSE损失，而不是像sam那样使用焦点损失和 dice 损失的组合来进行掩模预测。

  • 

  • Decoupled distillation for SAM.

• 关于掩码解码器微调的必要性。与半耦合蒸馏不同，上述解耦蒸馏产生了一个轻量级的图像编码器，该编码器可能无法与原始的冻结掩模解码器很好地对齐。根据经验，我们发现这不是真的，因为从学生图像编码器生成的图像编码可以足够接近原始教师编码器的图像编码，这使得在第二阶段中对组合解码器的微调是可选的。期望在冻结的轻量级图像编码器上微调掩模解码器或者将它们联合微调在一起可以进一步提高性能。

• 初步评估。在这里，我们进行了初步的研究，以比较耦合蒸馏和解耦蒸馏。在这里，为了进行绩效评估，我们计算教师SAM和学生SAM在同一提示点上生成的两个掩码之间的mIoU。直观地，通过假设由ViT-H生成的掩模是 GT，较高的mIoU指示较高的掩模预测性能。对于耦合蒸馏，我们采用原始SAM 中提供的带有ViT-B的SAM。它在128个GPU（每个GPU 1个样本）上的SA-1B（11M图像）上训练180k次迭代。相比之下，在我们的解耦蒸馏设置中，我们在2个GPU上训练模型（每个GPU两个样本以节省计算资源），在SA-1B数据集（11k）图像的0.1%样本上进行55k次迭代。总的来说，解耦蒸馏比耦合蒸馏占用的计算资源不到1%，同时实现了耦合sit的mIoU为0.75比0.72的优越性能（在200个样本上平均）。由于ViT-B对移动设备来说仍然是一个不小的负担，因此在下文中，我们基于我们提出的解耦蒸馏，对Tiny-ViT（具有5M参数）进行了实验。

  • 

  • 以ViT-B为图像编码器的SAM的耦合蒸馏和解耦蒸馏的比较。与耦合蒸馏相比，解耦蒸馏性能更好，所需计算资源少于1%。

Experiments

Experimental Setup

• 轻量级图像编码器。我们项目的目标是通过用移动设备的轻量级图像编码器取代默认的ViT-H来获得高效的SAM。作为一个基于ViT的主干，ViT Tiny具有与Deit Tiny相似的参数，但性能更好。例如，在ImageNet-1K上，Deit Tiny实现了72.2%的准确率，而ViT Tiny实现了79.1%的准确率。因此，我们采用ViT-Tiny进行概念验证，以证明我们提出的解耦蒸馏用于训练比原始SAM快得多的轻量级MobileSAM的有效性。所采用的轻量级图像编码器由四个阶段组成，这些阶段逐渐降低分辨率。第一级由具有反向残差的卷积块构建，而其余三级由 Transformer 块组成。在模型开始时，有2个卷积块，步长为2，用于对分辨率进行下采样。通过步长为2的卷积块来处理不同阶段之间的下采样操作。与Wu等人【Tinyvit】不同的是，我们将最后一次下采样卷积中的步长2设置为1，以使最终分辨率与原始SAM的ViT-H图像编码器的分辨率相匹配。MobileSAM的参数推断速度总结在下表中。注意，在第2节中讨论的其他高效图像编码器也可以用作图像编码器。

  • 

  • 原始SAM和MobileSAM中图像编码器的参数和速度的比较。推理速度是在单个GPU上测量的

• 训练和评估细节。对于图像编码器上的解耦KD，我们在单个GPU上用SA-1B数据集的1%训练轻量级编码器8个epoch。我们观察到，考虑到教师图像编码器比我们采用的学生图像编码器重得多，在正向处理上花费了更多的计算（见上文）。为了更快地进行提取，我们遵循Wu等人[TinyViT]的实践，预先保存图像嵌入，这样我们只需要运行一次正向过程。使用单个GPU，我们可以在不到一天的时间内获得MobileSAM。用更多的GPU在更长的时间内训练我们的MobileSAM有望产生更好的性能。执行掩码解码器微调的初步研究进一步提高了MobileSAM的性能，然而，为了简单起见，我们在本文的这个版本中省略了这一步骤。为了定量评估提取的SAM，我们计算原始SAM预测的掩码和我们的MobileSAM之间的mIoU。

MobileSAM performs on par with the orignal SAM

• 对于主要结果，我们报告了具有两种类型提示的预测遮罩：点和框。我们不使用文本提示报告结果，因为SAM的官方github项目没有为文本引导的掩码解码器提供预训练的模型。以point为提示的结果如下图所示，以box为提示的效果如下图所示。我们观察到MobileSAM做出了与原始SAM类似的令人满意的掩模预测。

  • 

  • 以单点为提示遮罩预测

  • 

  • 使用框作为提示来遮罩预测。

• 消融研究。在这里，我们对训练计算对SAM性能的影响进行了消融研究。下表中的结果表明，在相同的迭代次数下，增加批量大小会提高模型性能。此外，在批量大小下，通过增加训练时期，性能也受益于更多的更新迭代。请注意，所有实验都是在单个GPU上进行的。我们预计，增加GPU的数量以允许更大的批量或进一步增加迭代可以进一步提高性能。

  • 

  • 训练计算对MobileSAM性能影响的消融研究。

MobileSAM outperforms FastSAM

• Segment anything v.s. segment everything。请注意，原始SAM论文的标题是“segment anything”，而不是“segmente everything”。正如SAM中所强调的，SAM执行可提示分割的任务，该任务“returns a valid segmentation mask given any segmentation prompt”。提示的作用是指定在图像中分割什么。理论上，只要正确设置提示，任何对象都可以分割，因此称为“segmented anything（分割任何对象）”。相比之下，“segmenting anything“本质上是对象建议生成，其提示不是必要的。在SAM中，选择“segment anything”（对象建议生成）作为下游任务之一，以展示其零样本传输性能。总之，“segment everything”解决了任何对象的可提示分割的基础任务，而“segment everything”解决了所有对象的掩码建议生成的下游任务。由于“segment everything”不一定需要提示，FastSAM以无提示的方式直接用YOLO v8生成掩码提议。为了实现可提示的分割，设计了一个映射算法来从提议掩码集中选择掩码。值得强调的是，评估其泛化/稳健性或研究其多功能性的后续工作主要集中在任何而不是所有模式上，因为前者涉及基础任务。因此，与FastSAM的比较主要集中在“segment anything”上，但为了完整性，我们也提供了关于“segment everything”的比较。

  • 

  • Comparison between segment anything and segment everything.

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

  • 

  • Comparison between FastSAM and MobileSAM.

• mIoU comparison under segment anything mode. 我们进一步比较了预测掩模与原始SAM掩模之间的mIoU。建议FastSAM预测具有多个点的掩模，其中我们选择一个用于前景，另一个用于背景。下表中的结果显示，FastSAM的mIoU比MobileSAM的mIoU小得多，这表明FastSAM的掩模预测与原始SAM的掩膜预测非常不同。此外，当两个提示点之间的距离增加时，FastSAM的mIoU下降得非常快。这主要是由于当前台提示点设置得离后台提示点太近时，FastSAM往往无法预测对象。

  • 

  • mIoU比较。在假设来自原始SAM的预测掩码是 GT 的情况下，mIoU越高表示性能越好。

• Results for segment everything. “segment everything”的结果如下图所示。为了完整性，我们还报告了原始SAM的结果，它生成了一个令人满意的对象提案。我们有两个主要的观察结果。首先，我们的MobileSAM的结果与原始SAM的结果惊人地一致。相比之下，FastSAM的结果往往不那么令人满意。例如，FastSAM经常无法预测某些对象，例如第一张图像中的屋顶。此外，遮罩方案有时很难解释（参见第一幅图像中舞台的遮罩和第二幅图像中天空的遮罩）。其次，FastSAM通常生成具有非平滑边界的掩码。例如，第三个图像中的支柱具有非平滑边界，而原始SAM和我们的MobileSAM没有这个问题。

  • 

  • Comparison of segment everything results.

Conclusion

• 在这项工作中，我们的目标是通过用轻量级图像编码器取代重量级图像编码器，使SAM对移动友好。我们发现，原始SAM文件中训练这种新SAM的 naive 方式会导致不令人满意的性能，尤其是在训练来源有限的情况下。图像编码器和掩码解码器的耦合优化是原因，因此我们提出了解耦蒸馏，其中知识从原始SAM中的图像编码器ViT-H蒸馏到轻量级图像编码器。我们展示了由此产生的轻量级图像编码器可以自动与原始SAM中的掩码解码器兼容。我们的MobileSAM比原始SAM小60多倍，但性能与原始SAM不相上下。此外，我们与并发的FastSAM进行了比较，表明MobileSAM实现了卓越的性能。我们的MobileSAM也比并发的FastSAM快4倍，小7倍，使其更适合移动应用程序。由于我们的MobileSAM保留了原始SAM的所有管道，只替换了图像编码器，因此现有基于SAM的项目可以即插即用，几乎无需任何努力即可从重量级SAM转变为轻量级SAM。