TinySAM：突破有效分段任務模型的極限

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.