【论文阅读笔记】Lite-SAM Is Actually What You Need for Segment Everything

1.论文介绍

Lite-SAM Is Actually What You Need for Segment Everything
Lite-SAM是您实际上所需的分割一切的工具
2024年 arxiv
Paper

2.摘要

Segment Anything模型（SAM）以其优越的性能给分割领域带来了重大变化，但其对计算资源的巨大需求仍然是一个限制因素。 MobileSAM、Edge-SAM、MobileSAM-v2等许多作品都探索了轻量级解决方案。然而，他们使用传统的网格搜索采样策略或两阶段串联方法，不允许端到端训练，严重限制了分段一切（SegEvery）的性能。本文介绍了 Lite-SAM，这是一种针对 SegEvery 任务的高效端到端解决方案，旨在降低计算成本和冗余。 Lite-SAM由四个主要组件组成：精简的CNNTransformer混合编码器（LiteViT）、自动提示提议网络（AutoPPN）、传统提示编码器和掩码解码器。所有这些组件都集成在 SAM 框架内。我们的LiteViT是一个高性能的轻量级骨干网络，参数只有1.16M，比现有最轻的骨干网络Shufflenet减少了23%。我们还推出了 AutoPPN，这是一种用于提示框和积分生成的创新端到端方法。这是对传统网格搜索采样方法的改进，其独特的设计可以轻松集成到任何 SAM 系列算法中，从而扩展了其可用性。我们在大量公共和私人数据集上对 Lite-SAM 进行了彻底的基准测试。评估涵盖了广泛的通用指标，包括参数数量、SegEvery 执行时间和准确性。研究结果表明，Lite-SAM 以精简的 4.2M 参数运行，显着优于同类产品，与 SAM、MobileSAM、Edge-SAM、EfficientViT-SAM 和 EfficientViT-SAM 相比，性能提高了 43 倍、31 倍、20 倍、21 倍和 1.6 倍 MobileSAM-v2 分别始终保持有竞争力的准确性。这凸显了 Lite-SAM 在实现性能和精度之间的最佳平衡方面的能力，从而在该领域树立了新的最先进 (SOTA) 基准。

Keywords: SAM，轻量级SAM，SegEvery

3.Introduction

SAM 代表了视觉框架的重大进步，与 GPT 在 NLP 中的影响类似。它包含两个关键组件：基于 ViT 的图像编码器和提示引导掩模解码器，它们协同工作。 SAM 旨在处理两个分段任务：SegAny 和 SegEvery。这两项任务都涉及与类别无关的掩模分割，但目标不同。 SegAny 使用特定提示（例如点或框）来隔离和分割图像中感兴趣的特定项目。相比之下，SegEvery 的目标是描绘图像中所有可辨别的主题。

张超宁等提出了一种“解耦蒸馏”，旨在蒸馏 SAM 的 ViT-H 解码器，从而产生更高效的轻量级编码器，可以与 SAM 解码器集成。然而，该算法模型缺乏平台适配的鲁棒性，并且在翻译过程中表现出相当大的准确性损失，使其不太适合在移动设备上部署。赵等人介绍了基于 YOLOv8 [15] 构建的 Fast-SAM 模型，该模型展示了卓越的分割能力。然而，它的主要局限性是缺乏全方位的交互方式，特别是缺乏专用的框和点功能。李等人设计了 SemanticSAM，这是一种增强不同尺度图像的分割和识别多功能性的模型。但必须强调的是，其大量参数导致推理时间更长。韩才等人提出了 EfficientViT，引入了一种新颖的轻量级算法，该算法取得了有希望的结果。周崇等人。提出了 Edge-SAM ，这是一种在 iPhone 上实现 SegAny 任务实时执行的算法。所有上述方法[4,16,18,45,49,56,57]都针对SegAny进行了评估；然而，SegEvery 仍然非常耗时。

MobileSAM-v2 为 YOLOv8提出了一种创新的训练方法，使用预先生成的提示（对象感知提示采样）代替传统的 Gridsearch 采样策略，提高了 SegEvery 过程的效率。尽管有这样的改进，但这种方法需要使用单独的模型，这被认为是权宜之计。由于YOLOv8固有的推理和训练需求，总体节省的时间可能有限。为了解决上述问题，我们的贡献可总结如下：

我们推出了 LiteViT，这是一种轻量级 CNN-Transformer 编码器，通过减少参数来提高准确性，非常适合有限的计算环境。
开发 AutoPPN，一个自动提示提案网络，提高了网格搜索方法的效率，并与 SAM 系列算法轻松集成。
通过实验验证了 Lite-SAM 的性能，如图 1 所示，显示了 SegEvery 上的加速结果，同时保持了准确性。

4.网络结构详解

在这里插入图片描述
Lite-SAM架构，该架构由四个主要组件组成：LiteViT编码器、AutoPPN网络、标准提示编码器和掩码解码器。该配置如图（a）所示。新的AutoPPN模块专门设计用于简化自动化提示任务。该算法以端到端的方式同时对边界框提示和点提示进行回归，与已有的密集位置编码方案相比，显著减少了SegEvery任务的推理时间。这一进步是实现实时分割的关键。

Lite-SAM架构
标准的自注意力token混合器以其高计算成本而闻名。相比之下，组合CNN-Transformer混合结构在制作轻量级Vision Transformer（ViT）网络中发挥着重要作用。这种混合平衡了模型精度和计算效率。受现有研究中自注意力层的有效变化的启发，作者开发了LiteViT图像编码器，从PoolFormer-S12 基线开始。将图像编码器模型建立在一个称为LiteViT Block的新型构建模块上。详细的体系结构规范见图（b）。为了克服捕捉局部特征的挑战，本文将多尺度池化整合到了轻量级注意力模块中。具体地说，引入了多尺度池化模块（MSPM），以有效地增强网络架构的每个阶段的感受野。在LiteViT模块中，输入首先由MSPM模块处理，然后由卷积MLP（多层感知器）模块处理；每一级都通过跳跃连接进行连接。为了便于下采样和调整各级的输出通道，采用了一个称为补丁合并模块的专用模块，该模块可有效地充当主干卷积层。
在这里插入图片描述

AutoPPN：
对提示使用密集位置编码的标准方法可能不适合实时分割任务。为了增强SegEvery任务的推理性能，引入了AutoPPN模块，该架构在图（c）中详细描述。用位于边界框中心的单个点表示对象是一种简单有效的技术。基于这一概念，AutoPPN框架以端到端的方式从输出特征图中预测提示点和边界框。相应的损失由两个元素组成：点提示的置信度和边界框回归的准确性。我们已实施三项重大修改以完善我们的方法，详情如下：
（1）通过用更复杂的主干MSPM网络取代基本的主干卷积网络来增强我们的网络。这种更新的网络有效地整合了多尺度空间信息，大大提高了对天空、建筑物和水体等大型物体或实体的检测召回率。
（2）为了估计点提示的置信度，我们结合了距离变换的使用，以便于计算点与其对应掩模之间的距离。如果一个点位于多个掩模内，我们会选择面积最小的掩模。不同于基于高斯的方法，我们的技术能够在前景和背景之间创建柔和的逐点真值区分，而非生硬的二元分割。此外，本文的方法优先识别对象或实体的最中心点，而不是简单地关注其边界框的中心。这种修改有助于减轻在涉及不清楚的点提示的场景中存在的模糊性。在计算损失时，我们选择了MSE（Mean Squared Error）。
同时，对于边界框回归，采用平滑L1损失。同样重要的是要注意，任何未标记的区域都被排除在损失计算过程之外。在推理过程中，我们只应用基于点的非最大值抑制（Point-NMS），并坚持选择前N个点的做法。
（3）在端到端回归阶段，我们根据边界矩形的大小将目标掩蔽区域分为三组：大（ max（h Himg，w Wimg）≥ 0.25
），中等（ max（h Himg，w Wimg）∈（0.05，0.25））和小（ max（h Himg，w Wimg）≤ 0.05 ）。对每组进行单独的损失计算。这三项改进极大地提高了性能。
在这里插入图片描述

PPN回归的总损失，其中LH-MSE表示硬挖掘MSE损失，LS-L1表示盒回归的平滑L1损失。

Loss：
对于Lite-SAM的综合训练，结合了掩码损失，它结合了SAM中的原始焦点损失和骰子损失。这种组合量化了预测掩模和真值之间的逐像素对齐。此外，均方误差损失测量IoU预测与预测掩码与真实掩码的交集之间的差异。mask损失正式表示为：
在这里插入图片描述
总的训练损失由 $L_{total} = L_{ppn} + L_{mask}$ 定义。

说明

这个轻量化SAM的创新在于：

改掉原来的image encoder，创新一个LiteViT：它由多尺度池化和MLP代替自注意力，且每一层都用残差连接，还设计了一个patch merge模块，neck也是conv1*1+LN。
设计一个AutoPPN端到端地自动生成点提示或者框提示，并增加了点置信度损失和框回归损失以校正。有个问题在于它描述位置编码太过复杂费劲所以设置AutoPPN，但是位置编码是在mask decoder中使用的，如果说AutoPPN出来的点和框已经包含位置信息，那么提示编码器和mask decoder也会有改变，但文章中没有描述，所以这里存疑。