Faster Segment Anything: Towards Lightweight SAM for Mobile Applications
SAM代码:https://github.com/ChaoningZhang/MobileSAM
SAM论文:https://arxiv.org/pdf/2306.14289.pdf
1 概述
Faster SAM的目标是通过用轻量级图像编码器取代笨重的图像编码器,使得SAM适应移动设备。然而,按照原始SAM论文中的方式训练这样一个新的SAM会导致不令人满意的性能,尤其是在训练数据有限的情况下。我们发现,这主要是由于图像编码器和掩膜解码器的耦合优化所导致的。因此,为了解决这个问题,我们提出了一种名为解耦蒸馏的方法。具体而言,我们将原始SAM中的图像编码器ViT-H的知识蒸馏到一个轻量级图像编码器中,这个编码器可以自动与原始SAM中的掩膜解码器兼容。通过这种方式,我们能够在不到一天的时间内在单个GPU上完成训练,得到的轻量级SAM被称为MobileSAM。尽管MobileSAM的体积是原始SAM的60倍以上,但它的性能与原始SAM相当。
SAM由两个组件组成:
- ViT-based图像编码器
- Prompt-guided Mask解码器
它们按顺序工作
2 Mobile-Friendly SAM
2-1 Coupled distillation
用较小的图像编码器重新训练新的SAM。如SAM所述,使用ViT-H图像编码器训练SAM需要在256个A100 GPU上花费68小时。用ViT-L或ViT-B取代ViT-H将所需的GPU减少到128个,但是这仍然是一个不小的计算开销。
当然可以进一步采用更小的图像编码器,并用他们提供的11-T分割数据集(Mask是由预训练的SAM给出的)重新训练新的SAM。本质上,这种再训练过程是知识蒸馏,它将知识从基于ViT-H的SAM转移到具有较小图像编码器的SAM
2-2 From semi-coupled to decoupled distillation
在从原始的SAM到具有较小图像编码器的SAM进行知识蒸馏时,困难主要在于图像编码器和组合解码器的耦合优化。直观地说,图像编码器的优化取决于图像解码器的质量,反之亦然。受分治算法的启发,将知识蒸馏任务分为两个子任务:图像编码器蒸馏和掩膜解码器微调。
**图像编码器蒸馏:**冻结Mask解码器,从原始SAM中的ViT-H中蒸馏小图像编码器,而不使用组合解码器;
**掩膜解码器微调:**在冻结的轻量化图像编码器上对Mask解码器进行微调或将它们联合微调可能会进一步提高性能(从学生图像编码器生成的图像编码可以足够接近原始教师编码器的图像编码,这使得本阶段的微调是可选的);
此外,通过对图像嵌入进行蒸馏,我们可以采用简单的MSE loss,而不是使用focal loss和dice loss的组合来进行掩膜预测,这进一步简化了蒸馏过程。
3 Experiments
与FastSAM比较