比SAM快30倍的SAM-Lightening来了！推理一张图仅需7毫秒（北航）

今日应用

熊猫云简历

熊猫办公 - 时下流行的精品简历模板，免费在线制作，助您轻松展现职业精彩，一站式职业形象打造！

今日话题

比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与剪枝和量化技术整合。

文章来源

原文地址: 点我阅读全文
原文作者: 极市平台