TinySAM: Llevando al límite la eficiencia del modelo Segment Anything

TinySAM: Pushing the Envelope for Efficient Segment Anything Model

December 21, 2023

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

cs.AI

Resumen

Recientemente, el modelo Segment Anything (SAM) ha demostrado una potente capacidad de segmentación y ha atraído una gran atención en el campo de la visión por computadora. Numerosos trabajos posteriores han desarrollado diversas aplicaciones basadas en el SAM preentrenado, logrando un rendimiento impresionante en tareas de visión posteriores. Sin embargo, SAM está compuesto por arquitecturas pesadas y requiere una gran capacidad computacional, lo que dificulta su aplicación en dispositivos de borde con limitaciones de cálculo. Con este fin, en este artículo proponemos un marco para obtener un modelo Tiny Segment Anything (TinySAM) mientras se mantiene un fuerte rendimiento en tareas de zero-shot. Primero, proponemos un método de destilación de conocimiento en todas las etapas con una estrategia de muestreo de prompts difíciles en línea para destilar un modelo estudiantil ligero. También adaptamos la cuantización post-entrenamiento a la tarea de segmentación con prompts, reduciendo aún más el costo computacional. Además, proponemos una estrategia jerárquica de segmentación de todo para acelerar la inferencia de "todo" en 2 veces con casi ninguna degradación del rendimiento. Con todos estos métodos propuestos, nuestro TinySAM logra una reducción computacional de órdenes de magnitud y amplía los límites para la tarea eficiente de segmentar cualquier cosa. Experimentos extensos en diversas tareas de transferencia zero-shot demuestran el rendimiento significativamente ventajoso de nuestro TinySAM frente a métodos comparables. Los modelos preentrenados y los códigos estarán disponibles en https://github.com/xinghaochen/TinySAM y 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.