0. 简介

自从MetaAI提出的能够“分割一切”的视觉基础大模型SAM提供了很好的分割效果，为探索视觉大模型提供了一个新的方向。虽然SAM的效果很好，但由于SAM的backbone使用了ViT，导致推理时显存的占用较多，推理速度偏慢，对硬件的要求较高，在项目应用上有很大的限制。FastSAM通过更换检测头，能够较好的完成分割的实时性，但是将SAM的”重量级”解码器替换为”轻量级”以使其可在移动端部署应用仍然还是不现实的，所以《Faster Segment Anything: Towards Lightweight SAM for Mobile Applications》提出一种”解耦蒸馏”方案对SAM的ViT-H解码器进行蒸馏，同时所得轻量级编码器可与SAM的解码器”无缝兼容” 。相关的代码也已经在Github上开源了


图1：Segment Anything模型的概览

1. 主要贡献

考虑到SAM（Space-Aware Multi-Head Attention）中默认的图像编码器是基于ViT-H（Vision Transformer-H）的，获得MobileSAM的一个直接方法就是按照Kirillov等人[2023]的官方流程，用一个更小的图像编码器如ViT-L（Vision Transformer-L）或更小的ViT-B（Vision Transformer-B）来重新训练一个新的SAM。

  1. 延续SAM架构体系：采用轻量级ViT解码器生成隐特征，然后采用提示词引导解码器生成期望的Mask。
  2. 本文提出一种”解耦蒸馏”方案对SAM的ViT-H解码器进行蒸馏，同时所得轻量级编码器可与SAM的解码器”无缝兼容” 。
  3. 我们的方法具有简单、有效，同时可以以低成本（在单个GPU上少于一天的时间）进行再现的优点。结果得到的MobileSAM将编码器参数减少了100倍，总参数减少了60倍。在推理速度方面，MobileSAM处理一张图像仅需10ms(8ms@Encoder，2ms@Decoder)，比FastSAM的处理速度快4倍，这就使得MobileSAM非常适合于移动应用。

2. Mobile SAM背景和项目目标

2.1 SAM（Segment Anything Model，即“切割任何东西”的模型）背景

首先，我们总结一下SAM的结构和工作方式。SAM由基于ViT（Vision Transformer）的图像编码器和提示引导的掩模解码器组成。图像编码器将图像作为输入并生成嵌入，然后将其输入到掩模解码器中。掩模解码器根据提示（如一个点或框）生成一个掩模，用以从背景中切割出任何对象。此外，SAM允许为同一提示生成多个掩模，以解决歧义问题，这提供了宝贵的灵活性。考虑到这一点，我们保留了SAM的流程，首先采用基于ViT的编码器生成图像嵌入，然后采用提示引导的解码器生成所需的掩模。这个流程最适合“切割任何东西”，可用于“切割所有东西”这一下游任务

2.2 项目目标

该项目的目标是生成一个在轻量级方式下实现令人满意性能，并且比原始SAM更快的面向移动端的SAM（MobileSAM）。原始SAM中的提示引导掩模解码器的参数少于4M，因此被认为是轻量级的。如他们的公开演示所示，只要有编码器处理的图像嵌入，SAM就可以在资源受限的设备上运行，因为掩模解码器是轻量级的。然而，原始SAM中的默认图像编码器基于ViT-H，参数超过600M，这非常重，使得整个SAM流程与移动设备不兼容。因此，获得面向移动端的SAM的关键在于用轻量级的图像编码器替换重型的图像编码器，这也自动保留了原始SAM的所有功能和特性。接下来，我们将详述我们提出的方法，以实现这个项目目标。

3. Mobile SAM所提出的方法

3.1 耦合提炼

实现我们项目目标的一个直接方法是按照Kirillov等人[2023]的官方流程，用一个更小的图像编码器重新训练一个新的SAM。如Kirillov等人[2023]所述，使用ViT-H图像编码器训练SAM需要在256个A100 GPUs上运行68小时。用ViT-L或ViT-B替换ViT-H会将所需的GPU减少到128个，但对于社区中的许多研究人员来说，复制或改进他们的结果仍然是一个不小的负担。按照他们的方法，我们可以采用一个更小的图像编码器，并使用他们提供的分割数据集（11-T）重新训练一个新的SAM。需要注意的是，数据集中提供的掩模由预训练的SAM（带有ViT图像编码器）给出。从本质上讲，这个再训练过程是知识提炼（Hinton等人，2015），它将知识从基于ViT-H的SAM传递到一个带有更小图像编码器的SAM（见图2左）。


图2：SAM的耦合知识提炼。左图表示完全耦合的提炼，而右图代表半耦合的提炼。

3.2 从半耦合到解耦提炼

当从原始的SAM执行KD（知识提炼）到带有更小图像编码器的SAM时，主要困难在于图像编码器和组合解码器的耦合优化。直观地说，图像编码器的优化依赖于图像解码器的质量，反之亦然。当SAM中的两个模块都处于不良状态时，训练它们都达到良好状态就更具挑战性。受到分而治之算法（Zhang等人，2022c）的启发，我们提出将KD任务分为两个子任务：图像编码器提炼和掩模解码器微调。具体来说，我们首先在图像编码器上执行KD，将知识从ViT-H转移到一个更小的编码器。由于原始SAM中的掩模解码器已经是轻量级的，我们计划保留其架构。这带来了一个好处，即可以直接使用现成的组合解码器进行微调，而不是从头开始训练。为了缓解耦合提炼的优化问题，一个直接的方法是使用复制并冻结的掩模解码器优化图像编码器（见图2右）。冻结操作可以帮助防止掩模解码器的质量被差的图像编码器恶化。我们称这种提炼为半耦合，因为图像编码器的优化仍然没有完全从掩模解码器解耦。实际上，我们发现这种优化仍然具有挑战性，因为提示的选择是随机的，这使得掩模解码器变得可变，从而增加了优化的难度。因此，我们提议直接从原始SAM中的ViT-H中提炼小图像编码器，而不依赖于组合解码器，这被称为解耦提炼（见图3）。在图像嵌入上执行提炼的另一个优点是我们可以采用简单的MSE（Mean Squared Error）损失，而不是像Kirillov等人[2023]那样，为制作掩模预测而使用焦点损失（Lin等人，2017）和dice损失（Milletari等人，2016）的组合。


图3：用于SAM的解耦知识提炼

3.3 关于微调蒙版解码器的必要性

不同于半耦合蒸馏，经解耦合蒸馏训练得到的轻量级Encoder可能与冻结的Decoder存在不对齐问题。根据经验，我们发现：该现象并不存在。这是因为学生Encoder生成的隐特征非常接近于原始老师Encoder生成的隐特征，因此并不需要与Decoder进行组合微调。当然，进一步的组合微调可能有助于进一步提升性能。

3.4 初步评估

在这里，我们进行了初步的研究，比较了耦合提炼和解耦提炼。在这里，为了性能评估，我们计算了教师SAM和学生SAM在同一提示点上生成的两个蒙版之间的mIoU。直观地说，更高的mIoU表示了更高的蒙版预测性能，假设由ViT-H生成的蒙版是地面真实值。对于耦合提炼，我们采用的是原始SAM Kirillov等人[2023]提供的配有ViT-B的SAM。它在128个GPU上（每个GPU1个样本）对SA-1B（1100万图片）进行了18万次迭代训练。相比之下，在我们的解耦提炼设置中，我们在2个GPU上（每个GPU两个样本以节省计算资源）对SA-1B数据集的0.1%样本（11k张图片）进行了55k次迭代训练。总的来说，解耦提炼比耦合提炼需要的计算资源少1%，而且在mIoU上的性能优于耦合提炼（0.75 vs 0.72，基于200个样本的平均值）。由于ViT-B仍然对移动设备是一个不小的负担，因此在以下的实验中，我们基于我们提出的解耦提炼使用了TinyViT（参数为500万）Wu等人[2022]。


表2：使用ViT-B作为图像编码器的SAM的耦合提炼和解耦提炼的比较。解耦提炼表现更好，所需的计算资源不到耦合提炼的1%。


4. 参考链接

https://mp.weixin.qq.com