SAM轻量级改进RepViT-SAM论文解读RepViT-SAM: Towards Real-Time Segmenting Anything

现已总结SAM多方面相关的论文解读，具体请参考该专栏的置顶目录篇

一、总结

1. 简介

发表时间：2023年12月10日
论文：[2312.05760] RepViT-SAM: Towards Real-Time Segmenting Anything (arxiv.org)https://arxiv.org/abs/2312.05760代码：
github.comhttps://github.com/THU-MIG/RepViT

2. 摘要

        由于RepViT通过将高效的ViT架构设计融入CNN，在移动设备上实现了最先进的性能和延迟权衡。因此，RepViT-SAM选择用RepViT模型替换了SAM中的重量级图像编码器，延续MobileSAM 的方法，采用解耦蒸馏的方式，直接从SAM中的ViT-H提取图像编码器 RepViT-M2.3，利用简单的均方误差损失来训练RepViT-SAM。大量实验表明，RepViT-SAM具有比MobileSAM明显更好的零样本迁移能力，以及近10倍的推理速度。RepViT-SAM在保持各种下游任务传输性能的同时表现出出色的效率。

3. 前言

3.1 SAM

SAM由一个重量级的基于vit的图像编码器和一个轻量级的提示引导掩码解码器组成。其巨大的图像编码器占据了大部分的推理时间开销。

3.2 MobileSAM

MobileSAM建议将SAM中默认的ViT-H图像编码器替换为轻量级的TinyViT。TinyViT由四个阶段组成，这些阶段逐渐降低分辨率。TinyViT的初始阶段由利用反向残差块的卷积块组成。为了在模型开始时对分辨率进行下采样，使用了两个步幅为2的卷积块。类似地,采用步长为2的卷积块在相邻阶段之间进行空间下采样。为了使TinyViT的最终分辨率与原始SAM的ViT-H图像编码器的分辨率保持一致，MobileSAM将TinyViT中最后一次下采样卷积的步幅设置为1。此外，MobileSAM还提出了解耦蒸馏策略，以有效地训练轻量级图像编码器，其中TinyViT模型直接从原始SAM中的ViT-H中提取，而不需要提示引导的掩模解码器。尽管大大减少了分割任何东西的计算需求，但在移动设备上部署MobileSAM仍然面临相当大的挑战。

二、RepViT-SAM模型结构

1.RepViT

RepViT通过从ViT的角度重新审视cnn的高效设计，展示了移动设备上最先进的性能和延迟权衡。RepViT采用早期卷积作为主干，即两个步长为2的卷积，对输入进行4倍的下采样。它采用由结构重参数化深度卷积和前馈模块组成的RepViT块。在相邻阶段之间采用深度下采样模块，其利用步幅为2的深度卷积和点向卷积分别执行空间下采样和信道维度调制。此外，挤压激励层用于所有阶段以跨块方式进行。由于其纯卷积架构，RepViT在高分辨率视觉任务的延迟方面显示出巨大的优势。

2.解耦蒸馏

通过直接从SAM中的ViT-H提取图像编码器 RepViT-M2.3，利用简单的均方误差损失来训练RepViT-SAM。与MobileSAM一样，RepViT中最后一个降采样深度卷积的步幅设置为1，以使输出分辨率与原始SAM中的提示引导掩模解码器兼容。

3.后处理

使用16×16前景点的规则网格作为模型的提示。使用NMS来删除冗余的掩码。此外，对掩码的非阈值概率图采用Sobel滤波器，并进行标准的轻量级后处理，包括边缘NMS。

三、实验流程

 1. 推理速度对比

如下表所示，延迟（毫秒）是使用iPhone 12和Macbook M1 Pro上的Core ML工具，以标准分辨率[7] 1024×1024进行测量的。其中Platform表示使用平台，Image encoder表示图像编码器，Mask decoder表示掩码解码器，OOM表示内存不足。
MobileSAM由于占用大量内存，无法在iPhone 12上运行，在Macbook上，它处理单幅图像的推理时间为494毫秒，表明有很大的改进空间。
RepViT-SAM与其他模型相比，延迟明显降低，在iPhone 12上，RepViT-SAM可以平滑地执行模型推理，此外，在Macbook上，RepViT-SAM比MobileSAM快近10倍。

2.各方面性能对比

RepViT-SAM在与MobileSAM相同的设置下训练了8个epoch，只使用了SAM-1B数据集中1%的数据。为了加快训练过程，在蒸馏阶段之前预先计算并保存来自ViT-H图像编码器的图像嵌入，这样就不需要在蒸馏过程中运行ViT-H的前向过程。

2.1 zero-shot边缘检测

使用BSDS500分割数据集评估性能，其中ODS、OIS、AP均为边缘检测的评价指标。
ODS(optimal dataset scale)：称为全局最佳、固定轮廓阈值 数据集固定比例 、检测指标数据集尺度上最优，简单说就是为所有图像设置同样的阈值，即选取一个固定的阈值η应用于所有图片， 使得整个数据集上的F-score最大；
OIS(optimal image scale)：称为单图最佳、每幅图像的最佳阈值、图片尺度上最优，简单说就是在每一张图片上均选取不同使得该图片F-score最大的阈值η。
AP：平均准确率，是PR曲线的积分（即PR曲线下方的面积）。由于PR曲线很难积分，通常在PR曲线上采样求均值。
如下表所示，粗体表示最佳，下划线表示次佳。RepViT-SAM在所有指标上都优于MobileSAM和ViT-B-SAM。与最大的超过615M个参数的SAM模型ViT-H-SAM相比，小型RepViT-SAM在ODS和OIS方面可以获得相当的性能。

2.2 zero-shot实例分割

使用COCO数据集评估RepViT-SAM性能，我们利用最先进的H-Deformable-DETR与Swin-L作为目标检测器，并使用输出框提示模型。
如下表zero-shot instance segmention所示，粗体表示最佳，下划线表示次佳。RepViT-SAM分别比MobileSAM和ViT-B-SAM高出1.7和1.9 AP。此外，可以观察到解耦蒸馏策略得到的模型，即MobileSAM和RepViTSAM，在大对象上通常优于耦合蒸馏得到的模型，即ViT-B-SAM，但在小对象上表现相对较差。这种现象可能是由于解耦蒸馏策略能够更多地转移宏观层次的视觉特征，但在细粒度特征上却有所欠缺。

2.3 野生基准的分割

按照HO-SAM将Grounding-DINO设置为框提示，在 zero-shot轨道上对模型进行评估。
如下表SegInW所示，RepViTSAM比MobileSAM和ViT-B-SAM的平均AP分别高出2.2和1.3。这很好地说明了RepViTSAM在不同的下游分割任务中具有很强的可移植性。

2.4 zero-shot视频对象/实例分割

利用SAM-PT来评估RepViTSAM在DAVIS 2017和UVO两个数据集的验证分割上的性能。将SAM-PT中的原始SAM替换为RepViTSAM，使用CoTracker作为点跟踪器，而不使用重新初始化策略。
如下表所示，在DAVIS 2017上，RepViTSAM获得了73.5的平均J&F分数，分别显著优于MobileSAM和ViT-B-SAM 2.4和2.2。此外，在UVO上，RepViTSAM分别以2.6和6.2的显著优势优于MobileSAM和ViT-B-SAM。这些结果表明，RepViTSAM与其他轻量级SAM相比的优势不仅限于图像分割任务，还可以扩展到视频分割任务。

2.5 zero-shot显著目标分割

使用DUTS数据集评估RepViT-SAM性能，首先利用RepViTSAM为给定图像生成N个潜在的对象掩码。随后，通过评估其与地面真相的一致性来选择最合适的掩模。采用IoU (intersection over union)分数作为对齐度量。
如下表所示，较低的MAE(平均绝对误差)评分表示更好的分割结果，RepViTSAM在MAE得分方面比MobileSAM和ViT-B-SAM更具优势。

2.6 zero-shot异常检测

利用SAA+在MVTec-AD基准数据集上评估RepViTSAM的性能。只需将SAA+中的原始SAM替换为RepViTSAM。
如下表所示，max-F1-pixel (Fp)，它量化了在最佳阈值下逐像素分割的F1分数。RepViTSAM在Fp方面达到了最佳性能，甚至超过了ViT-H-SAM。这表明，对于某些特定的下游任务，解耦蒸馏模型可以达到比原始SAM更好的性能，并且延迟更低，显示了RepViTSAM的应用前景。

3. 可视化结果

下图中，用点和框提示展示了SAM、MobileSAM和RepViTSAM的预测掩码。这表明RepViTSAM可以在各种场景中生成类似于SAM的高质量掩码预测，在这些场景中，MobileSAM可能会在准确预测掩码方面遇到挑战。
下图中，给出了零镜头边缘检测的定性比较结果。RepViTSAM有效地生成合理的边缘映射，实现与SAM相当的性能。相比之下，MobileSAM可能难以在包含复杂细节的区域生成精确的边缘地图。