极致高效压缩，手机就能实时跑的分割一切模型TinySAM已在昇思MindSpore开源

分割一切模型（Segment Anything Model，SAM）展现出了强大的分割万物能力，但是网络结构复杂，计算代价高，在资源受限的端侧应用中无法直接使用。同时，我们发现在SAM的Everything推理模式中，原始的网格点prompt的方式会带来巨大的计算代价。针对这个问题，我们提出了TinySAM模型，利用知识蒸馏、量化等手段对模型做了极致的压缩，同时提出了再Everything推理模式下的层次化推理策略。TinySAM在计算代价降低多个数量级的情况下，仍然保持了SAM模型强大的zero-shot分割能力。论文见：https://arxiv.org/abs/2312.13789同时我们开源了TinySAM的MindSpore代码和模型，供大家试用：MindSpore代码：https://gitee.com/mindspore/models/tree/master/research/cv/TinySAM下面是demo效果展示，TinySAM在手机端侧可完成给定点或框作为prompt的自动分割功能。

TinySAM: 极致高效的分割一切模型

这个工作通过知识蒸馏、后量化等方法对SAM模型做极致压缩，同时提出一种针对Everything推理的加速方法，探索分割一切模型性能和效率的极限。

图1：(a)展示了TinySAM的框架，通过一种全阶段知识蒸馏、后量化和层次化推理加速方法，模型计算量可以降低到原来的1/50；（b）展示了模型在zero-shot下游分割任务上仍然保持了良好的性能。

(1) 全阶段知识蒸馏

针对分割一切模型多模块、难优化的结构特点，我们提出一种全阶段的知识蒸馏方法，从多个特征层面进行蒸馏学习，从图像特征（embedding）、最终监督特征（mask）、中间过渡特征（token）三个层面进行蒸馏损失的设计，最大程度减少轻量化学生网络的学习难度。

图2：全阶段知识蒸馏框架

(2) 模型后量化

为了进一步降低模型计算量，提升推理效率，我们将TinySAM量化到8bit，其中卷积层和反卷积层采用channel-wise的非对称均匀量化，matmul和linear层分别采用headwise和layerwise的对称均匀量化。在量化策略上，逐层对网络进行串行量化，并采用与任务Loss更一致的hessian guided metric作为量化前后输出的距离度量。量化后的模型计算量能进一步减少到全精度的1/4左右，大幅降低计算量。

(3) 层次化分割一切

SAM模型在Everything（分割一切）推理模式中，通过密集的网格点作为Prompt进行推理，我们发现这种密集的Prompt输入和推理耗费了大量的计算量。因此，我们设计了一种层次化的分割一切的推理策略。先用较少的初始采样点得到初始结果，然后把高置信度区域的密集采样点移除，得到二次推理结果，与初始推理结果进行后处理融合得到最终的推理结果。这种策略在保持分割一切效果的同时，大幅减少了采样点数量，推理加速一倍左右。

图3：我们设计的层次化分割一切推理模式，我们先用较少的采样点得到初始的结果，然后我们把高置信度区域的密集采样点移除，得到二次推理结果后通过后处理与初始结果融合得到最终的分层推理结果。

实验结果

我们在多个zero-shot下游任务上进行了实验，包括在COCO和LVIS数据集的zero-shot实例分割实验，我们按照SAM论文的设置通过ViTDet得到的检测框作为Prompt输入到TinySAM模型中。相比MobileSAM，TinySAM在COCO的平均准确率提升0.9%，在LVIS的平均准确率能提升1.6%。即使采用8bit量化后，TinySAM仍然保持了性能优势。更多的实验结果和可视化结果请参考论文。