【论文解读】SAM模型超级进化：面向移动端的轻量级SAM，比FastSAM快4倍！（附论文地址）

论文地址：https://arxiv.org/pdf/2306.14289.pdf

这篇论文的标题是《FASTER SEGMENT ANYTHING: TOWARDS LIGHTWEIGHT SAM FOR MOBILE APPLICATIONS》，由Chaoning Zhang等人撰写，发表于2023年。

文章主要探讨了如何将Segment Anything Model（SAM）优化，使其更适合移动设备应用。以下是对论文内容的解读：

研究背景与动机

Segment anything model (SAM) 是一个基于提示的视觉基础模型，用于将感兴趣的物体从其背景中切割出来。自Meta研究团队发布SAM(Segment anything model)项目以来，引起了极大的关注，因为它具有令人印象深刻的零样本迁移特性和与其他模型兼容的高通用性，可用于高级视觉应用，如具有细粒度控制的图像编辑。许多这样的用例都需要在资源受限的边缘设备上运行，如移动应用程序。在这项工作中，我们的目标是通过用一个轻量级的图像编码器取代重量级的图像编码器，使SAM成为移动友好型的。最初的SAM模型体积较为庞大，分割效果最好的ViT-H 模型达632M，因此在推理速度上不尽人意。我们发现，这主要是由图像编码器和掩码解码器的耦合优化造成的，在此基础上，我们提出了解耦蒸馏法。具体来说，我们将原始SAM中的图像编码器ViT-H的知识提炼成1个轻量级的图像编码器，它可以与原始SAM中的掩码解码器自动兼容。训练可以在1天之内在单个GPU.上完成，由此产生的轻量级SAM被称为MobileSAM， 它的体积小了60多倍，但性能与原始SAM相当。就推理速度而言，MobileSAM每 张图像的运行时间约为10ms:图像编码器为8ms,掩码解码器为2ms。凭借卓越的性能和更高的通用性,我们的MobileSAM比同期的FastSAM小7倍，快4倍，使其更适合移动端的应用。

SAM

上图给出了SAM架构示意图，它包含一个"重量级"ViT编码器与一个提示词引导Mask解码器。解码器以图像作为输入，输出将被送入Mask解码器的隐特征(embedding)；Mask解码器将基于提示词(如point、bbox)生成用于目标分割的Mask。此外，SAM可以对同一个提示词生成多个Mask以缓解"模棱两可"问题。

MobileSAM

延续SAM架构体系：采用轻量级ViT解码器生成隐特征，然后采用提示词引导解码器生成期望的Mask。

为了训练新的MobileSAM，作者提出了一种新颖的解耦蒸馏方法来训练MobileSAM。这个方法包括两个关键步骤，

其优势在于简单、有效，且成本低（在单GPU上不到一天的训练时间）。

图像编码器蒸馏

将原始SAM中的重型图像编码器（ViT-H）的知识转移到一个轻型图像编码器中。这一步骤通过简单的均方误差（MSE）损失来实现，而不是使用复杂的损失函数。

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

可选的掩码解码器微调

尽管这一步是可选的，但作者发现通过微调可以进一步提高MobileSAM的性能，尤其是在与轻型图像编码器的兼容性方面。

从半耦合到解耦蒸馏。

耦合蒸馏和解耦蒸馏对于使用ViT-B作为图像编码器的SAM的比较。解耦蒸馏表现更好，并且所需的计算资源不到耦合蒸馏的1%。

从上表对比了耦合蒸馏与解耦合蒸馏的初步对比。可以看到：

• 从指标方面，解耦合蒸馏方案指标稍高，0.75mIoU vs 0.72mIoU；

• 从训练GPU方面，解耦合蒸馏方案仅需两个GPU，远小于耦合蒸馏方案的128卡，大幅降低了对GPU的依赖；

• 从迭代次数方面，解耦合蒸馏方案仅需55k次迭代，远小于耦合蒸馏方案的180K，大幅降低了训练消耗；

• 从训练数据方面，解耦合蒸馏方案仅需11K数据量，远小于耦合蒸馏方案的11M，大幅降低了数据依赖。

尽管如此，但ViT-B对于移动端部署仍然非常困难。因此，后续实验主要基于TinyViT进行。

实验与评估

作者基于ViT-Tiny进行本文所提方案的有效性验证，所得MobileSAM与原生SAM的参数+速度的对比可参考上表。在训练方面，仅需SA-1B的1%数据量+单卡(RTX3090)，合计训练8个epoch，仅需不到一天即可完成训练。

具体的实验设置

• 使用1%的SA-1B数据集进行训练，共8个epoch。

• 训练过程中，作者采用了小批量大小和适量的迭代次数，以优化训练效率。

• 评估指标主要为平均交并比（mIoU），用于衡量MobileSAM生成的掩码与原始SAM生成的掩码之间的一致性。

实验结果

• MobileSAM的参数数量比原始SAM减少了60倍，推理速度提高了约5倍。

• MobileSAM在单GPU上每张图片的处理时间约为10ms，其中图像编码器处理时间为8ms，掩码解码器处理时间为4ms。

• MobileSAM在CPU上的运行也相对流畅，表明其适用于资源受限的环境。

• 通过对比实验，作者证明了MobileSAM在性能上与原始SAM相当，但在速度和大小上明显优于FastSAM。

下图给出了point与bbox提示词下MobileSAM与原生SAM的结果对比，可以看到：MobileSAM可以取得令人满意的Mask预测结果。

下图从Segment everything角度对比了SAM、FastSAM以及MobileSAM三个模型，可以看到：

• MobileSAM与原生SAM结果对齐惊人的好，而FastSAM会生成一些无法满意的结果

• FastSAM通常生成非平滑的边缘，而SAM与MobileSAM并没有该问题

结论与贡献

文章最后总结了MobileSAM的主要贡献：

• 提出了一种新的解耦蒸馏方法，使得SAM模型能够适应移动设备，同时保持与原始SAM相同的性能。

• MobileSAM显著减少了模型大小和提高了推理速度，实现了轻量化和快速化，非常适合移动设备应用。

• MobileSAM的代码和演示已经公开，便于研究者和开发者进一步研究和应用。

这篇论文的研究成果对于推动SAM在移动应用领域的应用具有重要意义，同时也为未来相关技术的发展提供了新的思路和方法。