[论文解读] RepViT-SAM: Towards Real-Time Segmenting Anything

Ao Wang, Hui Chen+3 more|arXiv (Cornell University)|Dec 10, 2023

Advanced Neural Network Applications被引用 4

一句话总结

RepViT-SAM 用轻量级的移动优化视觉变换器 RepViT-M2.3 替代了 Segment Anything Model (SAM) 中的重型 ViT 图像编码器，该模型采用高效的卷积架构和结构重参数化技术，在保持或超越 MobileSAM 的零样本分割性能的同时，实现接近 10 倍于 MobileSAM 的推理速度。

ABSTRACT

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 imes$ faster inference speed. The code and models are available at \url{https://github.com/THU-MIG/RepViT}.

研究动机与目标

• 在资源受限的移动设备上实现 Segment Anything Model (SAM) 的实时推理。
• 解决如 MobileSAM 等轻量级 SAM 变体中自注意力机制带来的高内存和计算开销问题。
• 提升轻量级 SAM 模型在多样化分割任务中的零样本迁移性能。
• 利用知识蒸馏技术，将大型 ViT-H 编码器的知识迁移至更小、更高效的基于 RepViT 的编码器。
• 为边缘部署的分割模型建立一个新的移动端优化基线，具备强大性能与低延迟特性。

提出的方法

• 将 SAM 中原始的 ViT-H 图像编码器替换为 RepViT-M2.3，这是一种轻量级、基于卷积的架构，结合了结构重参数化与深度可分离卷积。
• 采用解耦知识蒸馏：通过在特征图上使用均方误差（MSE）损失，直接从 ViT-H 编码器蒸馏 RepViT-M2.3 编码器。
• 将 RepViT 中最后一层下采样层的步长设为 1，以确保输出分辨率与 SAM 的提示引导掩码解码器兼容。
• 在蒸馏过程中预计算并缓存来自 ViT-H 编码器的图像嵌入，以加速训练过程，避免重复运行重型编码器。
• 使用与原始 SAM 相同的提示引导掩码解码器，保持推理流程的一致性。
• 采用 16×16 的前景点提示网格进行零样本边缘检测，并应用标准后处理操作，包括 Sobel 滤波和非极大值抑制。

Figure 1 : Mask predictions of SAM, MobileSAM, and RepViT-SAM with point prompts (top) and box prompts (bottom).

实验结果

研究问题

• RQ1 像 RepViT 这类轻量级卷积视觉变换器能否在保持强大零样本泛化能力的同时，实现在 SAM 中的实时推理？
• RQ2 RepViT-SAM 在多样化零样本分割基准上的性能与 MobileSAM 和 ViT-B-SAM 相比如何？
• RQ3 从大型 ViT-H 编码器蒸馏知识，能在多大程度上将高质量特征迁移至更小的基于 RepViT 的编码器？
• RQ4 RepViT 的纯卷积设计是否能在移动平台上实现比基于注意力机制的轻量级编码器更优的延迟与内存效率？
• RQ5 由于更优的特征蒸馏，RepViT-SAM 是否能在特定下游任务（如异常检测）中超越原始的 ViT-H-SAM？

主要发现

• 在 iPhone 12 上，RepViT-SAM 的推理延迟为 48.9ms，而 MobileSAM 因内存溢出（OOM）而失败，证明其在移动端部署的可行性。
• 在 Macbook M1 Pro 上，RepViT-SAM 的运行时间为 44.8ms，比 MobileSAM（482.2ms）快近 10 倍，比 ViT-H-SAM（6249.5ms）快 140 倍。
• 在零样本边缘检测任务中，RepViT-SAM 的 ODS 为 0.764，OIS 为 0.786，AP 为 0.773，优于 MobileSAM（0.756，0.768，0.746），且在 OIS 指标上与 ViT-H-SAM（0.768，0.786，0.794）持平。
• 在 COCO 数据集上的零样本实例分割任务中，RepViT-SAM 的平均 AP 达到 46.1，超过 MobileSAM（43.9）和 ViT-B-SAM（44.8），并接近 ViT-H-SAM（48.7）。
• 在 DAVIS 2017 数据集上的零样本视频对象分割任务中，RepViT-SAM 的平均 J&F 达到 73.5，优于 MobileSAM（71.1）和 ViT-B-SAM（71.3），较 MobileSAM 提升 2.4 个百分点。
• 在 MVTec-AD 数据集上的零样本异常检测任务中，RepViT-SAM 的 F1-像素得分达到 37.96，超过 ViT-H-SAM（37.65），并优于 MobileSAM（36.44），表明其在专业任务中具备更优的特征迁移能力。

Figure 2 : Visualization results of SAM, MobileSAM, and RepViT-SAM for zero-shot edge detection on BSDS500.