FASTER SEGMENT ANYTHING:TOWARDS LIGHTWEIGHT SAM FOR MOBILE APPLICATIONS

摘要

        基于提示符引导的视觉基础模型SAM可以从图片中分割出任何感兴趣的部分。自从Meta研究团队发布SA项目以来，SAM不仅可以实现零样本迁移，与其他模型的兼容性也很好，但在使用SAM时，可能会受限于计算资源和具体的设备。为了解决这个问题，很多人会想到重新训练一个与SAM类似的新模型，但由于图片编码器和掩码解码器之间存在耦合优化（the coupled optimization），这样做需要巨大的计算开销和庞大的数据集，因此本文提出了一种解耦蒸馏（decoupled distillation）方法，将原来图片编码器（教师编码器）的知识蒸馏到另一个轻量级的图片编码器（学生解码器）上，由此得到一个新的模型Mobile SAM，他和原来的SAM性能相当，但参数减少了60倍，在一个单核GPU上用不到一天的时间便可完成训练。对于推理速度，Mobile SAM处理一张图片仅需要0.01秒，比Fast SAM快了4倍（大小比Fast SAM小了7倍）。这使得Mobile SAM不仅性能更好、用途更广，而且对于移动应用也更适合。

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

1、introduction

        ChatGPT的出现，彻底改变了自然语言领域，也标志着生成式人工智能取得重大突破。在CV领域，受ChatGPT的启发，很多工作一直通过对比学习学习图片编码器和文本编码器，最近Meta研究团队发布了“Segment Anything”，也就是SAM，它被认为和NLP中的ChatGPT一样可以作为视觉领域的基础模型，由基于ViT的图片编码器和基于提示词引导的掩码解码器两部分构成。

        自SAM出现以来，便引起了人们的广泛关注。原因有很多：首先，它第一次向人们展示视觉领域也可以像NLP一样，可以将基础模型和特征提示工程结合起来（prompt engineering）；其次，它是第一个进行无标签分割的模型；此外，SAM可以与其他模型兼容，实现高级视觉应用，比如文本引导的图像分割、使用细粒度控制（fine-grained control）的图像编辑，但这些应用需要运行在资源受限的设备上（如：手机、平板等）。本文我们通过将原来SAM中的图片编码器换成一个轻量级的图片编码器，让SAM成功的运行在资源受限的设备上，本文主要研究如何获取一个轻量级的SAM，也就是Mobile SAM。

        我们解耦了SAM图像编码器和掩码解码器优化过程，将原来图像解码器（image encoder ViT-H）上的知识蒸馏到到一个tiny ViT上，然后微调掩码解码器，让掩码解码器更好的与新的图片编码器对齐。通过这样的操作，让SAM的复制变得简单高效，在保持性能不变的同时参数大大减少，性能大大超过了Fast SAM。

（最后两段的大部分内容和摘要是一样的，所以翻译的时候省略了很多）

2、相关工作

SAM: generalization and versatility.

ViT: lightweight and efficient.

3、Mobile-Friendly SAM

3.1 Background and Project Goal

Background on SAM：

        这里主要讲SAM的结构和工作原理。

Project goal：

        使用一个轻量级的图片编码器替换SAM原有的图片编码器，从而得到一个轻量级的SAM，也就是Mobile SAM。

3.2 Proposed Method

Coupled distillation.

        我们使用一个更小的图片编码器重新训练了一个SAM（Mobile SAM），在训练过程中，数据集中使用的掩码是由原来预训练好的SAM生成的。从本质上讲，重新训练这个过程就是知识蒸馏，它将ViT-H-based SAM的知识转移到Mobile SAM上。

From semi-coupled to decoupled distillation

        当将原始SAM模型的知识蒸馏到具有较小的图像编码器的SAM模型时（KD），面临的主要困难是图像编码器和联合解码器之间的耦合优化。受分治法思想的启发，我们将整个知识蒸馏（KD）过程分为两个子任务：图片编码器的蒸馏和掩码解码器的微调（这样我们就只需要将来自ViT-H的图像编码器的知识迁移到一个更小的图像解码器上，另外保持掩码解码器的结构不变，只需对其进行微调，无需从头训练解码器）。对于缓解耦合优化问题，一种方法是在优化图像编码器时，直接将掩码解码器冻结，这样以来就可以防止优化图像解码器时对掩码解码器产生影响，我们把这种蒸馏方式称为半耦合蒸馏（distillation semi-coupled）。但由于提示符的是随机的，所以不同的提示符会对掩码解码器产生不同的影响，这会给图像解码器的优化带来新的问题。因此，我们直接将原始SAM中图像解码器的知识蒸馏到一个更小的图像编码器中，这个过程不需要任何解码器的参与，我们把这种蒸馏方式称为解耦蒸馏。直接在图像嵌入上进行蒸馏的好处是：在进行图像掩码预测时，可以使用简单的均方差损失（MSE）代替focal损失和dice损失。

On the necessity of mask decoder finetuning

        得到的学生解码器和教师解码器很相似，它与原来的掩码解码器能很好的适配，因此不需要对掩码解码器进行微调，当然，对其进行微调也有可能进一步改善Mobile SAM的性能。

Preliminary evaluation（初步评估）

        这个部分我们对耦合蒸馏和解耦蒸馏进行了比较。使用相同的提示符，我们比较了教师SAM和学生SAM的掩码解码器生成的掩码之间的mIoU。具体结果如表2所示：

4、实验

4.1实验装置

        Lightweight Image Encoder：本文使用的baseline是ViT-Tiny，采用的轻量化图片编码器由四个部分组成，每个部分逐渐减小图片分辨率。第一个部分由带有倒置残差（inverted residual）的卷积块组成，剩下的三部分都由transformer模块组成。模型的前端由两个步幅为2的卷积块进行下采样分辨率。不同部分之间的下采样操作也是由步幅为2的卷积块完成。与Wu et al[2022]不同，为了让最终的分辨率和SAM中的ViT-H图片编码器匹配，我们将最后一个下采样卷积核的步幅为设为1。

        Training and evaluation details：我们在一个RTX3090的单核GPU上使用很少的数据训练了一个轻量级编码器，通过观察我们发现很多计算开销都是在教师图像编码器的前向处理过程中，为了加快蒸馏速度，我们提前保存了图片嵌入（image embedding），这样就只需要进行一次前向处理操作。当然，用更多的GPU训练更长的时间、对掩码编码器进行微调等都有可能进一步提高Mobile SAM的性能，但在这篇文章中我们没有这样做。为了对蒸馏SAM进行定量比较，我们计算了SAM和Mobile SAM预测得到的掩码之间的mIoU。

4.2MobileSAM performs on par with the orignal SAM

        Ablation study：通过消融实验，我们发现通过增加batch size和epoch都可以增加模型的性能，因此我们也期待使用更多的GPU来增加batch size和迭代次数，以进一步改善模型的性能。

4.3 MobileSAM outperforms FastSAM in All Aspects

        Segment anything v.s. segment everything：segment anything旨在利用提示词分割出任何想要的物体，而segment everything旨在生成任何物体的掩码，本文主要在segment anything方面集中比较MobileSAM和FastSAM的泛化性和通用性（versatility）。

        MobileSAM is faster and smaller：

        mIoU comparison under segment anything mode：

        Results for segment everything：

5、结论

结论和摘要差不多。