SAM-Lightening: A Lightweight Segment Anything Model with Dilated Flash Attention to Achieve 30 times Acceleration

• URL: http://arxiv.org/abs/2403.09195v2
• Date: Mon, 18 Mar 2024 02:30:17 GMT
• Title: SAM-Lightening: A Lightweight Segment Anything Model with Dilated Flash Attention to Achieve 30 times Acceleration
• Authors: Yanfei Song, Bangzheng Pu, Peng Wang, Hongxu Jiang, Dong Dong, Yongxiang Cao, Yiqing Shen,
• Abstract summary: Segment Anything Model (SAM) has garnered significant attention in segmentation tasks due to their zero-shot generalization ability. We introduce SAM-Lightening, a variant of SAM, that features a re-engineered attention mechanism, termed Dilated Flash Attention. Experiments on COCO and LVIS reveal that SAM-Lightening significantly outperforms the state-of-the-art methods in both run-time efficiency and segmentation accuracy.
• Score: 6.515075311704396
• License: http://arxiv.org/licenses/nonexclusive-distrib/1.0/
• Abstract: Segment Anything Model (SAM) has garnered significant attention in segmentation tasks due to their zero-shot generalization ability. However, a broader application of SAMs to real-world practice has been restricted by their low inference speed and high computational memory demands, which mainly stem from the attention mechanism. Existing work concentrated on optimizing the encoder, yet has not adequately addressed the inefficiency of the attention mechanism itself, even when distilled to a smaller model, which thus leaves space for further improvement. In response, we introduce SAM-Lightening, a variant of SAM, that features a re-engineered attention mechanism, termed Dilated Flash Attention. It not only facilitates higher parallelism, enhancing processing efficiency but also retains compatibility with the existing FlashAttention. Correspondingly, we propose a progressive distillation to enable an efficient knowledge transfer from the vanilla SAM without costly training from scratch. Experiments on COCO and LVIS reveal that SAM-Lightening significantly outperforms the state-of-the-art methods in both run-time efficiency and segmentation accuracy. Specifically, it can achieve an inference speed of 7 milliseconds (ms) per image, for images of size 1024*1024 pixels, which is 30.1 times faster than the vanilla SAM and 2.1 times than the state-of-the-art. Moreover, it takes only 244MB memory, which is 3.5\% of the vanilla SAM. The code and weights are available at https://anonymous.4open.science/r/SAM-LIGHTENING-BC25/.

Related papers

• Lite-SAM Is Actually What You Need for Segment Everything [4.696541976769272]
  Lite-SAM is an efficient end-to-end solution for the SegEvery task. Lite-SAM is composed of four main components: a streamlined CNN-Transformer hybrid encoder (LiteViT), an automated prompt proposal network (AutoPPN)
  arXiv  Detail & Related papers  (2024-07-12T03:28:46Z)
• HRSAM: Efficiently Segment Anything in High-Resolution Images [59.537068118473066]
  This study proposes HRSAM that integrates Flash Attention and incorporates Plain, Shifted and newly proposed Cycle-scan Window attention to address these issues. The cycle-scan window attention adopts the recently developed State Space Models (SSMs) to ensure global information exchange with minimal computational overhead. Experiments on the high-precision segmentation datasets HQSeg44K and DAVIS show that high-resolution inputs enable the SAM-distilled HRSAM models to outperform the teacher model.
  arXiv  Detail & Related papers  (2024-07-02T09:51:56Z)
• FastSAM3D: An Efficient Segment Anything Model for 3D Volumetric Medical Images [7.2993352400518035]
  We present FastSAM3D, which accelerates SAM inference to 8 milliseconds per 128*128*128 3D volumetric image on an NVIDIA A100 GPU. FastSAM3D achieves a remarkable speedup of 527.38x compared to 2D SAMs and 8.75x compared to 3D SAMs on the same volumes without significant performance decline.
  arXiv  Detail & Related papers  (2024-03-14T19:29:44Z)
• TinySAM: Pushing the Envelope for Efficient Segment Anything Model [76.21007576954035]
  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 hard prompt sampling and hard mask weighting 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.
  arXiv  Detail & Related papers  (2023-12-21T12:26:11Z)
• EdgeSAM: Prompt-In-the-Loop Distillation for On-Device Deployment of SAM [71.868623296582]
  EdgeSAM is an accelerated variant of the Segment Anything Model (SAM) Our approach involves distilling the original ViT-based SAM image encoder into a purely CNN-based architecture. It is the first SAM variant that can run at over 30 FPS on an iPhone 14.
  arXiv  Detail & Related papers  (2023-12-11T18:59:52Z)
• Stable Segment Anything Model [79.9005670886038]
  The Segment Anything Model (SAM) achieves remarkable promptable segmentation given high-quality prompts. This paper presents the first comprehensive analysis on SAM's segmentation stability across a diverse spectrum of prompt qualities. Our solution, termed Stable-SAM, offers several advantages: 1) improved SAM's segmentation stability across a wide range of prompt qualities, while 2) retaining SAM's powerful promptable segmentation efficiency and generality.
  arXiv  Detail & Related papers  (2023-11-27T12:51:42Z)
• Faster Segment Anything: Towards Lightweight SAM for Mobile Applications [47.177751899636164]
  This work aims to make Segment Anything Model (SAM) mobile-friendly by replacing the heavyweight image encoder with a lightweight one. We distill the knowledge from the heavy image encoder to a lightweight image encoder, which can be automatically compatible with the mask decoder in the original SAM. The resulting lightweight SAM is termed MobileSAM which is more than 60 times smaller yet performs on par with the original SAM.
  arXiv  Detail & Related papers  (2023-06-25T16:37:25Z)
• Towards Efficient and Scalable Sharpness-Aware Minimization [81.22779501753695]
  We propose a novel algorithm LookSAM that only periodically calculates the inner gradient ascent. LookSAM achieves similar accuracy gains to SAM while being tremendously faster. We are the first to successfully scale up the batch size when training Vision Transformers (ViTs)
  arXiv  Detail & Related papers  (2022-03-05T11:53:37Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.