EfficientViT-SAM: Modello Segment Anything Accelerato Senza Perdita di Prestazioni

EfficientViT-SAM: Accelerated Segment Anything Model Without Performance Loss

February 7, 2024

Autori: Zhuoyang Zhang, Han Cai, Song Han

cs.AI

Abstract

Presentiamo EfficientViT-SAM, una nuova famiglia di modelli accelerati per il segmentazione di qualsiasi oggetto. Manteniamo il prompt encoder leggero e il decoder delle maschere di SAM, sostituendo il pesante image encoder con EfficientViT. Per l'addestramento, iniziamo con la distillazione della conoscenza dall'image encoder SAM-ViT-H a EfficientViT. Successivamente, conduciamo un addestramento end-to-end sul dataset SA-1B. Grazie all'efficienza e alla capacità di EfficientViT, EfficientViT-SAM offre un'accelerazione misurata di 48.9x su GPU A100 con TensorRT rispetto a SAM-ViT-H, senza compromettere le prestazioni. Il nostro codice e i modelli pre-addestrati sono disponibili all'indirizzo https://github.com/mit-han-lab/efficientvit.

English

We present EfficientViT-SAM, a new family of accelerated segment anything models. We retain SAM's lightweight prompt encoder and mask decoder while replacing the heavy image encoder with EfficientViT. For the training, we begin with the knowledge distillation from the SAM-ViT-H image encoder to EfficientViT. Subsequently, we conduct end-to-end training on the SA-1B dataset. Benefiting from EfficientViT's efficiency and capacity, EfficientViT-SAM delivers 48.9x measured TensorRT speedup on A100 GPU over SAM-ViT-H without sacrificing performance. Our code and pre-trained models are released at https://github.com/mit-han-lab/efficientvit.