EfficientViT-SAM: Modelo Segment Anything Acelerado Sem Perda de Desempenho

EfficientViT-SAM: Accelerated Segment Anything Model Without Performance Loss

February 7, 2024

Autores: Zhuoyang Zhang, Han Cai, Song Han

cs.AI

Resumo

Apresentamos o EfficientViT-SAM, uma nova família de modelos acelerados de segmentação universal. Mantemos o codificador de prompts leve e o decodificador de máscaras do SAM, enquanto substituímos o pesado codificador de imagens pelo EfficientViT. Para o treinamento, iniciamos com a destilação de conhecimento do codificador de imagens SAM-ViT-H para o EfficientViT. Posteriormente, realizamos o treinamento de ponta a ponta no conjunto de dados SA-1B. Beneficiando-se da eficiência e capacidade do EfficientViT, o EfficientViT-SAM alcança um aumento de velocidade de 48,9x medido com TensorRT na GPU A100 em comparação com o SAM-ViT-H, sem sacrificar o desempenho. Nosso código e modelos pré-treinados estão disponíveis em 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.