今日应用
今日话题
比SAM快30倍的SAM-Lightening来了!推理一张图仅需7毫秒(北航)
重点标签 图像分割、SAM-Lightening、Dilated Flash Attention、动态分层蒸馏、内存优化
文章摘要
摘要:
本文介绍了一种名为SAM-Lightening的新型图像分割模型,它是Segment Anything Model(SAM)的变体,旨在解决原始SAM在实际应用中的低推理速度和高计算内存需求的问题。SAM-Lightening通过重新设计的注意力机制——Dilated Flash Attention,提高了并行性和处理效率,同时保留了与现有Flash Attention的兼容性。研究者还提出了一种渐进式蒸馏方法,有效传递了普通SAM的知识,无需从头开始训练。实验结果显示,SAM-Lightening在COCO和LVIS数据集上的性能明显优于现有技术,对于1024×1024像素的图像,推理速度可达7毫秒,比普通SAM快30.1倍,内存需求仅为244MB,是普通SAM的3.5%。
介绍:
SAM因其zero-shot泛化能力受到关注,但其在实际应用中的使用受限于低推理速度和高内存需求。以往的工作主要集中在优化编码器上,但未充分解决注意力机制的低效率问题。SAM-Lightening通过引入Dilated Flash Attention和动态分层蒸馏(DLD)方法,显著降低了计算复杂性,同时保持了分割性能。
相关工作:
SAM的图像编码器是模型中参数最多、处理时间最长的部分。FastSAM、MobileSAM、SAMFast和EfficientSAM等方法都试图通过不同方式优化SAM,但存在精度下降、训练过程复杂等问题。FlashAttention机制通过策略性切片和重计算技术,显著减少了内存读写需求。知识蒸馏技术旨在将复杂模型的知识转移到简单模型,以减少计算占用和模型大小。
方法:
Dilated Flash Attention通过分割和稀疏化输入、并行处理和输出重组,提高了计算效率。动态分层蒸馏(DLD)方法通过逐渐分配特征权重,实现了从SAM到轻量级模型的知识转移,避免了从头训练的高成本。
实验:
实验使用SA-1B数据集的1%进行蒸馏和微调,结果表明SAM-Lightening在推理延迟和内存使用方面优于所有对手,且在边界框和点提示模式下的分割性能与原始SAM相当。消融研究表明,FlashAttention的适用性取决于模型的输入尺寸和硬件配置。
结论:
SAM-Lightening通过优化图像编码器和应用动态层次蒸馏,显著提高了推理速度和内存优化,使其更适合部署在资源受限的设备上。未来的研究方向包括将SAM-Lightening与剪枝和量化技术整合。
文章来源
原文地址: 点我阅读全文
原文作者: 极市平台