点击下方卡片，关注 「AI视界引擎」 公众号
Segment Anything Model (SAM) 最近在各种计算机视觉任务上展现了令人瞩目的零样本迁移性能 [3, 9, 19, 26, 28]。然而，其高昂的计算成本对于实际应用仍然具有挑战性。MobileSAM 提出通过使用蒸馏替换 SAM 中的重图像编码器，使用 TinyViT，从而显著降低了计算需求。然而，由于自注意力机制导致的内存和计算开销，其部署在资源受限的移动设备上仍面临挑战。
最近，RepViT 通过将 ViTs 的高效架构设计集成到 CNN 中，实现了移动设备上的最佳性能与延迟权衡。在这里，为了在移动设备上实现实时分割任何事物，作者遵循 [27] 的方法，用 RepViT 模型替换 SAM 中的重型图像编码器，最终得到了 RepViT-SAM 模型。
大量实验表明，RepViT-SAM 可以在零样本迁移能力方面显著优于 MobileSAM，同时推理速度快近 10 倍。
1 Methodology
SAM 主要由一个基于 Transformer（Transformer）的重量级图像编码器和一个轻量级的提示引导掩码解码器组成。它的巨大图像编码器占据了推理时间开销的大部分。因此，MobileSAM 建议用一个轻量级的图像编码器替换 SAM 中的默认 ViT-H 图像编码器，即 TinyViT 。TinyViT 包括四个阶段，逐渐降低分辨率。TinyViT 的初始阶段由使用倒置残差模块 的卷积块组成。为了在模型开始时降低分辨率，使用了两个步长为 2 的卷积块。类似地，使用步长为 2 的卷积块在相邻阶段之间进行空间降采样。为了使 TinyViT 的最终分辨率与原始 SAM 的 ViT-H 图像编码器对齐，MobileSAM 将 TinyViT 最后一个降采样卷积的步长设置为 1。
此外，MobileSAM 提出了解耦蒸馏策略来有效地训练轻量级图像编码器，其中 TinyViT 模型直接从原始 SAM 中的 ViT-H 解耦蒸馏，无需提示引导掩码解码器。尽管 MobileSAM 显著降低了分割任何事物的计算需求，但将其部署在移动设备上仍然面临相当大的挑战。如表 1 所示，由于其较大的内存占用，MobileSAM 无法在 iPhone 12 上运行。此外，在 Macbook 上，处理单张图像的推理时间为 494ms，表明仍有改进空间。
最近，RepViT 通过从 ViT（Vision Transformer）的角度重新审视 CNN（卷积神经网络）的高效设计，在移动设备上实现了最先进的性能与延迟权衡。RepViT 采用早期卷积 作为Stem，使用步长为 2 的两个卷积进行 降采样处理输入。它采用了 RepViT 模块，该模块由结构正则化的深度卷积 和前馈模块组成。相邻阶段之间使用深度降采样模块，该模块利用步长为 2 的深度卷积和点卷积分别进行空间降采样和通道维度调制。
此外，SE层以交叉块方式在所有阶段中使用。由于其纯卷积架构，RepViT 在高分辨率视觉任务 的延迟方面具有显著优势。如表 1 所示，在用 RepViT-M2.3 替换 ViT-H 图像编码器后，RepViT-SAM 与其他方法相比延迟显著降低。在 iPhone 12 上，RepViT-SAM 可以流畅地进行模型推理。此外，在 Macbook 上，RepViT-SAM 比 MobileSAM 快近 10 倍。
遵循 [27] 的方法，作者通过直接从原始 SAM 的 ViT-H 图像编码器中蒸馏 RepViT-M2.3 模型来进行 RepViT-SAM 的训练，利用简单均方误差（MSE）损失。与 [27] 一样，RepViT 在最后一个降采样深度卷积中的步长设置为 1，以便使输出分辨率与原始 SAM 中的提示引导掩码解码器兼容。
2 Experiments
Implementation details
RepViT-SAM 在与 [27] 相同的环境下进行了 8 轮训练。与 MobileSAM 一样，作者在 SAM-1B 数据集 的 1% 数据上进行训练。为了加快训练过程，作者在蒸馏阶段之前预计算并保存了 ViT-H 图像编码器的图像嵌入，这样在蒸馏过程中就不需要运行 ViT-H 的前向过程，就像 [27] 一样。作者在 BSDS500上评估 RepViT-SAM 的零样本边缘检测性能，在 COCO 上评估零样本实例分割性能，以及在野外观测基准 (SegInW) 上的分割性能，零样本视频目标/实例分割使用 DAVIS 2017/UVO v1.0，零样本显著目标分割使用 DUTS，以及零样本异常检测使用 MVTec-AD，遵循 [3, 9, 12, 13, 19] 的方法。
Zero-shot edge detection
遵循 [13] 的方法，作者使用一个 1616 的规则网格作为提示，为模型提供 foreground 点。作者使用 NMS 去除冗余掩码。此外，作者使用 Sobel 滤波器对掩码的不阈值概率图进行处理，并采用标准的轻量级后处理，包括边缘 NMS，如 [13] 所述。
如表 2 所示，作者的 RepViT-SAM 在所有指标上均优于 MobileSAM 和 ViT-B-SAM。与具有超过 615M 参数的最大的 SAM 模型 ViT-H-SAM 相比，作者的小 RepViT-SAM 在 ODS 和 OIS 方面可以获得可比的性能。
Zero-shot instance segmentation
作者利用最先进的 H-Deformable-DETR 作为目标检测器，并使用 Swin-L 作为提示，向模型提供其输出框，就像 [12, 13] 一样。
如表 3 所示，作者的 RepViT-SAM 分别比 MobileSAM 和 ViT-B-SAM 在 1.7 和 1.9 个 AP 上显著优于它们。此外，可以观察到，通过解耦蒸馏策略获得的模型（即 MobileSAM 和 RepViT-SAM），在大型物体上的性能通常优于通过耦合蒸馏获得的模型（即 ViT-B-SAM），但在小型物体上表现出相对较差的性能。这种现象可能归因于解耦蒸馏策略使宏观视觉特征的转移更多，但在细粒度特征方面有所不足。
Segmentation in the wild benchmark
作者遵循 [12] 的方法，将 Grounding-DINO 作为Box提示，在零样本轨道上评估模型。如表 3 所示，作者的 RepViT-SAM 分别比 MobileSAM 和 ViT-B-SAM 在 2.2 和 1.3 个平均 AP 上显著优于它们。这很好地表明了 RepViT-SAM 在各种下游分割任务上的强迁移能力。
Zero-shot video object/instance segmentation
作者利用 SAM-PT 在 DAVIS 2017 和 UVO 数据集的验证分片中评估 RepViT-SAM 的性能。作者简单地将 SAM-PT 中的原始 SAM 用作者的 RepViT-SAM 替换，使用 CoTracker 作为点跟踪器，但不需要重初始化策略。
如表 4 所示，在 DAVIS 2017 上，作者的 RepViT-SAM 获得了平均 分数为 73.5，显著优于 MobileSAM 和 ViT-B-SAM，分别提高了 2.4 和 2.2。此外，在 UVO 上，RepViT-SAM 分别比 MobileSAM 和 ViT-B-SAM 具有 2.6 和 6.2 个显著的优点。这些结果表明，与轻量级方法相比，RepViT-SAM 的优越性不仅限于图像分割任务，还可以扩展到视频分割任务。
Zero-shot salient object segmentation
作者遵循 [9] 的方法，首先使用 RepViT-SAM 为给定图像生成 N 个潜在目标掩码。然后，作者通过评估其与真实值的对齐来选择最合适的掩码。采用交点与 Union（IoU）分数作为对齐指标。遵循 [9] 的方法，作者报告平均绝对误差（MAE）分数。请注意，较低的 MAE 表示更好的分割结果。如表 5 所示，作者的 RepViT-SAM 在 MAE 分数方面优于 MobileSAM 和 ViT-B-SAM。
Zero-shot anomaly detection
作者利用 SAA+ 在 MVTec-AD 基准数据集 上评估 RepViT-SAM 的性能。作者简单地将 SAA+ 中的原始 SAM 用作者的 RepViT-SAM 替换。作者遵循 [3] 报告最大 F1-像素（），它衡量了在最佳阈值下像素级分割的 F1-score。如表 5 所示，作者的 RepViT-SAM 在 方面取得了最佳性能，甚至超过了 ViT-H-SAM。这表明对于某些特定特定的下游任务，解耦蒸馏模型可以达到比原始 SAM 更好的性能，同时延迟更低，展示了 RepViT-SAM 的前景。
3 Visualization
在图 1 中，作者展示了使用点提示和Box提示预测的 SAM、MobileSAM 和 RepViT-SAM 的掩码。它表明 RepViT-SAM 可以在各种场景下生成类似于 SAM 的高质量掩码预测，而 MobileSAM 可能在准确预测掩码方面遇到挑战。
此外，作者在图 2 中呈现了零样本边缘检测的定性比较结果。RepViT-SAM 有效地生成了合理的边缘图，实现了与 SAM 相当的表现。相比之下，MobileSAM 在包含复杂细节的地区可能难以生成准确的边缘图。
4 Conclusion
总之，作者的 RepViT-SAM 模型在保持对各种下游任务 impressive 的迁移性能的同时，展示了出色的效率。作者希望作者的 RepViT-SAM 模型可以在边缘部署中作为 SAM 的强大 Baseline ，展示其在实际应用中的潜力。
参考
[1].RepViT-SAM: Towards Real-Time Segmenting Anything.
点击上方卡片，关注 「AI视界引擎」 公众号