EdgeSAM: Prompt-In-the-Loop Distillation for On-Device Deployment of SAM
EdgeSAM论文：https://arxiv.org/pdf/2312.06660.pdf
EdgeSAM代码：https://github.com/chongzhou96/EdgeSAM
1 概述
- 在2080Ti上，相比原生SAM，EdgeSAM推理速度快40倍；
- 在iPhone14上，相比MobileSAM，EdgeSAM推理速度快14倍，达到了38.7fps。
2 EdgeSAM
EdgeSAM的目标是将SAM的功能转移到一个更为紧凑的模型中，使得在边缘设备上的部署变得可行。因此，EdgeSAM保留了SAM的编码器-解码器架构，并旨在保持零样本交互式分割（通过框和点提示）的性能。EdgeSAM的核心是在知识蒸馏过程中考虑提示词，以便学生模型接受任务特定的指导，并专注于更难的训练目标，如更精细的边界。为此引入了一种动态提示词采样策略，旨在实现三个关键目标：（1）从初始提示词（无论是框还是点）中动态生成一组多样化的提示词组合；（2）准确识别学生模型在掩膜内表现不准确的区域，从而将其注意力集中在这些特定部分；（3）迫使教师模型（即SAM）产生高质量的掩膜，以提供更精确的指导。
2-1 仅编码器知识蒸馏 Encoder-Only Knowledge Distillation
由于学生模型和SAM图像编码器的下采样步长和特征通道并不一致，MobileSAM移除了学生模型最后两个阶段的下采样操作，并使用一个投影层来对齐通道维度。EdgeSAM构建了一个小型FPN（特征金字塔网络），将特征上采样到所需的分辨率，并与前几个阶段的特征进行逐元素相加来保持下采样层不变，最后也使用投影层进行通道对齐。
实验发现，使用不同的backbones、更长的训练计划或者使用为密集预测任务设计的蒸馏损失并没有表现出明显的性能提升。
2-2 循环Prompt知识蒸馏 Prompt-In-the-Loop Knowledge Distillation
SAM的掩膜解码器是一个双流双向转换器，其中来自图像编码器的输出特征图f作为其中一个流的输入，而稀疏提示嵌入p与四个掩膜标记m和一个IoU标记c连接在一起，作为另一个流的输入。稀疏提示可以是正/负点的任意组合以及一个单独的框，用于指示用户感兴趣的对象。
2-3 粒度先验 Granularity Priors
由于SA-1B是一个类别无关、多粒度、自动标注的数据集，其标注分布可能与由人力密集标注的数据集（如COCO）存在很大差异。
因此，对于模糊提示词，如单点，SAM很难确定所需的输出粒度。如下图所示，使用框提示词时，SAM可以轻松地确定目标粒度。
构建了一个轻量级的区域提议网络（RPN）[41]，该网络由特征金字塔网络（FPN）和一个共享检测头组成。RPN在特定数据集（例如COCO）上进行训练，以捕获其粒度先验。在推理过程中，合并那些中心点为点提示词的K个最近邻（根据置信度得分加权）的提议框。最后，将合并后的框与点输入一起作为提示词，输入到掩膜解码器中。
3 Experiments
