【论文阅读笔记】Faster Segment Anything: Towards Lightweight SAM for Mobile Applications

最新推荐文章于 2024-09-27 19:06:31 发布

咔叽布吉

最新推荐文章于 2024-09-27 19:06:31 发布

阅读量708

点赞数 29

分类专栏：论文阅读学习文章标签：论文阅读笔记

本文链接：https://blog.csdn.net/qq_46056318/article/details/142492966

版权

论文阅读学习专栏收录该内容

34 篇文章 2 订阅

订阅专栏

1.论文介绍

Faster Segment Anything: Towards Lightweight SAM for Mobile Applications
更快的细分市场：面向移动的应用的轻量级SAM
2023年arxiv
Paper Code

2.摘要

Segment Anything Model（SAM）因其令人印象深刻的zero-shot传输性能和众多视觉应用（如具有细粒度控制的图像编辑）的高通用性而引起了广泛关注。许多这样的应用程序需要在资源受限的边缘设备上运行，比如移动设备。在这项工作中，我们的目标是使SAM移动设备友好，通过把重量级图像编码器替换为轻量级的。一训练这样一个新的SAM会导致不满意的性能，特别是当有限的训练资源。我们发现，这主要是由于图像编码器和掩码解码器的耦合优化，我们提出了解耦蒸馏的动机。具体地说，我们从重型图像编码器（原始SAM中的ViT-H）中提取的知识，以轻量级的图像编码器，它可以自动兼容原始SAM中的掩码解码器。训练可以在不到一天的时间内在单个GPU上完成，由此产生的轻量级SAM被称为MobileSAM，它比原来的SAM小60倍以上，但性能却与原来的SAM相当。对于推理速度，使用单个GPU，MobileSAM每幅图像运行约10 ms：图像编码器上为8 ms，掩码解码器上为4 ms。凭借上级性能，我们的MobileSAM比并发FastSAM快5倍左右，体积小7倍，更适合移动的应用。此外，我们表明，MobileSAM可以在CPU上运行相对平稳。

Keywords：解耦蒸馏、轻量级编码器、Faster SAM

3.Introduction

SAM由两个组件组成：基于ViT的图像编码器和图像引导掩码解码器，它们按顺序工作（也可以说是三个：Image Encoder、Prompt Encoder、Mask Decoder）。自其出现以来，SAM因多种原因引起了极大的关注。首先，这是第一次表明视觉可以遵循NLP，追求一条将基础模型与快速工程相结合的道路。其次，它是第一个执行无标签分割的方法，这是一项与标签预测并行的基本视觉任务。此外，这一基本任务使SAM与其他模型兼容，以实现高级视觉应用，例如文本引导的分割和具有细粒度控制的图像编辑。然而，许多此类用例需要在资源受限的边缘设备上运行，例如移动的应用程序。如官方演示所示，通过处理图像嵌入，SAM可以在资源受限的设备上工作，因为掩码解码器是轻量级的。导致SAM流水线计算繁重的原因在于庞大的图像编码器。在这项工作中，我们研究如何获得一个轻量级的SAM适合资源受限的移动的设备，因此被称为MobileSAM。

鉴于SAM中的默认图像编码器是基于VIT-H的，获得MobileSAM的直接方法是遵循SAM的流程，用更小的图像编码器重新训练新的SAM，如用更小的VIT-L或甚至更小的VIT-B取代VIT-H。用VIT-L或VIT-B作为图像编码器训练一个新的SAM需要多天128个GPU。这种资源密集的再培训对于复制或改进其结果可能是一个不小的负担。这一优化难点主要来自于图像编码器和掩模解码器的耦合优化。基于这种理解，我们提出了将图像编码器和掩码解码器的优化解耦的建议。具体地说，我们首先将默认图像编码器VIT-H的知识提取到一个微小的VIT中。之后，我们可以微调原始SAM中的掩码解码器，以更好地与提取的图像编码器对齐。值得强调的是，对齐优化是可选的，因为轻量级图像编码器是从默认图像编码器提取的，这一事实确保了其与默认掩码解码器的固有对准。通过将寻找新的SAM流水线的问题转化为解耦精馏，我们的方法具有简单、有效、可重复使用的优势，同时成本较低(在单个GPU上不到一天)。由此产生的MobileSAM将编码器参数减少了100倍，而总参数减少了60倍。令人惊讶的是，如此轻量级的MobileSAM的表现与最初的重量级SAM不相上下，这构成了将SAM推向移动应用的重要一步。对于MobileSAM的推断，单个图像仅运行10ms左右：在图像编码器上运行8ms，在掩码解码器上运行4ms。值得一提的是，我们的MobileSAM比并发的FastSAM快约5倍，小约7倍。同时取得优异的业绩。

4.网络结构详解

一些尝试：
在这里插入图片描述

耦合蒸馏Coupled distillation。如上图左边，直接知识蒸馏，把SAM的Image Encoder替换成小的encoder，然后用他们提供的分割数据集11-T重新训练新的SAM。这里的GT是由预先训练的SAM(ViT图像编码器)给出的，所以等同于蒸馏。

从半耦合到解耦蒸馏。 如上图右边，当执行从原始SAM到具有较小图像编码器的SAM的知识蒸馏时，困难主要在于图像编码器和组合解码器的耦合优化。本文将KD任务分为两个子任务：图像编码器提取和掩码解码器微调。具体来说，首先通过将知识从ViT-H转移到较小的编码器来对图像编码器执行KD。为了缓解耦合蒸馏的优化问题，使用复制和冻结的掩码解码器。称这种提取为半耦合，因为图像编码器的优化仍然没有完全与掩码解码器解耦。
优化仍然具有挑战性，因为提示的选择是随机的，这使得掩码解码器可变，从而增加了优化难度。
在这里插入图片描述
预期微调。冻结的轻量级图像编码器上的掩码解码器或将它们一起联合微调可能进一步提高性能。直接从原始SAM中的ViT-H中提取小型图像编码器，而不求助于组合解码器，这被称为解耦提取（如上图）。对图像嵌入进行蒸馏的另一个优点是，我们可以采用简单的MSE损失，而不是像SAM那样使用焦点损失和骰子损失的组合来进行掩模预测。

使用用ViT-B验证蒸馏方法，依旧困难。所以采用ViT-Tiny进行概念验证，以证明提出的解耦蒸馏的有效性，用于训练比原始SAM快得多的轻量级MobileSAM。所采用的轻量级图像编码器包括四个阶段，逐步降低分辨率。第一阶段由具有反转残差的卷积块构成，而其余三个阶段由Transformer块组成。在模型的开始，有2个步长为2的卷积块用于对分辨率进行下采样。不同阶段之间的下采样操作由步长为2的卷积块处理。将最后一次下采样卷积中的步幅2设置为1，以使最终分辨率与原始SAM的ViT-H图像编码器的分辨率相匹配。