SAM-Lightening：一种具有扩张式快速注意力的轻量级分割一切模型，实现30倍加速
计算机视觉与模式识别
2024-03-19 v2
摘要
分割一切模型（SAM）因其零样本泛化能力在分割任务中引起了广泛关注。然而，SAM在现实世界实践中的更广泛应用受到其低推理速度和高计算内存需求的限制，这主要源于注意力机制。现有工作集中于优化编码器，但尚未充分解决注意力机制本身的低效问题，即使将其蒸馏为更小的模型后依然如此，这留下了进一步改进的空间。为此，我们引入了SAM-Lightening，这是SAM的一个变体，其特点是重新设计的注意力机制，称为扩张式快速注意力（Dilated Flash Attention）。它不仅有助于实现更高的并行性，提高处理效率，而且还保留了与现有FlashAttention的兼容性。相应地，我们提出了一种渐进式蒸馏方法，能够从原始SAM进行高效的知识迁移，而无需从头开始进行昂贵的训练。在COCO和LVIS数据集上的实验表明，SAM-Lightening在运行效率和分割精度方面均显著优于最先进的方法。具体来说，对于尺寸为1024*1024像素的图像，它可以实现每张图像7毫秒（ms）的推理速度，比原始SAM快30.1倍，比最先进的方法快2.1倍。此外，它仅占用244MB内存，是原始SAM的3.5%。代码和权重可在https://anonymous.4open.science/r/SAM-LIGHTENING-BC25/获取。
引用
@article{arxiv.2403.09195,
title = {SAM-Lightening: A Lightweight Segment Anything Model with Dilated Flash Attention to Achieve 30 times Acceleration},
author = {Yanfei Song and Bangzheng Pu and Peng Wang and Hongxu Jiang and Dong Dong and Yongxiang Cao and Yiqing Shen},
journal= {arXiv preprint arXiv:2403.09195},
year = {2024}
}