SAM轻量级改进Mobile-SAM论文解读FASTER SEGMENT ANYTHING: TOWARDS LIGHTWEIGHT SAMFOR MOBILE APPLICATIONS

现已总结SAM多方面相关的论文解读，具体请参考该专栏的置顶目录篇

一、总结

1. 简介

发表时间：2023年7月1日

论文：[2306.14289] Faster Segment Anything: Towards Lightweight SAM for Mobile Applications (arxiv.org)https://arxiv.org/abs/2306.14289

代码：ChaoningZhang/MobileSAM: This is the official code for MobileSAM project that makes SAM lightweight for mobile applications and beyond! (github.com)https://github.com/ChaoningZhang/MobileSAM

2. 摘要 

       SAM在可用的训练资源有限的情况下性能不能令人满意，这主要是由图像编码器和掩膜解码器的耦合优化引起的，为此MobileSAM提出了解耦蒸馏。具体来说，是将重型图像编码器的知识提取为轻型图像编码器来使SAM适合移动设备，该轻型编码器可与原SAM中的掩码解码器自动兼容。由于我们的MobileSAM保留了原始SAM的所有管道，只是替换了图像编码器，因此它可以即插即用地为现有的基于SAM的项目从重量级SAM转变为轻量级SAM。
        MobileSAM可以在不到一天的时间内在单个GPU上完成训练，比原来的SAM编码器参数减少了100倍，总参数减少了60倍，但性能与原来的SAM相当。MobileSAM推理速度：使用单个GPU, 每张图像运行约10ms，其中图像编码器运行8ms，掩码解码器运行4ms。还将MobileSAM与同期FastSAM进行了比较，结果表明MobileSAM比FastSAM快5倍左右，小7倍，使其更适合移动应用。此外，MobileSAM还可以在CPU上相对平稳地运行。

3. 前言

        SAM工作流程如下图所示，SAM由两个组件组成：基于vit的图像编码器和提示引导的掩码解码器，图像编码器将图像作为输入并生成嵌入，然后将其馈送到提示引导的掩码解码器。掩码解码器生成一个掩码，根据点(或框)等提示从背景中剪切任何对象。此外，SAM允许为解决歧义问题的同一提示生成多个掩码​​​​​​。
        SAM图像编码器参数超过600M是重量级的，而掩码解码器参数小于4M是轻量级的，使得整个SAM流水线与移动设备不兼容，仅可以在资源受限的设备上工作。因此，获得适合移动设备的SAM的关键在于将重型图像编码器替换为轻量级图像编码器，并自动保留原SAM的所有功能和特性。

         下表总结了不同图像编码器变体的SAM参数量，用ViT-L或ViT-B作为图像编码器训练一个新的SAM需要128个gpu，持续数天。

二、Mobile SAM模型结构

        为了可在资源受限的边缘设备上运行，由此推出MobileSAM。鉴于SAM中的默认图像编码器基于ViT-H，MobileSAM采用解耦蒸馏的方式，用更小的图像编码器重新训练新的SAM，例如用更小的ViT-L甚至更小的ViT-B替换ViT-H。想要了解知识蒸馏可以参考下面我写的这篇文章。Mobile-SAM使用的知识蒸馏方法论文讲解Distilling the Knowledge in a Neural Network-CSDN博客https://blog.csdn.net/sunshineine/article/details/135815511?spm=1001.2014.3001.5501

1. 耦合蒸馏

        使用ViT-H图像编码器训练SAM需要在256个A100 gpu上花费68小时。用更小的图像编码器ViT-L或ViT-B取代ViT-H可以将所需的gpu减少到128个，重现或改进该结果仍然是一个不小的负担。按照SAM的方法，我们可以进一步采用更小的图像编码器，并使用他们提供的分割数据集重新训练新的SAM。请注意，所提供数据集中的掩码是由预训练的SAM(使用ViT图像编码器)给出的。从本质上讲，这种再训练过程是知识蒸馏，它将知识从基于ViT-H的SAM转移到具有较小图像编码器的SAM。
        如下图所示，为SAM的耦合知识蒸馏，左子图为全耦合蒸馏，右子图为半耦合蒸馏。半耦合蒸馏是使用复制和冻结掩膜解码器来优化图像编码器，冻结操作可以帮助防止掩码解码器的质量因较差的图像编码器而恶化，但图像编码器的优化仍然没有与掩模解码器完全解耦。

2. 解耦蒸馏

        这种资源密集的再培训可能是复制或改进其结果的一个不小的负担，这种优化难点主要来自于图像编码器和掩码解码器的耦合优化。基于这种理解，提出对图像编码器和掩码解码器进行解耦优化。具体来说，首先从默认图像编码器ViT-H提取知识到一个微小的ViT。之后，我们可以对原始SAM中的掩码解码器进行微调，以更好地与蒸馏后的图像编码器对齐。值得强调的是，对齐优化是可选的，因为轻量级图像编码器是从默认图像编码器中提取出来的，这保证了它与默认掩码解码器的固有对齐。

3. 从半耦合蒸馏到解耦蒸馏

        当使用较小的图像编码器对原始SAM进行KD转换时，难点主要在于图像编码器和组合解码器的耦合优化。直观地说，图像编码器的优化取决于图像解码器的质量，反之亦然。当SAM中的两个模块都处于不良状态时，将它们训练到良好状态更具挑战性。因此提出将KD任务分为两个子任务:图像编码器蒸馏和掩膜解码器微调。具体来说，我们首先通过将ViT-H的知识转移到更小的编码器上，对图像编码器进行KD。由于原始SAM中的掩码解码器已经是轻量级的，我们计划保留掩码解码器的架构。这带来了一个方便使用的组合解码器微调的好处，而不是从头开始训练它。耦合蒸馏的优化仍然具有挑战性，因为提示符的选择是随机的，这使得掩码解码器可变，从而增加了优化难度。因此，我们建议直接从原始SAM中的ViT-H提取小图像编码器，而不使用组合解码器，这被称为解耦蒸馏(如下图所示)。对图像嵌入进行蒸馏，因此可以采用简单的MSE损失，而不是使用焦损失或组合来进行mask预测。

 4. 论掩码解码器微调的必要性

        与半耦合蒸馏不同，上述解耦蒸馏产生轻量级图像编码器，也可以与原始冻结掩膜解码器很好地对齐。因为从学生编码器生成的图像编码可以足够接近原始教师编码器的编码，这使得第二阶段对组合解码器的微调是可选的。  期望在冻结的轻量级图像编码器上微调掩码解码器或将它们共同微调可以进一步提高性能。

5. 初步评估

        为了进行性能评估，对耦合蒸馏和解耦蒸馏进行了初步的比较研究。计算了教师SAM和学生SAM在同一提示点上生成的两个掩码之间的mIoU。直观上，假设ViT-H生成的掩码为基真值，mIoU越高，掩码预测性能越好。
        如下表所示，为以ViT-B作为图像编码器的SAM的耦合蒸馏和解耦蒸馏的比较。对于耦合蒸馏，在128个GPU(每个GPU 1个样本)上对SA-1B (11M张图像)进行了180k次迭代训练；对于解耦蒸馏，我们在2个GPU(每个GPU两个样本以节省计算资源)上对SA-1B数据集(11k)图像的0.1%样本进行了55k次迭代。
        总体而言，与耦合蒸馏相比，解耦蒸馏所需的计算资源不到1%，而耦合蒸馏的mIoU性能为0.75 vs 0.72(平均200个样本)。由于ViT-B对于移动设备来说仍然是一个不容忽视的负担，因此在接下来的实验中，我们基于我们提出的解耦蒸馏，使用一个TinyViT(具有5M参数)进行实验。   

三、实验流程

1. 实验设置

1.1 轻量级图像编码器

        该项目的目标是通过使用轻量级图像编码器替换默认的ViT-H来获得高效的SAM。作为基于ViT的骨干网络，ViT-Tiny的参数与Deit-Tiny相似，但性能更好。因此，我们采用ViT-Tiny进行概念验证。
        轻量化图像编码器由四个阶段组成，逐步降低分辨率。第一阶段由具有倒残差的卷积块构建，其余三个阶段由变压器块组成。在模型开始时，有2个卷积块，步幅为2，用于下采样分辨率。各阶段间的降采样操作采用步长为2的卷积块进行处理。将最后一次下采样卷积的步长2设置为1，以使最终分辨率与原始SAM的ViT-H图像编码器的分辨率相匹配。
        下表为原始SAM和MobileSAM图像编码器参数和在单个GPU上测量的推理速度的比较。

1.2 培训和评估细节

        对于图像编码器上的解耦KD，我们使用1%的SA-1B数据集在单个GPU上训练轻量级编码器8个epoch。我们观察到教师图像编码器的前向过程花费了更多的计算，因为它比我们采用的学生图像编码器要重得多。
        为了使蒸馏更快，预先保存图像嵌入，这样我们只需要运行一次前向过程。使用单个GPU，我们可以在不到一天的时间内获得我们的MobileSAM。用更多的GPU长时间训练我们的MobileSAM有望产生更好的性能。
        为了对提取的SAM进行定量评价，我们计算了原始SAM预测的掩码与我们的MobileSAM之间的mIoU。

2. SAM与MobileSAM的性能

2.1 SAM消融实验

        如下表所示，在单个GPU上训练训练SAM性能，在相同的迭代次数下，增加批大小可以提高模型的性能。此外，在批处理规模下，通过增加训练epoch，更多的更新迭代也会使性能受益。

2.2 MobileSAM性能

        使用两种类型的提示来报告预测掩码：点提示和框提示。点提示结果如图4所示，框提示结果如图5所示。实验表明，MobileSAM做出了与原始SAM相似的令人满意的掩模预测。

         

3. MobileSAM优于FastSAM

3.1 segment anything VS segment everything

      “segment anything”执行提示式分割任务，根据任何分割提示返回有效的分割掩码，提示符的作用是指定在图像中分割什么，从理论上讲，只要设置好提示符，任何对象都可以被分割，因此称为“分割任何对象”。而“segment everything”本质上是对象候选框生成，不需要提示，选择“segment everything”作为下游任务可以展示其零射击传输性能。综上所述，“segment anything”解决的是对任意对象进行及时分割的基础任务，而“segment everything”解决的是对所有对象生成掩码提议的下游任务。
        因为segment everything并不一定需要提示，FastSAM直接用YOLO v8以一种不需要提示的方式生成掩码建议。为了实现即时分割，设计了映射算法从建议掩码集中选择掩码。与FastSAM的比较主要集中在“segment anything”上，但为了完整性，也提供了关于“segment everything”的比较。如下表所示，为“segment anything”与“segment everything”的比较。

3.2 MobileSAM更快、更小

        FastSAM由一个基于yolov8的检测分支和一个基于yolact的分割分支组成，用于生成无提示的掩码建议。它有68M个参数，处理一张图像需要40毫秒。相比之下，MobileSAM的10M参数更少，明显更小。对于推理速度，在单个GPU上，FastSAM处理图像需要40ms，而MobileSAM只需要10ms，比FastSAM快4倍。具体如下表所示。

3.3 基于SAM为ground-truth的mIoU比较

        进一步比较了预测掩模与原始SAM的mIoU，假设原始SAM的预测掩模为真实值，mIoU越高性能越好。FastSAM建议用多个点来预测蒙版，选择了一个点作为前景，另一个点作为背景。 如下表所示，FastSAM的mIoU比MobileSAM的mIoU小得多，这表明FastSAM的掩码预测与原始SAM的掩码预测有很大的不同。此外，当两个提示点之间的距离增大时，FastSAM的mIoU减小得非常快。主要是由于当前景提示点设置得太接近背景提示点时，FastSAM往往无法预测目标。 

3.4 分割一切的结果

        “分割一切”的结果如下图所示。首先，我们的MobileSAM的结果与原始SAM的结果惊人地一致，而FastSAM的结果往往不那么令人满意。
（1）FastSAM经常无法预测某些物体，比如第一张图像中的屋顶。
（2）FastSAM掩码候选框有时很难解释，比如第一张图片中舞台的掩码和第二张图片中天空的掩码。
（3）FastSAM经常生成具有非光滑边界的掩码，比如第三张图片中的柱子有不光滑的边界。