RepViT-SAM: Towards Real-Time Segmenting Anything

📄 arXiv: 2312.05760v2 📥 PDF

作者: Ao Wang, Hui Chen, Zijia Lin, Jungong Han, Guiguang Ding

分类: cs.CV

发布日期: 2023-12-10 (更新: 2024-02-29)

备注: Technical report of RepViT+SAM in our CVPR 2024 work. Project page: https://jameslahm.github.io/repvit-sam/

🔗 代码/项目: GITHUB

💡 一句话要点

提出RepViT-SAM以解决移动设备实时分割问题

🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture) 支柱九：具身大模型 (Embodied Foundation Models)

关键词: 实时分割 移动设备 计算机视觉 自注意力机制 模型蒸馏 高效架构 零-shot迁移

📋 核心要点

  1. 现有的SAM模型在计算成本上过于高昂，限制了其在实际应用中的广泛使用。
  2. 本文提出RepViT-SAM，通过用RepViT替代SAM中的重型图像编码器，显著降低计算需求。
  3. 实验结果显示，RepViT-SAM在零-shot迁移能力上优于MobileSAM，并实现了近10倍的推理速度提升。

📝 摘要（中文）

Segment Anything Model (SAM) 最近在各种计算机视觉任务中展现了令人印象深刻的零-shot迁移性能。然而，其高昂的计算成本在实际应用中仍然是一个挑战。MobileSAM通过蒸馏技术用TinyViT替代了SAM中的重型图像编码器，从而显著降低了计算需求。然而，由于自注意力机制带来的内存和计算开销，MobileSAM在资源受限的移动设备上的部署仍面临困难。最近，RepViT通过将高效的ViT架构设计融入CNN，达到了移动设备上性能与延迟的最佳平衡。为实现移动设备上的实时分割，本文在MobileSAM的基础上，将SAM中的重型图像编码器替换为RepViT模型，最终形成RepViT-SAM模型。大量实验表明，RepViT-SAM在零-shot迁移能力上显著优于MobileSAM，并且推理速度几乎提高了10倍。

🔬 方法详解

问题定义：本文旨在解决现有SAM模型在移动设备上实时分割的计算成本过高的问题，尤其是自注意力机制带来的内存和计算开销。

核心思路：通过将SAM中的重型图像编码器替换为RepViT模型，利用RepViT的高效架构设计，降低计算需求并提升推理速度。

技术框架：RepViT-SAM模型主要由图像编码器、分割头和后处理模块组成。图像编码器负责特征提取，分割头进行像素级分类，后处理模块优化输出结果。

关键创新：RepViT-SAM在保持高性能的同时，显著降低了计算复杂度，尤其是在移动设备上实现了实时分割能力，这是与现有方法的本质区别。

关键设计：在模型设计中，采用了高效的卷积神经网络结构，结合了自注意力机制的优势，并对损失函数进行了优化，以提高模型的训练效率和推理速度。

📊 实验亮点

实验结果表明，RepViT-SAM在零-shot迁移能力上显著优于MobileSAM，推理速度提升接近10倍。这一性能提升使得RepViT-SAM在移动设备上的应用成为可能，极大地推动了实时分割技术的发展。

🎯 应用场景

RepViT-SAM模型具有广泛的应用潜力，特别是在移动设备上进行实时图像分割的场景，如智能手机应用、无人机图像处理和增强现实等领域。其高效的计算性能和快速的推理速度将推动相关技术的普及和发展。

📄 摘要（原文）

Segment Anything Model (SAM) has shown impressive zero-shot transfer performance for various computer vision tasks recently. However, its heavy computation costs remain daunting for practical applications. MobileSAM proposes to replace the heavyweight image encoder in SAM with TinyViT by employing distillation, which results in a significant reduction in computational requirements. However, its deployment on resource-constrained mobile devices still encounters challenges due to the substantial memory and computational overhead caused by self-attention mechanisms. Recently, RepViT achieves the state-of-the-art performance and latency trade-off on mobile devices by incorporating efficient architectural designs of ViTs into CNNs. Here, to achieve real-time segmenting anything on mobile devices, following MobileSAM, we replace the heavyweight image encoder in SAM with RepViT model, ending up with the RepViT-SAM model. Extensive experiments show that RepViT-SAM can enjoy significantly better zero-shot transfer capability than MobileSAM, along with nearly $10\times$ faster inference speed. The code and models are available at \url{https://github.com/THU-MIG/RepViT}.