EfficientViT-SAM: Beschleunigtes Segment-Anything-Modell ohne Leistungseinbußen

EfficientViT-SAM: Accelerated Segment Anything Model Without Performance Loss

February 7, 2024

Autoren: Zhuoyang Zhang, Han Cai, Song Han

cs.AI

Zusammenfassung

Wir stellen EfficientViT-SAM vor, eine neue Familie von beschleunigten Segment-Anything-Modellen. Wir behalten den leichten Prompt-Encoder und Mask-Decoder von SAM bei, ersetzen jedoch den rechenintensiven Bild-Encoder durch EfficientViT. Für das Training beginnen wir mit der Wissensdestillation vom SAM-ViT-H-Bild-Encoder zu EfficientViT. Anschließend führen wir ein End-to-End-Training auf dem SA-1B-Datensatz durch. Dank der Effizienz und Kapazität von EfficientViT bietet EfficientViT-SAM eine gemessene TensorRT-Beschleunigung von 48,9x auf der A100-GPU im Vergleich zu SAM-ViT-H, ohne dabei an Leistung einzubüßen. Unser Code und die vortrainierten Modelle sind unter https://github.com/mit-han-lab/efficientvit verfügbar.

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.