[论文解读] Faster Segment Anything: Towards Lightweight SAM for Mobile Applications

Chaoning Zhang, Dongshen Han+5 more|arXiv (Cornell University)|Jun 25, 2023

Advanced Neural Network Applications被引用 141

一句话总结

本论文提出 MobileSAM，是 Segment Anything 的一种轻量、适合移动端的变体，通过解耦蒸馏来替代重量级图像编码器，在尺寸更小、推理更快的情况下达到与原始 SAM 相当的分割效果，与 FastSAM 相比。

ABSTRACT

Segment Anything Model (SAM) has attracted significant attention due to its impressive zero-shot transfer performance and high versatility for numerous vision applications (like image editing with fine-grained control). Many of such applications need to be run on resource-constraint edge devices, like mobile phones. In this work, we aim to make SAM mobile-friendly by replacing the heavyweight image encoder with a lightweight one. A naive way to train such a new SAM as in the original SAM paper leads to unsatisfactory performance, especially when limited training sources are available. We find that this is mainly caused by the coupled optimization of the image encoder and mask decoder, motivated by which we propose decoupled distillation. Concretely, we distill the knowledge from the heavy image encoder (ViT-H in the original SAM) to a lightweight image encoder, which can be automatically compatible with the mask decoder in the original SAM. The training can be completed on a single GPU within less than one day, and the resulting lightweight SAM is termed MobileSAM which is more than 60 times smaller yet performs on par with the original SAM. For inference speed, With a single GPU, MobileSAM runs around 10ms per image: 8ms on the image encoder and 4ms on the mask decoder. With superior performance, our MobileSAM is around 5 times faster than the concurrent FastSAM and 7 times smaller, making it more suitable for mobile applications. Moreover, we show that MobileSAM can run relatively smoothly on CPU. The code for our project is provided at \href{https://github.com/ChaoningZhang/MobileSAM}{ extcolor{red}{MobileSAM}}), with a demo showing that MobileSAM can run relatively smoothly on CPU.

研究动机与目标

• 推动在资源受限设备上移动部署 SAM。
• 通过替换重量级图像编码器来减小模型规模并提升推理速度。
• 在不进行大量再训练的情况下，保持与原始 SAM 掩码解码器的兼容性。
• 证明基于蒸馏的训练能够产生轻量且准确的分割模型。

提出的方法

• 通过对教师 ViT-H 的知识蒸馏，用轻量级编码器替换 SAM 中的 ViT-H 图像编码器。
• 应用解耦蒸馏将知识转移给小型学生编码器，同时保持原始掩码解码器冻结或轻微微调。
• 在图像嵌入对齐中使用均方误差（MSE）损失，而不是相关工作中使用的聚焦损失和 Dice 损失的组合。
• 可选地对掩码解码器进行微调，但需展示解耦蒸馏已能在学生编码器与解码器之间实现良好的对齐。
• 使用 mIoU 和推理速度指标，对 MobileSAM 与原始 SAM 和 FastSAM 进行评估。

实验结果

研究问题

• RQ1 从重量级 SAM 编码器蒸馏得到的轻量级图像编码器是否能够达到与原始 SAM 相当的分割质量？
• RQ2 解耦蒸馏在训练轻量级 SAM 时是否优于耦合蒸馏或半耦合蒸馏？
• RQ3 在 segment-anything 任务中，MobileSAM 在准确性（mIoU）和效率（参数量、速度）方面与 FastSAM 的比较如何？
• RQ4 MobileSAM 是否能够在 CPU 上高效运行以用于设备端应用？

主要发现

• MobileSAM 将编码器参数大约减少 100 倍，总参数量大约减少 60 倍，同时实现与原始 SAM 相当的性能。
• 在单个 GPU 上，MobileSAM 处理一张图像大约需要 10 ms（编码器 8 ms，解码器 4 ms）。
• MobileSAM 比 FastSAM 快约 5 倍、尺寸小 7 倍，同时在 segment-anything 场景中提供更强的性能。
• 解耦蒸馏（直接用重教师训练轻量级编码器而不进行解码器联合训练）在初步实验中获得更好的 mIoU（0.75），比耦合蒸馏（0.72）更好。
• MobileSAM 能在 CPU 上相对流畅地运行，实现移动设备部署。