“EdgeSAM: Prompt-In-the-Loop Distillation for On-Device Deployment of SAM”
尽管SAM性能优异，但是运行速度慢，且难以在移动端运行。近日，洋理工大学联合上海AI Lab发表了EdgeSAM，可以在移动端上实时运行的SAM变体。在iPhone 14上运行，速度可以达到38 FPS，速读提升40倍。
- 项目主页：https://mmlab-ntu.github.io/project/edgesam/
- 论文地址：https://arxiv.org/pdf/2312.06660.pdf
- Github地址：https://github.com/chongzhou96/EdgeSAM
摘要
本文介绍了EdgeSAM，这是一种优化后的Segment Anything Model（SAM）变体，专门针对边缘设备进行了加速，同时在性能上没有太大的妥协。该方法将原始的基于ViT的SAM图像编码器提炼为纯CNN架构，更适合于边缘设备。作者通过精心的基准测试，证明了任务无关的编码器提炼无法捕捉SAM所体现的全部知识。为了克服这个瓶颈，作者将提示编码器和掩码解码器都包含在提炼过程中，并在循环中使用框和点提示，以便提炼模型可以准确地捕捉用户输入和掩码生成之间的复杂动态。为了缓解点提示提炼中的数据集偏差问题，作者在编码器中加入了一个轻量级模块。EdgeSAM相比原始SAM实现了40倍的速度提升，并且在边缘设备上部署时比MobileSAM快14倍，同时在COCO和LVIS上提高了2.3和3.2的mIoUs。它也是第一个可以在iPhone 14上以超过30 FPS运行的SAM变体。
简介
本研究探讨了在边缘设备（如智能手机）上直接部署Segment Anything Model（SAM）的可行性，以实现实时交互式分割并促进其在各种下游任务中的集成。然而，由于SAM的计算量和内存消耗较大，我们无法在iPhone 14上运行SAM，即使在NVIDIA 2080 Ti上，吞吐量也只有每秒四张图片。为了提高速度，一种看似简单的解决方案是用更紧凑的版本替换SAM的巨大的ViT-based图像编码器，但这会显著降低性能。为了解决这个问题，我们提出了EdgeSAM，它在iPhone 14上以超过30 FPS的速度运行，并且与原始SAM的准确性相当。我们的关键思路是在知识蒸馏过程中考虑提示，使得学生模型能够接收到任务特定的指导，并专注于更难的训练目标，如更精细的边界。
知识蒸馏的关键因素不是特定于密集任务或基于查询的检测器的损失函数，而是在蒸馏过程中选择合适的提示。我们引入了一种新的方法：动态提示循环策略，通过与SAM的多粒度输出掩码对齐，引入新的提示来指导学生模型，强调错误分割的区域，从而增强学习过程。我们还研究了不同类型的提示、冻结编码器/解码器组件以及蒸馏目标的选择。此外，我们进行了详细的消融研究，重点关注选择骨干架构，特别是在考虑到在设备部署中关键的吞吐量和性能平衡方面。我们发现，纯CNN架构比ViT架构更有优势，这归因于当前设备上的AI加速器，如Apple Neural Engine (ANE)，主要针对CNN进行了优化。我们还注意到，SAM在处理模糊提示时遇到了输出粒度的挑战，我们提出了一个模块来解决这个问题，增强SAM对不同级别提示的解释和响应能力。
EdgeSAM模型在NVIDIA 2080 Ti GPU上比原始SAM快40倍，比MobileSAM快1.6倍，在iPhone 14上每张图像处理时间仅为14毫秒，比MobileSAM快14倍。EdgeSAM在SA-1B、COCO和LVIS数据集上的表现与原始SAM相似，在点提示性能方面比MobileSAM平均高2 mAP和2-3 mIoU。EdgeSAM能够在边缘设备上实时运行，为设备上的视频编辑和视频实例分割等下游应用开辟了多种可能性。
相关工作
高效模型设计：目前研究主要集中在设计高效的CNNs、transformers和它们的混合架构，以实现视觉表示学习。
知识蒸馏在检测和分割中的应用：大部分知识蒸馏的研究集中在分类任务上，也有一些研究将知识蒸馏技术应用于语义分割和目标检测等密集预测任务。最近，对于基于查询的检测器（如DETR）的专门知识蒸馏损失的研究也越来越多。
高效分割模型：以往的研究主要集中在特定领域内的封闭集分割，最近有一些研究开始设计适用于移动平台的高效分割模型。然而，对于移动设备上的交互式分割领域仍然较少探索。现有的一些尝试仍然存在计算效率和性能下降的挑战，因此在这个方向上还有很大的发展空间。
方法
EdgeSAM通过编码器蒸馏、prompt-in-the-loop蒸馏和轻量级模块嵌入粒度偏好来精简SAM模型。
SAM
SAM模型由图像编码器、提示编码器和掩码解码器组成，其中图像编码器是最耗费计算成本和参数的部分。SAM可以处理四种类型的提示，包括点、框、掩码和文本。SA-1B数据集是目前最大的分割数据集，具有超过10亿个掩码注释。该数据集的特点使得SAM的蒸馏不同于以前的分割模型。
EdgeSAM
EdgeSAM旨在将SAM的能力转移到更紧凑的模型上，以便在边缘设备上部署。它保留了SAM的编码器-解码器架构，并旨在保持零-shot交互分割的性能。使用SA-1B数据集的1%数据训练EdgeSAM，并在COCO和LVIS数据集上评估其零-shot可转移性。
Encoder-Only知识蒸馏
MobileSAM采用了像素级特征蒸馏损失函数，用于对SAM的图像编码器和高效网络进行对齐。为了实现通道对齐，MobileSAM移除了学生模型中最后两个阶段的下采样操作，并使用投影层进行通道对齐。此外，MobileSAM还构建了一个小型的FPN网络，用于将特征上采样到所需的分辨率，并与前一阶段的特征进行逐元素相加。然而，尽管采用了不同的高效主干网络，但仍存在较大的性能差距。因此，MobileSAM提出了在蒸馏过程中考虑提示信息以提供任务特定的指导。
Prompt-In-the-Loop知识蒸馏
SAM模型的原始解码器非常轻量级，仅占总参数的0.6％，因此保留其架构以继承预训练权重。解码器包括两个流的双向Transformer，其中一个流的输入是来自图像编码器的输出特征图f，另一个流的输入是稀疏提示嵌入p与四个掩码标记m和一个IoU标记c相连接。通过实验研究，发现以教师模型的掩码输出作为学生模型的监督目标最有效。采用Dice损失和BCE损失的组合作为掩码损失。注意，允许梯度反向传播到图像编码器中，以便共同学习。
对于知识蒸馏过程，选择训练迭代的提示需要仔细设计，因为微调掩码解码器会威胁到零样本泛化能力。冻结掩码解码器或使用LoRA进行调节可以缓解问题，但也限制了训练期间观察到的提示的性能上限。引入动态提示采样策略可以提高蒸馏过程的效果，实现三个关键目标：动态生成多样化的提示组合、准确识别学生模型不准确的掩码区域、促使SAM产生高质量的掩码以提供更精确的指导。
本方法从交互式分割方法的最新进展中汲取灵感，通过在蒸馏阶段的循环中迭代采样新的提示来提高学习过程。研究发现，与依赖于密集预测或基于查询的方法的蒸馏相比，动态地将适当的提示组合输入到蒙版解码器中的蒸馏方法更加有效。实验结果表明，该方法优先考虑了提示的战略使用，从而增强了学习过程。
粒度先验
本文提出了一种基于区域建议网络的简单有效的模块，可以显式地嵌入特定数据集的粒度先验，并且可以选择关闭。该模块可以帮助SAM确定所需的输出粒度，特别是在使用单点提示时。该模块在图像编码器保持冻结的情况下构建，使用特定数据集（如COCO）进行训练，通过合并置信度加权的K个最近邻的提议框来确定输出粒度。最后，将合并的框与点输入一起作为提示输入到掩码解码器中。
训练和应用
EdgeSAM模型的训练分为三个阶段。第一阶段使用MobileSAM进行编码器知识蒸馏，对1%的SA-1B图像进行训练。第二阶段在相同的图像集上应用prompt-in-the-loop蒸馏，使用点和框提示作为输入的一部分。最后一个阶段是可选的，除了轻量级的RPN之外，其他模块都被冻结，并使用常用的focal loss和Huber loss对类别无关的真实框进行训练。
EdgeSAM模型可以使用点或框提示进行推理，并且可以逐步添加额外的点进行进一步的细化。为了展示EdgeSAM的实际效用，我们开发了一个设备上的演示。强烈建议探索这个演示，以充分了解EdgeSAM的能力。
实验
本文介绍了一种基于视觉提示的图像分割模型，主要通过吞吐量和准确性来评估模型效率。准确性分为三种情况，包括使用真实框和逐步添加更多点提示进行细化，使用真实掩模中心点作为初始提示，以及利用外部物体检测器提供框提示。评估数据集包括SA-1B、COCO和LVIS，其中SA-1B使用1%的图像进行蒸馏，构建了一个名为SA-1K的测试集。
实现细节
EdgeSAM分为三个阶段进行训练。第一阶段使用编码器的蒸馏损失函数进行训练，第二阶段使用Prompt-in-the-loop蒸馏损失函数进行训练，第三阶段是可选的RPN训练。具体训练细节包括使用AdamW优化器、余弦衰减学习率等。
定量结果
效率。我们比较了SAM、MobileSAM和EdgeSAM在桌面和移动平台上的编码器FPS。在使用NVIDIA 2080 Ti进行评估时，我们使用ONNX编译模型，以避免Python解释器的开销，并确保GPU以100%的负载运行。在iPhone 14上，我们使用Core ML Tools作为编译器，并使用Xcode提供的基准工具来测量吞吐量。如Tab. 1所示，EdgeSAM比SAM和MobileSAM快得多，特别是在移动平台上。特别是，在NVIDIA 2080 Ti上，EdgeSAM的编码器比SAM快40多倍，在iPhone 14上比MobileSAM快大约14倍。从基准报告中，我们观察到MobileSAM没有很好地优化设备上的AI加速器，如ANE，这解释了桌面和移动平台之间的巨大差异。此外，结合解码器，每次运行需要12毫秒，EdgeSAM在iPhone 14上的总体速度为38 FPS，使其成为首个在边缘设备上实时运行的SAM变体。
使用ground truth框提示。尽管非常轻量级，EdgeSAM能够生成准确的分割掩码。如表2所示，EdgeSAM在各种提示组合和数据集上始终优于MobileSAM，证明了其在实际应用中的有效性，包括零样本可迁移性和迭代改进。令人惊讶的是，在COCO数据集上，通过一个或两个点作为改进提示，EdgeSAM甚至超过了SAM。
使用中心点提示。与使用ground truth框提示类似，如表3所示，EdgeSAM在大多数情况下明显优于MobileSAM。然而，正如我们在第3.2.3节中讨论的
使用外部目标检测器提示。本文研究了结合外部物体检测器的各种SAM变体在COCO数据集上的mAP表现。结果表明，EdgeSAM表现优于MobileSAM，但落后于SAM。这种性能差距可能反映了模型容量的固有限制，也可能是由于仅使用groundtruth boxes进行训练，在推理过程中可能会出现差异。未来的研究方向是解决这些差异并探索其影响。同时，本文还报告了边界IoU以反映边界准确性。
消融分析
每个提议组件的有效性。本文介绍了两种技术：prompt-KD和轻量级RPN，它们都能有效提高模型性能。prompt-KD通过提供任务特定的监督来优化性能，而轻量级RPN在可管理的计算成本下运行良好。此外，作者还强调了在RPN的训练阶段冻结其他组件的重要性，以确保模型的泛化能力。
骨干选择。MobileSAM提出了一种解决SAM输出分辨率与高效骨干网络不匹配的方法，通过FPN将低分辨率特征图与高分辨率特征图融合。在比较了基于ViT、基于CNN和混合设计的高效骨干网络的性能和速度权衡后，发现基于纯CNN的RepViT-M1结合FPN可以达到最佳平衡。此外，由于许多边缘设备上的AI加速器（如ANE）对CNN进行了高度优化，因此在部署到边缘设备上时，速度差距变得更大。因此，EdgeSAM采用RepViT-M1作为图像编码器。在iPhone 14上，使用FPN进行分辨率对齐的RepViT-M1只需要14毫秒来编码1024×1024的输入，比TinyViT5M和EfficientViT-B1分别快14倍和4倍。
其他。本文介绍了一种基于知识蒸馏的文本生成模型Fine-Tuning方法，该方法通过在大型语料库上训练一个大型模型，然后将其知识蒸馏到小型模型中，以提高小型模型的性能。作者通过实验发现，在冻结解码器的情况下进行蒸馏可以获得最佳的领域内准确性，而微调所有模块可以更好地推广到其他数据集。作者还探索了多种改进方法，包括在蒸馏过程中变化的提示采样循环次数、训练时间表的长度和应用于密集预测任务的知识蒸馏损失等。最终，作者提出了一种基于提示感知的蒸馏方法，以进一步提高模型性能。
定性结果
EdgeSAM是一个具有零样本泛化能力的模型，可以通过一个点、两个点或一个框作为提示来生成遮罩质量的可视化结果。这些结果不是来自EdgeSAM训练集SA-1B数据集，证明了其零样本泛化能力。此外，作者还实现了一个演示应用程序，展示了EdgeSAM的更多定性结果，并建议读者探索该演示以更好地理解EdgeSAM的能力和应用。
总结
本文介绍了EdgeSAM，这是第一个在边缘设备上实时运行的SAM变体。通过将SAM蒸馏为轻量级CNN架构，实现了这一目标。作者提出了一种prompt-in-the-loop知识蒸馏方法，考虑了SAM的编码器和解码器，并提供了任务特定的监督信号。实验结果表明，EdgeSAM在效率和准确性方面比SAM和MobileSAM更有效。作者还提出了一些未来的研究方向，包括量化、模型剪枝、设备上的优化、混合精度推理等。作者希望EdgeSAM能够鼓励更多的实际应用，并继续努力改进。
评论0