TinySAM：推动高效Segment Anything模型的发展

TinySAM: Pushing the Envelope for Efficient Segment Anything Model

December 21, 2023

作者: Han Shu, Wenshuo Li, Yehui Tang, Yiman Zhang, Yihao Chen, Houqiang Li, Yunhe Wang, Xinghao Chen

cs.AI

摘要

最近，分割任意模型（Segment Anything Model，SAM）展现出强大的分割能力，在计算机视觉领域引起了极大关注。许多后续研究基于预训练的SAM开发了各种应用，并在下游视觉任务中取得了令人印象深刻的性能。然而，SAM由于庞大的架构和需要大量计算资源，限制了SAM在计算受限边缘设备上的进一步应用。因此，在本文中，我们提出了一个框架，以获得一个微型分割任意模型（TinySAM），同时保持强大的零样本性能。我们首先提出了一种全阶段知识蒸馏方法，结合在线硬提示采样策略，来蒸馏出一个轻量级的学生模型。我们还将后训练量化方法应用于可提示的分割任务，进一步降低计算成本。此外，我们提出了一种分层分割一切策略，通过几乎没有性能降级，将一切推理加速了2倍。通过所有这些提出的方法，我们的TinySAM实现了数量级的计算减少，并推动了高效分割任意任务的发展。对各种零样本迁移任务的大量实验表明，我们的TinySAM在性能上明显优于对应方法。预训练模型和代码将在https://github.com/xinghaochen/TinySAM 和 https://gitee.com/mindspore/models/tree/master/research/cv/TinySAM 上提供。

English

Recently segment anything model (SAM) has shown powerful segmentation capability and has drawn great attention in computer vision fields. Massive following works have developed various applications based on the pretrained SAM and achieved impressive performance on downstream vision tasks. However, SAM consists of heavy architectures and requires massive computational capacity, which hinders the further application of SAM on computation constrained edge devices. To this end, in this paper we propose a framework to obtain a tiny segment anything model (TinySAM) while maintaining the strong zero-shot performance. We first propose a full-stage knowledge distillation method with online hard prompt sampling strategy to distill a lightweight student model. We also adapt the post-training quantization to the promptable segmentation task and further reduce the computational cost. Moreover, a hierarchical segmenting everything strategy is proposed to accelerate the everything inference by 2times with almost no performance degradation. With all these proposed methods, our TinySAM leads to orders of magnitude computational reduction and pushes the envelope for efficient segment anything task. Extensive experiments on various zero-shot transfer tasks demonstrate the significantly advantageous performance of our TinySAM against counterpart methods. Pre-trained models and codes will be available at https://github.com/xinghaochen/TinySAM and https://gitee.com/mindspore/models/tree/master/research/cv/TinySAM.