【读点论文】EdgeSAM: Prompt-In-the-Loop Distillation for On-Device Deployment of SAM，知识蒸馏，更适合部署的轻量卷积网络作为主干

EdgeSAM: Prompt-In-the-Loop Distillation for On-Device Deployment of SAM

Abstract

• 本文介绍了EdgeSAM，这是分段任意模型（SAM）的一种加速变体，它经过优化，可在边缘设备上高效执行，性能折衷最小。我们的方法包括将原始的基于ViT的SAM图像编码器提取为纯基于CNN的架构，更适合边缘设备。我们仔细地对各种提取策略进行了基准测试，并证明了任务不可知的编码器提取无法捕获SAM中包含的全部知识。为了克服这一瓶颈，我们在提取过程中同时包括提示编码器和掩码解码器，并在环圈中包含框和点提示，以便提取的模型能够准确地捕获用户输入和掩码生成之间的复杂动态。为了缓解点提示提取引起的数据集偏差问题，我们在编码器中加入了一个轻量级模块。与原始SAM相比，EdgeSAM的速度提高了40倍，而且它的性能也优于MobileSAM，部署在边缘设备上的速度是MobileSAM的14倍，同时将COCO和LVIS上的mIoU分别提高了2.3和3.2。这也是第一个可以在iPhone 14上以超过30 FPS的速度运行的SAM变体。
• 论文地址：[2312.06660] EdgeSAM: Prompt-In-the-Loop Distillation for On-Device Deployment of SAM (arxiv.org)
• 项目地址：chongzhou96/EdgeSAM: Official PyTorch implementation of “EdgeSAM: Prompt-In-the-Loop Distillation for On-Device Deployment of SAM” (github.com)
• 由于SAM最重的部分就在于image encoder的vit，一般蒸馏都会选择先蒸馏encoder部分,如MobileSAM和EfficientSAM，只不过蒸馏方法不同，论文此处借鉴了MobileSAM的方法，只不过将MobileSAM的tinyvit改成了RepVit，相比tinyvit增加了fpn融合多尺度特征。
• 在仅蒸馏encoder提升到了一定瓶颈的情况下，论文提出了对SAM的第二部分decoder部分蒸馏，此处保留了mask decoder的结构作为学生网络，用原生的SAM权重初始化该部分权重。

Introduction

• 在这项研究中，我们探索了直接在智能手机等边缘设备上部署分段任何模型（SAM）的可行性，以实现实时交互式分段，并促进其在各种下游任务中的集成。应对这一挑战并非易事。首先，SAM不是为设备上部署而设计的。它的总参数计数为641M，在固定的1024×1024输入分辨率下需要2735个GFLOP。在没有任何专门优化的情况下，由于计算量和内存消耗大，我们无法在iPhone 14上运行SAM。即使在NVIDIA 2080 Ti上，吞吐量也只有每秒四个图像。
• 如MobileSAM所示，一个看似简单的提高速度的解决方案是用更紧凑的版本取代SAM的巨大的基于ViT的图像编码器。这种方法确实导致了相当大的速度提高，大约快了26倍，但也大大降低了性能。例如，随着框的提示，COCO数据集上的掩码mIoU从77.3下降到74.5，其中 GT 框。此外，当部署在边缘设备上时，MobileSAM的速度仍然远远不能实时。例如，在iPhone 14上，其吞吐量仅为每秒5张图像。
• 为了缓解性能下降并进一步实现边缘设备上的实时推理，我们提出了EdgeSAM，它在iPhone 14上运行超过30 FPS，与原始SAM的精度相当。我们的关键见解是在知识提取过程中考虑提示，以便学生模型获得特定任务的指导，并专注于更难的训练目标，如更精细的边界。
• 在我们广泛的实证研究中，出乎意料地清楚地发现，知识提取的关键因素不是专门为密集任务或基于查询的检测器定制的损失。相反，至关重要的是在蒸馏过程中对提示的战略性选择。针对这一问题，我们引入了一种新颖的方法：动态提示循环策略。这项技术有效地使学生模型能够吸收SAM中封装的复杂知识。它包括主动将学生模型与SAM的多粒度输出掩码对齐，并在学生模型显示不准确的区域反复引入新提示。这些量身定制的提示引导掩码解码器，强调不正确分割的区域，从而增强学习过程。此外，我们研究了各种训练配置，检查了不同类型的提示、冻结编码器/解码器组件的影响以及提取目标的选择。实验部分详细介绍了这些探索的细节和发现。
• 此外，我们进行了彻底的消融研究，重点是骨干架构的选择，特别是考虑到对设备部署至关重要的吞吐量-性能平衡。我们发现，纯基于CNN的架构是基于ViT的主干网实现最佳权衡的更有利选择。这归因于当前设备上人工智能加速器的现状，如苹果神经引擎（ANE），它们主要针对CNN而非ViT架构进行优化。这一观察结果还强调了我们提出的即时感知知识提取方法的多功能性，突出了其在不同架构中的适用性。
• 在我们的最后观察中，我们注意到，SAM在具有多粒度注释的数据集上进行了训练，当面临模糊提示（如单点）时，在解决输出粒度方面遇到了挑战。当在评估期间用COCO数据集上的中心点提示SAM时，这一点尤其明显；该模型并不一致地生成实例级遮罩，而是生成零件级遮罩。当SAM作为教师模型发挥作用时，这个问题变得更加明显。为了解决这一问题，我们提出了一个简单而有效的模块，旨在明确识别和适应特定于给定测试集或应用场景的粒度先验。该模块增强SAM准确解释和响应不同级别的提示歧义的能力。
• 因此，我们的EdgeSAM型号实现了显著的性能提升，在NVIDIA 2080 Ti GPU上的运行速度是原始SAM的40倍，是MobileSAM的1.6倍。值得注意的是，在iPhone 14上，EdgeSAM利用其纯粹基于CNN的主干[Repvit]来处理每张图像仅需14毫秒的图像编码。与MobileSAM在同一平台上的性能相比，这一速率快了14倍。即使与掩码解码器相结合，EdgeSAM也能保持超过30 FPS的总体速度，这标志着它是第一个能够在边缘设备上实时操作的SAM变体。就准确性而言，EdgeSAM在SA-1B、COCO 和 LVIS 数据集的盒子提示性能上与原始SAM非常相似。此外，它在这些数据集的点提示性能上平均超过MobileSAM 2mAP时和2-3 mIOU时。这种在边缘设备上实时运行SAM的能力为下游应用程序开辟了多种可能性，包括设备上的视频编辑和视频实例分割。

Related Work

• 高效的模型设计。该研究方向主要集中于设计高效的卷积神经网络[Ghostnet，mobilenetv3，Mobilenets，Squeezenet，Shufflenetv2，Mobilenetv2，Shufflenet]、Transformer[Edgenext，Mobilevit，Edgevits]及其混合架构，以视觉表示学习为目标。我们的工作也采用了高效的模型作为图像编码器，但与这些工作正交，因为它可以应用于各种有效的主干。
• 检测和分割中的知识提取。知识提炼的大多数研究都集中在分类任务上[Differentiable feature aggregation search for knowledge distillation，Self-training with noisy student improves imagenet classification，Revisiting knowledge distillation via label smoothing regularization，Distilling effective supervision from severe label noise]。一些研究已经将知识提取技术应用于密集预测任务，如语义分割和对象检测。这些工作中的常见方法包括利用教师和学生模型的密集特征之间的像素相关性或通道交互。最近，人们对开发基于查询的检测器（如DETR）的专门知识提取损失越来越感兴趣，如[Detrdistill，D3etr，Teach-detr]等工作所证明的那样。MobileSAM与我们的工作密切相关，在SAM编码器和紧凑主干之间实现像素特征提取。然而，它没有解决提示编码器和掩码解码器的问题，导致与原始SAM相比存在显著的性能差异。另一方面，FastSAM使用SA-1B数据集训练基于YOLACT的实例分割模型，并使用启发式规则进行后处理对象选择，这种方法与SAM原则基本一致。
• Efficient Segmentation Models. 先前关于有效分割的研究[Lpsnet，You only segment once，Sfnet，Espnet，Seaformer，Bisenet]主要集中在特定领域内的闭集分割，这些研究的很大一部分专门针对驾驶场景。最近，一些工作[Seaformer，Topformer]大胆地设计了适用于设备上实现的分割模型，能够在移动平台上高效运行。然而，设备上交互式分割的领域在很大程度上仍未被探索。MobileSAM代表了一个初步的尝试，但它仍然面临着计算效率方面的挑战和性能的显著下降，这表明在这个方向上有很大的发展机会。

Method

• 在本节中，我们首先简要介绍SAM（第3.1节），然后介绍我们提出的方法EdgeSAM，该方法通过编码器蒸馏（第3.2.1节）、循环蒸馏提示（第3.2.2节）和嵌入粒度偏好的轻量级模块（第3.2.3节）来精简SAM。

Preliminaries: SAM

• SAM模型由三个组件组成：图像编码器、提示编码器和掩码解码器。特别地，图像编码器占用了最多的计算成本和参数。它遵循了ViTDet的主干设计，本质上是一个层次视觉Transformer。图像编码器的输入和输出大小保持固定在（3、1024、1024）和（256、64、64）。
• SAM可以处理四种类型的提示，包括点、框、掩码和文本。点和框被视为稀疏提示，用随机位置嵌入编码，并用类型指示符标记求和。掩码和自由格式文本分别嵌入了一个微小的CNN和CLIP的文本编码器。在本文中，我们重点讨论点和框提示。
• 掩码解码器将图像特征、嵌入的点和/或框提示、掩码预测令牌和IoU预测 token 作为输入。所有输入都将使用轻量级双向transformer进行混合和编码。结果，每个掩码标记被转换为动态线性分类器，该分类器可以计算图像特征的每个空间位置的前景概率。IoU token 产生每个掩码的掩码置信度分数。
• 除了基于交互式提示的模型设计外，SA-1B数据集也对SAM的零样本通用性做出了重大贡献。它是迄今为止最大的分割数据集，在1100万张图像中有超过10亿个掩模注释。在这里，我们强调了SA-1B的几个特征，以更好地理解SAM的能力：（1）掩码注释是通过提示模糊感知模型来获得的，该模型是在半自动标记的数据上训练的，具有网格点；（2） mask 是阶级不可知论的；（3） 并且注释在实例和部件级别都是多粒度的。正如我们稍后将讨论的，SA-1B的这些特性使SAM上的蒸馏与以前的分割模型不同，并且更具挑战性。

EdgeSAM

• EdgeSAM的目标是将SAM的功能转移到一个更紧凑的模型中，这使得在边缘设备上部署变得可行。特别是，如下图所示，EdgeSAM保持了SAM的编码器-解码器架构，旨在保持具有框和点提示的零样本交互式分割的性能。我们仅使用SA-1B数据集的1%数据训练EdgeSAM，并评估其在COCO[28]和LVIS[11]数据集上的零样本可转移性。

  • 

  • EdgeSAM概述。我们首先在SAM和EdgeSAM的图像编码器的输出特征之间仅应用编码器KD。然后，我们采用循环提示KD，从错误分割的区域中交互式地采样点提示。请注意，初始提示也可以是一个框。

Encoder-Only Knowledge Distillation

• 受MobileSAM的启发，我们在SAM的图像编码器 $T_{enc}$ 和高效网络 $S_{enc}$ 之间采用了像素特征提取损失Lp，如下所示：

  • $$L_p=MSE(T_{enc}(I),S_{enc}(I)),$$

  • 其中 I 表示输入图像。由于学生模型和SAM图像编码器的下采样步长和特征通道没有对齐，MobileSAM删除了学生模型最后两个阶段的下采样操作，并使用投影层对齐通道维度。我们还使用投影层进行通道对齐，但保持下采样层不变。相反，我们构建了一个微小的FPN，它将特征上采样到所需的分辨率，并与前一阶段的特征进行逐元素相加。

• 我们探索了遵循这一范式的各种高效骨干，包括基于ViT的[Tinyvit]、基于CNN的[Repvit]和混合网络[Efficientvit]。然而，我们发现总是存在相当大的性能差距。针对密集预测任务设计的更长时间表或使用蒸馏损失[Structured knowledge distillation for semantic segmentation，Channel-wise knowledge distillation for dense prediction]的训练并没有显示出明显的改善。因此，我们进一步建议在蒸馏过程中考虑提示，以提供特定任务的指导。

Prompt-In-the-Loop Knowledge Distillation

• 由于原始SAM解码器足够轻，仅占SAM总参数的0.6%，因此我们保留了其架构，以便可以继承预先训练的权重。在这里，我们重新访问SAM的掩码解码器。它是一个双流双向 Transformer，来自图像编码器的输出特征图 f 作为一个流的输入，稀疏提示嵌入 p 与四个掩码 token m和一个IoU token c级联，作为另一个流的输入。稀疏提示可以是正/负点和单个框的任何组合，指示用户感兴趣的对象。

• 有了这些输入，教师和学生模型就有了许多潜在的提取目标，包括细化的特征图、掩码/IoU标记、两个输入流之间的交叉注意力以及输出掩码逻辑。通过实证研究，我们发现以教师 mask 输出为 GT 来监督学生是最有效的。更多细节见附录。我们将解码器损耗公式化如下：

  • $$L_d=L_{mask}(\varphi (T_{dec}(f_t,p,m,c)),S_{dec}(f_s,p,m,c))$$

  • 其中，$\varphi （\cdot ）$ 是二值阈值，ft 和 fs 分别表示来自教师和学生图像编码器的特征。教师和学生共享同一组p、m和c，这些数据在训练期间被冻结。我们采用 Dice损失 和 BCE损失的组合作为掩模损失Lmask。注意，我们允许对图像编码器进行梯度反向传播，以便对其进行联合学习。

• 虽然蒸馏损失简单得令人沮丧，但每次训练迭代的即时选择都需要仔细设计。我们观察到，微调掩码解码器对零样本泛化能力构成威胁。也就是说，当使用训练中未使用的提示组合进行推理时，使用特定提示组合进行训练会破坏其能力。例如，当使用框提示进行测试时，仅使用分数进行训练会导致性能显著下降。虽然冻结掩码解码器或在解码器上使用LoRA进行调节可以缓解问题，但它也限制了训练期间观察到的提示的性能上限。同时，我们发现，当提示不明确时，即使是原始SAM也会输出不令人满意的掩码预测，例如单点。如图所示，在COCO数据集上，以对象中心为提示，SAM的掩码mIoU仅为53.6。因此，在这种情况下调整学生的产出可能不是最佳的。

  • 

  • SAM在不同提示下的行为。当提示不明确时，例如一个点，SAM会产生次优结果。更多信息提示可以解决问题。

  • 如何选择prompt，对于原始的SAM，如果选点不合适，输出的mask效果也不好。为了解决这个问题，论文提出了动态采样prompt的方法，思路如下：动态生成一部分box和点作为prompt；识别出蒸馏的学生模型表现不好的区域，然后让模型聚集在这个地方；直到教师模型(原生的SAM)生成更准确的结果。

• 为了提高蒸馏过程的效率，我们引入了一种动态即时采样策略。这种方法旨在实现三个关键目标：（1）从最初的提示（无论是方框还是点）动态生成一组不同的提示组合，（2）准确识别mask 内学生模型显示不准确的区域，从而将其注意力集中在这些特定的部分，以及（3）迫使教师模型，即SAM，生产高质量的mask，以获得更精确的指导。

• 从交互式分割方法[Getting to 99% accuracy in interactive segmentation，Reviving iterative training with mask guidance for interactive segmentation]的最新进展中获得灵感，我们的策略涉及在蒸馏阶段对循环中的新提示进行迭代采样。详见算法1，我们从一个初始提示开始。在相同的概率下，SA-1B数据集提供的框或点提示被输入到教师和学生模型的解码器中。随后，我们确定了教师和学生的掩码预测存在分歧的区域。

• 

• 如下图所示，使用教师的输出作为参考，我们对新提示进行了统一采样：标记为假阴性的区域中的一个正点，或标记为假阳性区域中的负点。然后将这些新采样的点与现有提示合并，用于随后的解码迭代。需要注意的是，每个提示都会导致不同粒度级别的四个掩码预测。在我们的分析中，差异是专门计算出IoU得分最高的教师 mask 与其对应的学生 mask 之间的差异。

  • 

  • 在蒸馏过程中，我们在循环中从FN/FP区域随机采样新的正/负点提示，以便学生动态地关注这些区域。

• 总之，我们的研究结果表明，在蒸馏过程中，将适当的提示组合动态馈送到掩码解码器中被证明是更有效的，而不是依赖于具有专门用于密集预测或基于查询的方法的损失的蒸馏。我们的提示循环蒸馏方法优先考虑提示的战略性使用，以增强学习过程。在实验部分，我们通过全面的消融研究进行了详细的讨论。

Granularity Priors

• 由于SA-1B是一个类不可知、多粒度、自动标记的数据集，其注释分布可能与人工密集标记的数据集中（如COCO）的注释分布非常不同。因此，对于不明确的提示，例如单个点，SAM很难确定所需的输出粒度。同时，如上文图所示，通过框提示，SAM可以轻松地精确定位目标粒度。此外，与反复点击或与框交互相比，在智能手机上的许多情况和应用程序中，单次点击是受欢迎的，例如点击和拖动。因此，我们提出了一个简单高效的模块，该模块显式嵌入某些数据集的粒度先验，如果SAM的原始行为是首选的，则可以选择性地关闭该模块。
• 在图像编码器保持冻结的情况下，我们在其上构建了一个轻量级区域建议网络（RPN），该网络由特征金字塔网络（FPN）和共享检测头组成。为了提高效率，我们遵循EfficientDet提出的设计。RPN在特定的数据集上进行训练，例如COCO，以预先捕获其粒度。在推理过程中，我们合并以点提示的K个最近邻居为中心的建议框，这些建议框由它们的置信度得分加权。最后，我们将合并后的框与点输入组合在一起，作为输入到掩码解码器的提示。
• 对于点prompt，由于SA-1B数据集是无类别的，SAM很难区分要输出的mask是多大粒度的，输入一个点prompt，结果可能是不准确的。此处增加了一个轻量的RPN和目标检测头。推理时，从RPN的候选框中选取和点prompt比较近的K个box作为prompt，一起送入mask decoder中。

Training and Application

• Training Pipeline. 我们把训练分为三个阶段。在第一阶段，我们在MobileSAM之后对1%SA-1B图像执行仅编码器知识提取。在第二阶段，我们对同一图像集应用循环中提示蒸馏，但将点和框提示作为输入的一部分。在可选的最后阶段，我们冻结了轻量级RPN以外的模块，并使用具有常用焦点损耗和Huber损耗的类不可知 GT 盒进行训练。

• 推理和设备演示。我们的EdgeSAM模型能够通过点或框提示启动推理，并且，与SAM类似，它可以逐步包含额外的点以进行进一步的细化。为了展示EdgeSAM的实用性，我们开发了一个设备演示。我们强烈建议您探索此演示，以充分了解EdgeSAM的功能。

Experiments

• Efficiency. 我们主要通过单个NVIDIA 2080 Ti和iPhone 14的每秒吞吐量来衡量模型效率，因为它反映了真实世界的延迟。我们还报告了FLOP和总参数，以供侧面参考。
• Accuracy. 根据提示类型和来源，我们将精度测量分为三种场景。（1） 我们使用 GT 框作为初始提示，并迭代添加更多点提示进行细化。（2） 类似于（1），但使用 GT 遮罩的中心点作为初始提示。（3） 利用外部对象检测器来提供框提示。对于（1）和（2），我们计算所有实例的mIoU，对于（3），我们报告mAP和边界IoU。对于所有评估，我们将第一个掩码 token 的输出作为预测。
• 数据集。我们考虑SA-1B、COCO和LVIS来评估域内性能和零样本可转移性。对于蒸馏，使用整个SA-1B数据集的1%的图像。我们从SA-1B构建了一个名为SA-1K的测试集。特别是，我们从SA-1B中随机选择1000张图像，这些图像在训练过程中没有使用。对于每个图像，我们随机采样64个实例以避免内存溢出。对于COCO和LVIS，我们使用它们的默认验证集。

Implementation Details

• 如第3.3节所述，我们将训练分为三个阶段。在第一阶段中，我们对1%的SA-1B图像进行训练，仅使用编码器的蒸馏损失 $L_p$。我们将该模型训练10个epoch，批量大小设置为64。我们采用AdamW优化器，并将初始学习率设置为1.25e-2。我们采用余弦衰减时间表，它将学习率逐渐降低到6.25e-7。在第二阶段中，我们加载在第一阶段中训练的编码器权重，并从SAM继承解码器权重。然后，我们在1%SA-1B上使用提示环路蒸馏损失 $L_d$ 训练5个epoch。批量大小为16，学习率从1e-4衰减到1e-5。我们将每个图像允许的最大实例数设置为16，提示采样的循环数为1。在RPN训练的最后可选阶段，我们遵循COCO上MMDetection的1x时间表。

Quantitative Results

• Efficiency. 在下表中，我们比较了桌面和移动平台上SAM、MobileSAM和EdgeSAM的编码器FPS。当使用NVIDIA 2080 Ti进行评估时，我们使用ONNX编译模型以避免Python解释器的开销，并确保GPU以100%的有效负载运行。在iPhone 14上，我们使用Core ML Tools作为编译器，并使用Xcode提供的基准测试工具来测量吞吐量。如下表所示，所提出的EdgeSAM明显快于SAM和MobileSAM，尤其是在移动平台上。特别是，EdgeSAM的编码器比NVIDIA 2080 Ti上的SAM快40多倍，比iPhone 14上的MobileSAM快约14倍。从基准报告中，我们观察到MobileSAM没有被设备上的人工智能加速器（如ANE）很好地优化，这解释了桌面和移动平台之间的巨大差异。此外，EdgeSAM与每次运行花费12毫秒的解码器相结合，在iPhone 14上提供了38 FPS的整体速度，这使其成为第一个在边缘设备上实时运行的SAM变体。更多比较见附录A。

  • 

  • 效率比较。对于FPS，我们只报告编码器FPS，因为所有型号共享相同的解码器，iPhone 14上的每张图像只需要12毫秒。*表示带有RPN的EdgeSAM。

• GT框提示。尽管EdgeSAM非常轻，但它可以生成准确的分割掩码。如下表所示，所提出的EdgeSAM在各种提示组合和数据集上始终优于MobileSAM，证明了其在现实世界应用中的有效性，包括零样本可转移性和迭代细化。令我们惊讶的是，在COCO数据集上，有一两个点作为细化提示，EdgeSAM甚至超过了SAM。

  • 

  • 以GT盒子作为提示进行表演。我们报告测试集中所有实例的mIoU+1 pt。表示附加一个额外的细化点作为提示。粗体表示最好，下划线表示第二好。

• 使用中心点提示。与GT框的提示类似，EdgeSAM在大多数情况下显示出明显优于MobileSAM的优势，如下表所示。然而，正如我们在第3.2.3节中所讨论的，对于不明确的提示，例如单个点，原始SAM并不总是以所需的粒度产生掩码。此外，EdgeSAM是在SA-1B数据集上训练的，其粒度分布与COCO的粒度分布非常不同。因此，使用所提出的任务感知蒸馏，与MobileSAM相比，EdgeSAM在处理模糊提示时更有可能符合SAM粒度分布。为了明确地捕捉这样的粒度先验，我们提出了一种轻量级RPN，它如预期的那样工作，并使COCO上的单点性能从48.0提高到54.3。

  • 

  • 以中心点为提示的性能。类似于上表，但使用遮罩中心点作为初始提示。*表示将RPN用于粒度先验。

• Prompting with External Object Detectors. 最后，我们将所有SAM变体与外部对象检测器相结合（我们使用MMDetection中的Detic），并评估其在COCO上的mAP性能。此设置检查了一个实际用例，其中自动分割比交互式分割更可取，例如，在跟踪框内生成掩码。我们使用对象检测框架来生成边界框和相关的类标签，确保观察到的mAP变化仅归因于生成的掩码的质量。这种方法还评估了我们的模型在处理不精确的边界框输入时的弹性。我们还报告了边界IoU，以反映边界精度。下表显示，虽然EdgeSAM以相当大的优势优于MobileSAM，但它落后于SAM。这一性能差距不仅反映了模型容量的固有限制，还反映了仅使用 GT 框训练的结果，这可能导致推理过程中的差异。解决这些差异并探索其含义是未来研究的一个有希望的方向。

  • 

  • 使用来自外部对象检测器的框作为提示的性能。我们报告了COCO数据集上的掩码mAP和边界IoU。物体探测器的盒mAP为47.4。

Ablation Studies

• 各拟议组成部分的有效性。我们研究了引入的技术的有效性，即prompt-in-the-loop知识提取（prompt KD）和结合粒度先验的轻量级区域建议网络（RPN）。如下表a所示，与仅编码器的知识提取相比，提示KD在性能上有了显著的增强，特别是当补充了额外的细化点时。Prompt KD的优势在于其提供特定任务的监督，这比仅编码器KD提供的一般指导更明确、更有针对性。此外，它在不准确分割的区域动态生成新提示的策略将更多的注意力放在这些区域上，在这个过程中创建了不同的提示组合。

  • 

• 此外，如下表b所示，轻量级RPN以可控的计算成本有效运行。需要注意的是，在RPN的训练阶段，我们会冻结其他组件，包括骨干网络。这种方法允许在推理过程中动态停用RPN，确保模型的泛化能力保持不变。

  • 

• 骨干的选择。如第3.2.1节所述，SAM的输出分辨率和有效主干的输出分辨率不一致。MobileSAM提出删除最后几个阶段的下采样操作，而我们实验通过FPN将低分辨率特征图与高分辨率特征图融合。此外，我们比较了基于ViT的[Tinyvit]、基于CNN的[Repvit]和混合[Efficientvit]设计后高效主干的性能-速度权衡。如下表c所示，纯基于CNN的RepViT-M1与用于分辨率对准的FPN实现了最佳平衡。此外，由于许多设备上的人工智能加速器，如ANE，都是针对卷积神经网络进行高度优化的，因此当部署在边缘设备上时，速度差距会变得更大。例如，在iPhone 14上，使用FPN进行分辨率对齐，RepViT-M1对1024×1024输入的编码只需14毫秒，分别比TinyViT-5M和EfficientViT-B1快14倍和4倍。因此，我们的EdgeSAM采用它作为图像编码器。

  • 

• 其他消融。在下表d中，我们测试了关于模块冻结和LoRA的几种配置。特别地，LoRA被应用于解码器中的关注块的查询和值投影层。结果表明，在蒸馏过程中冻结解码器可以获得最佳的域内精度，而微调所有模块在其他数据集上可以更好地推广。为了获得更实用的设置，我们选择微调编码器和解码器。

  • 

• 此外，我们在蒸馏过程中改变了即时采样回路的数量，在表e中，我们发现一个额外的回路就足够了，因此默认情况下将回路数量设置为1。

  • 

• 此外，由于仅编码器蒸馏和loop中提示蒸馏的最佳优化设置不同，我们发现与两个阶段的训练相比，联合训练产生次优结果（见表f）。

  • 

• 因此，我们建议按顺序训练他们。最后，在表g中，我们探索了几种提高纯编码器蒸馏性能的方法，包括针对更长时间表的训练和应用为密集预测任务设计的知识蒸馏损失。虽然它们提高了点提示的性能，但框提示的性能略有下降，这表明仅编码器提取的性能是上限。这促使我们探索 prompt-aware 蒸馏。

  • 

Qualitative Results

• 在下图中，我们分别用一个点、两个点和一个框作为提示来可视化EdgeSAM的掩码质量。请注意，没有任何图像来自EdgeSAM训练的SA-1B数据集，这表明了其零样本泛化能力。我们在附录C中提供了SAM和EdgeSAM之间的更多视觉比较。我们还实施了一个设备上演示应用程序，展示了EdgeSAM的更多定性结果。我们强烈建议您探索此演示，以便更好地了解EdgeSAM的功能和应用程序。

  • 

  • 使用点框提示的EdgeSAM的定性结果。绿色和红色的星星分别表示正点和负点。附录C提供了EdgeSAM和SAM之间的更多比较结果。

Conclusion

• 在本文中，我们介绍了EdgeSAM，这是第一个在边缘设备上实时运行的SAM变体。我们通过将SAM提取到基于CNN的轻量级架构中来实现这一点。我们的初步实验表明，现有的提取方案仅涉及图像编码器，因此是任务不可知的，无法向学生模型揭示SAM的全部知识谱。因此，我们提出了一种即时循环知识提取方法，该方法同时考虑SAM的编码器和解码器，并提供特定任务的监督信号。在各种设置下，SA-1B、COCO和LVIS的实验证明了与SAM和MobileSAM相比，EdgeSAM在效率和准确性方面的有效性。
• 尽管能够在边缘设备上以相当快的速度运行EdgeSAM，但我们仍有几个研究方向，这些方向可能会提供提升，还有待研究，包括量化、模型修剪、设备优化、混合精度推理等。此外，我们在训练过程中没有应用任何增强，因此适当的数据或即时增强也可能是有希望的方向。我们希望EdgeSAM将通过其设备上的交互式分割功能鼓励更多真实世界的应用程序，并将继续进行进一步的改进。

A. More Quantitative Results

• 在我们的主要论文中，我们将我们的方法与SAM和MobileSAM进行了比较。在这里，与最近的一项工作FastSAM相比，我们补充了更多的结果。与MobileSAM和EdgeSAM不同，FastSAM不遵循SAM的编码器-解码器架构。相反，它生成所有实例的掩码，并通过post-prediction handcrafted rules选择感兴趣的对象。因此，在掩码预测期间不考虑提示，这限制了输出的灵活性。根据FastSAM的实验设置，我们使用ViTDet-H作为检测器在COCO数据集上测量平均精度（AP），该检测器提供方框和类标签。
• 在这里，我们没有根据原始SAM论文的设置应用额外的掩码细化迭代，因此数字比SAM中报告的略差。请注意，ViTDet-H比我们在主要论文中使用的检测器强得多（方框AP:58.7对47.4）。此外，FastSAM使用2%的SA-1B图像进行训练，这是MobileSAM和EdgeSAM的两倍。如下表所示，EdgeSAM始终以可观的优势超过FastSAM和MobileSAM，但仍落后于原始SAM。下表还表明，有了更多的训练数据，EdgeSAM可以获得更好的性能，尤其是对于大型对象。此外，在效率方面，EdgeSAM实现了最佳效果，比所有同行都快得多，尤其是在iPhone 14等边缘设备上。

  • 

  • 带方框提示的性能。我们在COCO数据集上报告了掩码mAP。使用ViTDet-H作为检测器，其 box mAP为58.7，以提供盒提示。对于速度基准测试，我们推断编码器和解码器（而不是仅编码器）。因此，这些数字与主要论文中的数字不同。FLOP是基于1024×1024的输入分辨率计算的。*表示的数字是从MobileSAM复制的。3x和10x表示使用更多数据进行训练。粗体表示最好，下划线表示第二好。

B. More Ablation Studies

• Training Datasets. 在我们的默认设置中，我们使用SA-1B数据集的1%图像来训练EdgeSAM。在下表a中，我们研究了训练数据集对可转移性的影响。具体而言，我们在1%SA-1B、COCO和LVIS之间切换训练数据集时，保持模型和训练管道不变。请注意，所有数据集都包含相似数量的图像（大约110K）。结果表明，使用SA-1B进行的训练在使用框提示的所有测试集中都能获得最佳性能，而使用点提示，在同一领域中的训练和测试会产生最佳结果。这表明，通过信息提示（如方框），SA-1B数据集上的蒸馏足以进行零样本转移。然而，对于不明确的提示，例如单个点，需要努力捕获测试数据集的粒度先验。这促使我们利用额外的RPN来显式地捕获此类先验。

• 解码器的训练损失。在我们的主要论文中，我们提到以教师掩码输出作为 GT 来监督学生是解码器损失的最有效的方法。在这里，我们提供了我们实证研究结果的更多细节。特别地，我们比较了各种配置，包括（1）将掩码损失与来自教师和学生解码器的注意力图对齐相结合；（2） 将掩码损失与教师和学生之间的IoU预测对齐相结合；（3） 在教师解码器和学生解码器之间对准输出特征图。如下表b所示，普通掩模损耗监督产生了最佳性能。因此，我们不再探索不同的损失组合，而是转而研究在蒸馏过程中选择适当提示的策略，并进一步提出环管蒸馏中的提示。

  • 

  • 更多消融研究。框和中心表示测试EdgeSAM，GT框和GT掩码的中心分别作为提示。通过Detic对象检测器提供框提示和类标签（框mAP 47.4）来获得COCO-mAP和边界IoU。请注意，对于所有消融，我们将prompt采样循环的数量设置为0（无重新采样）。SA-1B*表示1%的SA-1B。粗体表示最佳数字和我们采用的默认配置。

• 提示的数量。假设SA-1B的每个样本包含大约100个实例，在一个批次中训练所有实例会导致VRAM溢出。因此，我们为每次迭代随机抽取一个子集。在上表c中，我们研究了采样提示数量的影响，发现16个提示实现了最佳折衷。因此，我们将其设置为默认值。

C. Qualitative Results Compared with SAM

• 在下图中，我们提供了EdgeSAM的更多定性结果，并将其与单点、框和一切模式下的原始SAM进行了比较。请注意，由于EdgeSAM遵循与SAM相同的编码器-解码器架构，因此一切模式都很昂贵。具体来说，在 everything 模式下，32×32网格中的点被馈送到解码器，这导致1024次解码器推理。由于EdgeSAM是为推理速度敏感设备构建的，我们不主张在实践中使用EdgeSAM的 everything 模式。因此，我们提供可视化示例主要讨论其与SAM相比的行为。正如预期的那样，EdgeSAM的整体掩码质量不如SAM，但它们是可比较的，尤其是对于框提示。EdgeSAM的一个常见故障案例是错误分割对象内的孔，例如图中的圆环示例（第1行，第4列）。我们认为这可能是由于缺乏训练样本。但在实际使用中，可以在空穴区域中添加一个负点，以明确强制EdgeSAM排除这些区域。此外，在所有模式中，只保留自信分数超过一定阈值的 mask。因此，没有掩模的区域表明掩模预测不令人满意。与SAM相同，EdgeSAM覆盖了场景中的大多数对象，这证明了它的能力。

  • 

  • Qualitative results of EdgeSAM and SAM under point, box, and everything mode.