发表时间：2023.06.27
论文地址：https://arxiv.org/pdf/2306.14289.pdf
项目代码：https://github.com/ChaoningZhang/MobileSAM
Segment anything model(SAM）是一个快速引导的视觉基础模型，用于从其背景中找出感兴趣的对象。自从Meta研究团队发布了SA项目以来，SAM因其令人印象深刻的zero-shot性能和兼容的其他模型而吸引了广泛的关注，如细粒度控制的图像编辑。许多这样的用例需要在资源约束的边缘设备上运行，比如移动应用程序。在这项工作中，我们的目标是使SAM移动友好取代轻量级图像编码器。像最初的SAM论文中这样的训练新SAM的方法会导致不令人满意的性能，特别是在训练来源有限的情况下。我们发现，这主要是由图像编码器和mask解码器的耦合优化引起的，因此我们提出了解耦蒸馏方法。具体地说，我们将原始SAM中的图像编码器ViT-H提取到轻量级图像编码器中，可以自动与原始SAM中的mask解码器相兼容。训练可以在不到一天的时间内完成，由此产生的轻量级SAM被称为MobileSAM，它要小60倍以上，但性能与原来的SAM相当。为了提高推理速度，MobileSAM对每幅图像运行约10 ms：在图像编码器上运行8 ms，在mask解码器上运行2 ms。由于具有优越的性能和更高的通用性，我们的MobileSAM比并发的FastSAM要小7倍，速度要快4倍，使其更适合移动应用程序。
关键解读
基本背景
SAM类似于视觉界的chatgpt模型，其提供了图像的编码的基本能力，是视觉界的大模型，可以应用于各种图像处理任务中；
SAM通过在服务器端进行图像特征编码可以运行在轻量化的web端，但其图像编码器其过于巨大。
亮点
1、使用解耦蒸馏方法（只对图像编码器进行蒸馏），使backbone与原始的解码器相适应，整个训练在一个GPU上不到一天，将编码器参数减少100倍，总参数减少60倍。
2、蒸馏后的图像编码器运行为8 ms，mask解码器运行为2 ms，总体运行时间为10ms，比FastSAM快4倍。
3、其基于conv和transformer设计了轻量化的图像编码器；同时，为了加快训练，保存了教师模型预测的特征编码，减少了知识蒸馏中教师模型forward的时间。
4、只使用1%的SA-1B数据集进行训练，而FastSAM使用了2%的数据。
5、在4.2节的消融实验中证明了，在海量数据集下batch size越大，效果越好；同时，海量数据的epoch数基本上都是各位数。
MobileSAM本质就是对SAM中ViT模型的知识蒸馏，使用了原来SAM中的mask解码器，然后在训练数据上只使用了原来的1%，故才可实现在一天内训练好模型
1 Introduction
ChatGPT彻底改变了NLP领域，标志着生成式人工智能（AIGC，人工智能生成内容）的突破。Bommasani等人[2021]在web规模的文本数据集上训练的基础模型，构成了gpt系列模型。随着NLP基础模型的成功，大部分学者通过对比学习训练了图像编码器和文本编码器。最近，Meta研究团队发布了“Segment Anything”项目，其中提出了一个被命名为SAM的即时引导视觉基础模型，并被认为是视觉的GPT时刻。SAM由两个组件组成：基于vit的图像编码器和提示引导的mask解码器，它们可以按顺序工作（见图1）。
自出现以来，SAM因多种原因引起了广泛的关注。首先，它是第一个表明视觉可以像NLP一样去追求一条将基础模型与快速工程相结合的道路。其次，它是第一个进行无标签分割，这是一个与标签预测并行的基本视觉任务。此外，这一基本任务使SAM与其他模型兼容，实现高级视觉应用，如文本引导分割、细粒度控制的图像编辑。然而，许多此类用例都需要在资源受限的边缘设备上运行，比如移动应用程序。如官方演示所示，通过在服务器端处理的图像嵌入，SAM可以在资源受限的设备上工作，因为掩模解码器是轻量级的。使SAM pipeline计算量较大的原因在于这个巨大的图像编码器。在这项工作中，我们研究了如何获得一个适合于资源受限的移动设备的轻量级SAM，因此它被称为MobileSAM。
鉴于默认图像编码器是基于ViT-H，直接获得MobileSAM是遵循官方pipeline在基里洛夫等人[2023]再训练一个新的SAM与一个较小的图像编码器，取代ViT-H更小的ViT-L甚至更小的ViT-B。表1总结了不同规模下图像编码器SAM的SAM参数。
如Kirillov等人[2023]所述，作为图像编码器，使用ViT-L或ViT-B训练新的SAM需要128个gpu。这种资源密集型的再培训可能是重现或改进其结果的一个重要负担。这种优化困难主要来自于图像编码器和mask解码器的耦合优化。基于这种理解，我们提出了解耦图像编码器和mask解码器的优化。具体地说，我们首先将知识从默认的图像编码器ViT-H提取到一个很小的ViT。然后，我们可以微调原始SAM中的mask解码器，以更好地与蒸馏后的图像编码器对齐。值得强调的是，对齐优化是可选的，因为轻量级图像编码器是从默认图像编码器中提取出来的，这保证了它与默认mask解码器的固有对齐。
通过将寻找新的SAM pipeline的问题转化为解耦的蒸馏，我们的方法具有简单和有效的优点，同时以低成本的成本重复（在一个GPU上不到一天）。所得到的MobileSAM使编码器参数减少100倍，总参数减少60倍。令人惊讶的是，这种轻量级的MobileSAM的性能与最初的重量级SAM相当，这是推动移动应用程序的SAM的重要一步。对于MobileSAM的推理，单个图像只运行大约10毫秒：图像编码器上运行8 ms，mask解码器上运行2 ms。值得强调的是，我们的MobileSAM比并发的FastSAM Zhao等人[2023]小7倍，快4倍，同时实现了卓越的性能。
2 Related work
SAM: generalization and versatility. 自今年4月初出现以来，已经出现了许多项目和论文，从不同的角度调查SAM。鉴于SAM声称可以分割任何内容，一系列作品报告了其在现实世界中的表现，包括医学图像、伪装对象和透明对象。研究结果一致表明，SAM在一般设置中工作良好，但在上述具有挑战性的任务中无效。另一个重要的研究方向集中在提高SAM，以提高其实用性。Attack-SAM已经表明，SAM的输出mask可以通过恶意生成的对抗性攻击轻松操作。Qiao等人[2023b]进一步对SAM进行了全面的稳健性评估，从风格转移、常见腐蚀到局部闭塞和对抗性扰动。Qiao等人[2023b]发现SAM具有较高的鲁棒性，但对对抗性扰动没有影响，这与Zhang等人[2023e]的发现一致。另一系列工作的重点是展示SAM的多功能性。 Grounded SAM IDEA[2023]是将Grounding DINO Liu等人[2023a]与SAM结合起来进行任何文本输入分割的开创性工作。具体来说，它依赖于Grounding DINO从文本中生成一个边界框，然后生成的边界框可以用作分割mask的提示。SAM预测的mask没有类别信息和多种作品，Chen等人[2023]，Park [2023]结合SAM与其他类似CLIP的模型，以对任何东西进行语义分割。除了对象分割，多部作品也在其他领域也显示了其多功能性，包括图像编辑Rombach等人[2022]，以及Yu等人[2023]，视频对象跟踪杨等人[2023]，Zxyang [2023]。除了二维视觉之外，SAM的研究也扩展到三维物体重建Shen等[2023]，Kang等[2022]，展示了其辅助从单个图像生成三维模型的能力。关于SAM的完整调查， proposal 读者参考Zhang等人[2023c]。
对于SAM的研究有很多，不愧是新一代基础视觉大模型
ViT: lightweight and efficient 早期的移动视觉应用主要由轻量级cnn驱动，如MobileNet霍华德等人[2017]及其改进的变异，桑德勒等人[2018]，霍华德等人[2019]。MobileNet的核心思想是将一个正规卷积块分为深度卷积和点卷积，这大大减少了模式参数和计算时间。自从VIT出现以来，许多研究者都试图使其轻量级和高效。根据原始ViT论文，Touvron等人[2020]介绍了较小的改动，提出了Deit-Small (Deit-S) and Deit-Tiny (Deit-T) ViT-Small and ViT-Tiny。MobileViT[2021]是一项结合ViT和标准卷积以提高其性能的开创性工作，它优于MobileNet v2。主要动机是利用CNN的局部表示能力，该实践之后是多个后续工作，旨在提高模型速度。最近在轻量级和更快的ViT方面的进展是对我们提出的解耦蒸馏的补充，以使下一代SAM适用于资源有限的移动设备。
提出了ViT模型的轻量化
3 Mobile-Friendly SAM
3.1 Background and Project Goal
Background on SAM. 在这里，我们首先总结了SAM的结构和它的工作原理。SAM由一个基于vit的图像编码器和一个提示引导的mask解码器组成。图像编码器将图像作为输入并生成嵌入，然后将其输入mask解码器。mask解码器生成一个mask，根据点（或框）等提示从背景中删除任何对象。此外，SAM允许为相同的提示生成多个mask，以解决模糊性问题，这提供了宝贵的灵活性。考虑到这一点，本工作维护了SAM的 pipeline，首先采用基于vit的编码器生成图像嵌入，然后采用提示引导解码器生成所需的mask。这个 pipeline是为“segment anything”优化设计的，它可以用于“segment anything”的下游任务(见章节4.3以供更多讨论。）
Project goal. 这个项目的目标是生成一个移动友好的SAM（MobileSAM），它以轻量级的方式实现令人满意的性能，并且比原来的SAM要快得多。原始SAM中的提示引导mask解码器的参数小于4M，因此被认为是轻量级的。给定一个由编码器处理的图像嵌入，如他们的公开演示所示，SAM可以在资源有限的设备中工作，因为mask解码器是轻量级的。然而，原始SAM中的默认图像编码器是基于ViT-H，参数超过600M，这是非常重量级的，使得整个SAM pipeline与移动设备不兼容。因此，获得移动友好型SAM的关键在于用轻量级图像编码器替换为轻量级图像编码器，它也会自动保留其原始SAM的所有功能和特性。在下面，我们将详细阐述我们为实现这个项目目标所提出的方法。
论述了SMA中ViT轻量化的必要性
3.2 Proposed Method
Coupled distillation. 实现我们的项目目标的一种直接的方法是遵循Kirillov等人[2023]的官方 pipeline，用一个更小的图像编码器重新训练一个新的SAM。如Kirillov等人[2023]所述，使用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）。
From semi-coupled to decoupled distillation. 当执行从原始SAM到使用较小的图像编码器的KD时，困难主要在于对图像编码器和组合解码器的耦合优化。直观地说，图像编码器的优化取决于图像解码器的质量，反之亦然。当SAM中的两个模块都处于坏状态时，将它们都训练到良好状态就更具挑战性了。受Zhang等人[2022c]的启发，我们提出将KD任务划分为两个子任务：图像编码器蒸馏和mask解码器微调。具体地说，我们首先在图像编码器上通过将知识转移到一个较小的编码器来执行KD。由于原始SAM中的mask解码器已经是轻量级的了，所以我们计划保留它的架构。这带来了一个好处，一个容易使用的组合解码器进行微调，而不是从头开始训练它。为了缓解耦合蒸馏的优化问题，一种简单的方法是使用复制和冻结的mask解码器来优化图像编码器（见右图2）。
冻结操作可以帮助防止掩模解码器的质量被糟糕的图像编码器恶化。我们称这种蒸馏为半耦合，因为图像编码器的优化仍然没有完全与mask解码器完全解耦。根据经验，我们发现这种优化仍然具有挑战性，因为提示符的选择是随机的，这使得mask解码器成为变量，从而增加了优化的难度。因此，我们 proposal 直接从原始SAM中的ViT-H中提取小图像编码器，而不采用组合解码器，称为解耦蒸馏（见图3）。在图像嵌入上执行蒸馏的另一个优点是，我们可以采用一个简单的MSE损失，而不是结合使用focal loss林等人[2017]和dice loss米列塔里等人[2016]来预测。
On the necessity of mask decoder finetuning. 与半耦合蒸馏不同，上述解耦蒸馏产生了一个轻量级的图像编码器，它可能与原始的冻结掩模解码器不能很好地对齐。根据经验，我们发现这是不正确的，因为从学生图像编码器生成的图像编码可以与原始教师编码器的图像编码足够接近，这使得在第二阶段对组合解码器的微调是可选的。期望对冻结的轻量级图像编码器上的mask解码器进行微调或联合微调，可以进一步提高性能。
Preliminary evaluation. 在此，我们进行了一个初步的研究，以比较耦合蒸馏和解耦蒸馏。在这里，为了进行绩效评估，我们计算了由教师SAM和学生SAM在同一提示点上生成的两个面具之间的mIoU。直观地看，假设ViT-H生成的掩模是GT的，较高的mIoU表明掩模预测性能较高。对于耦合蒸馏，我们采用了原始SAM Kirillov等人[2023]中提供的SAM和ViT-B。在128个GPU（每个GPU1个样本）上的SA-1B（11M图像）上进行训练，进行180k次迭代。相比之下，在我们的解耦蒸馏设置中，我们在2个GPU（每个GPU2个样本以节省计算资源）和0.1%的SA-1B数据集（11k）图像样本上训练模型，进行55k次迭代。总的来说，解耦蒸馏比耦合蒸馏需要的计算资源不到1%，mIOU分别为0.75比0.72（平均200个样本）。由于ViT-B对于移动设备来说仍然是一个重要的负担，因此下面我们实验使用tiny基于我们提出的解耦蒸馏的Vit(Wu等人(2022]。
论证了解耦蒸馏的必要性，其可以大量节省训练时间
4 Experiments
4.1 Experimental Setup
Lightweight Image Encoder. 我们项目的目标是通过用用于移动设备的轻量级图像编码器替换默认的ViT-H来获得一个高效的SAM。作为一种基于vit的骨干，ViT-Tiny具有与Deit-Tiny具有相似的参数，但性能更好。例如，在ImageNet-1K上，Deit-Yiny的准确率达到72.2%，而ViT-Tiny的准确率达到79.1%。因此，我们采用ViT-Tiny来证明概念，以证明我们提出的解耦蒸馏的有效性，以训练一个轻量级的MobileSAM的SAM，可以比原来的SAM快得多。所采用的轻量级图像编码器由四个阶段组成，逐渐降低了分辨率。第一阶段由conv和倒残差结构组成，而其余三个阶段由transformer块组成。在模型的开始，有2个conv，步幅为2，用于降采样分辨率。不同阶段之间的降采样操作采用卷积块处理，步积块为2。与Wu等人[2022]不同，我们将上采样卷积的步幅设置为2，使最终分辨率与原始SAM的ViT-H图像编码器相匹配。请注意，在第2节中讨论的其他高效的图像编码器也可以被采用作为图像编码器。
Training and evaluation details. 对于图像编码器上的解耦KD，我们使用1%的SA-1B数据集Kirillov等人[2023]在单个GPU（RTX3090）上训练轻量级编码器。我们观察到，考虑到它明显比我们所采用的学生图像编码器更重，在教师图像编码器的前向过程上花费了更多的计算量（见上文）。为了加快蒸馏速度，我们遵循Wu等人[2022]的实践，预先保存图像嵌入，以便我们只需要运行一次正向过程。使用一个GPU，我们可以在不到一天的时间内获得MobileSAM。用更长时间的gpu训练我们的MobileSAM，有望产生更好的性能。执行mask解码器微调的初步研究进一步提高了MobileSAM的性能，但是，为了简单起见，我们在本文的这个版本中省略了这一步。为了定量评价蒸馏后的SAM，我们计算了原始SAM预测的掩模和我们的MobileSAM之间的mIoU。
基于conv和transformer设计了轻量化的图像编码器；同时，为了加快训练，保存了教师模型预测的特征编码
4.2 MobileSAM performs on par with the orignal SAM
对于主要结果，我们用两种类型的提示来报告预测的mask：point和bboxes。我们没有使用文本提示来报告结果，因为SAM的官方github项目没有为文本引导的mask解码器提供预先训练过的模型。以点为提示的结果如图4所示，
以方框为提示的结果如图5所示。我们观察到MobileSAM做出了与原始SAM相似的令人满意的掩模预测。
Ablation study. 在这里，我们对训练计算对SAM性能的影响进行了消融研究。表4中的结果表明，在相同的迭代次数下，增加批处理大小可以提高模型的性能。此外，在批处理大小下，性能还可以通过增加训练时代来获得更多的更新迭代。请注意，所有的实验都是在单个GPU上进行的。我们期望增加gpu的数量以允许更大的批处理大小或进一步增加迭代可以进一步提高性能。
4.3 MobileSAM outperforms FastSAM in All Aspects
Segment anything v.s. segment everything . 请注意，最初SAM的论文Kirillov等人[2023]的标题是“分割任何东西”，而不是“分割一切”。正如Kirillov等人[2023]所强调的，SAM执行快速分割任务，“根据输入的提示符返回有效的分割mask”。提示符的作用是指定要在图像中分割的内容。理论上，只要提示符设置正确，任何对象都可以被分割，因此，它被称为“分割任何东西”。相比之下，“分割一切”本质上是对象proposal生成，因此不需要提示。
综上所述，“分割任何东西”解决了任何对象的快速分割的基础任务，而“分割一切”解决了所有对象的mask proposal生成的下游任务。由于“分割一切”并不一定需要提示输入，因此FastSAM会以无提示的方式与YOLO v8直接生成mask proposal 。为了实现可提示分割，设计了一种映射算法，从 proposal mask集中选择mask。值得强调的是，评估其泛化/鲁棒性或研究其多用途性的后续工作主要集中在任何模式上，而不是所有模式，因为前者解决了基础任务。因此，与FastSAM的比较主要集中在“分割的任何东西”上，但为了完整性，我们也提供了一个关于“分割的一切”的比较。
MobileSAM is faster and smaller. FastSAM包括一个基于yolov8的检测分支和一个基于YOLACT的分割分支，用于执行无提示mask proposal 生成。它有68个M的参数，需要40个ms来处理一个图像。相比之下，MobileSAM的10M参数更少，这是明显更小的。就推理速度而言，在一个GPU上，处理一个图像需要40 ms，而我们的图像只需要10 ms，比FastSAM快4倍。
mIoU comparison under segment anything mode. 我们进一步比较了预测的掩模与原始SAM之间的mIoU。请注意，FastSAM不能用单个点作为原始SAM来预测mask。相反，它需要至少两个提示点：一个是前景，另一个是背景。表6的结果显示，FastSAM的mIoU比MobileSAM要小得多，说明FastSAM的掩模预测与原始SAM有很大的不同。此外，当两个提示点之间的距离增大时，FastSAM的mIoU减小得非常快。这主要是由于当前景提示点设置得离背景提示点太近时，FastSAM往往无法预测对象。
Results for segment everything.
“分割一切”的结果如图6所示。为了完整性，我们还报告了原始SAM的结果，它生成了一个令人愉快的对象 proposal 。我们有两个主要的观察结果。首先，我们的MobileSAM的结果与最初的SAM的结果惊人地一致。相比之下，FastSAM的结果往往不太令人满意。例如，FastSAM经常无法预测一些对象，比如第一张图像中的屋顶。此外，面具的提议有时很难解释(参见第一个图像中的舞台的面具和在第二幅图像中的天空). 其次，FastSAM经常生成具有非平滑边界的mask，为此，我们建议读者放大以查看图6中的详细信息。例如，第三幅图像中的支柱具有非平滑的边界，而原始的SAM和我们的MobileSAM则没有这个问题。
5 Conclusion
在这项工作中，我们的目标是使SAM移动友好取代轻量级图像编码器。我们发现，像在最初的SAM论文中那样，训练这样一个新的SAM的天真方法会导致不令人满意的性能，特别是在训练源有限的设置下。图像编码器和掩模解码器的耦合优化是原因，因此我们提出解耦蒸馏，从原始SAM中的图像编码器ViT-H提取知识到轻量级图像编码器。我们表明，得到的轻量级图像编码器可以自动兼容原始的mask解码器。我们的MobileSAM比原来小60多倍，但性能与原来的SAM相当。此外，我们还与并发的FastSAM进行了比较，结果表明MobileSAM实现了优越的性能。我们的MobileSAM也比并发的FastSAM快4倍，速度小7倍，使它更适合移动应用程序。由于我们的MobileSAM保留了原始SAM的所有 pipeline，并且只是替换了图像编码器，因此对于现有的基于SAM的项目，它可以从重量级的SAM转换为轻量级的SAM。