【多模态】15、MobileSAM | 首个专为移动端设计的更快的 SAM

论文：FASTER SEGMENT ANYTHING: TOWARDS LIGHTWEIGHT SAM FOR MOBILE APPLICATIONS

代码：https://github.com/ChaoningZhang/MobileSAM

出处：韩国庆熙大学

时间：2023.06.27

一、背景

Meta 提出的 SAM 以其出色的对任何感兴趣目标都可以进行分割的能力引起了广泛的关注，SAM 的结构如图 1 所示，包括两个部分：

• ViT-based image encoder
• prompt-guided mask decoder

SAM 是一个 label-free 的分割模型，可以和其他模型结合来进行进一步的下游任务，如 text-guided 分割、图像编辑等

在移动设备非常普及的情况下，很多图像编辑都是在移动端，但 SAM 的 image encoder 很大，所以设计一个对移动端友好的 SAM 是很迫切的。

故本文提出了 MobileSAM，主要就是设计了一个适用于移动端设备的轻量级 SAM

二、方法

按一般的想法来说，既然 image encoder 太大了，那么就缩小 image encoder

比如，将 ViT-H 使用 ViT-B 来替换，使用不同大小的 image encder 的模型参数量如表 3 所示：

从头训练一个 SAM（使用 ViT-L 或 ViT-B 做 image encoder）需要使用 128 个 GPU 训练数天，所以，重新训练代价也很大。

作者认为，优化的难度在于 image encoder 和 mask encoder 被耦合在一起了

所以，作者将 image encoder 和 mask encoder 进行了解耦：

• 首先，将 ViT-H 的知识蒸馏到 tiny ViT
• 然后，finetune mask encoder 来对齐蒸馏的小 image encoder

基于此，设计轻量级 SAM 的任务就转变为解耦蒸馏了，就很简单高效

MobileSAM 将 encoder 的参数降低了 100 倍，将整体参数降低了 60 倍

MobileSAM 推理速度：

• 单个图像推理速度大概 10ms（8ms 是 image encoder，2ms 是 mask encoder）

MobileSAM 和 FastSAM 的速度对比：

• MobileSAM 比 FastSAM 小 7 倍，快 4 倍

2.1 耦合蒸馏

要实现对移动端友好的 SAM 的一个直观方法是重新训练一个 image encoder 小的 SAM，但训练代价太大了，所以可以考虑蒸馏的的方法，如图 2 左侧所示，是用最终大模型的 mask 来指导小模型的 mask。

2.2 从半蒸馏到解耦蒸馏

当直接用 mask 来指导蒸馏时，其难点在于 image encoder 和 mask decoder 是联系在一起的，两者是互相依赖的，所以可以：

• image encoder：蒸馏
• mask encoder：finetuned（因为 SAM 中的 mask encoder 本来就很小，故保持结构不变）

如图 2 右侧所示，这种也叫半解耦蒸馏（semi-coupled），对 image encoder 蒸馏的同时对 mask encoder 参数冻结，冻结能够使得 mask decoder 的效果不变，不会被 image encoder 效果的好坏影响。

但这种方法还会有问题，因为 prompt 是随机的，会导致 mask decoder 多变，难以优化

所以，本文的蒸馏方法：

• 完全解耦蒸馏
• 直接对 image embedding 进行蒸馏
• 完全将 image encoder 和 mask decoder 解耦开来了
• 这样也可以直接使用 MSE loss，不用再结合 focal 和 dice loss 了

解耦的蒸馏方法计算资源是耦合蒸馏方法的 1%，但却达到了 0.75 mIoU：

三、效果