0. 简介
Segment Anything Model (SAM) 最近在各种计算机视觉任务上展现了令人瞩目的零样本迁移性能 。然而,其高昂的计算成本对于实际应用仍然具有挑战性。MobileSAM 提出通过使用蒸馏替换 SAM 中的重图像编码器,使用 TinyViT,从而显著降低了计算需求。然而,由于自注意力机制导致的内存和计算开销,其部署在资源受限的移动设备上仍面临挑战。
最近,RepViT 通过将 ViTs 的高效架构设计集成到 CNN 中,实现了移动设备上的最佳性能与延迟权衡。在这里,为了在移动设备上实现实时分割任何事物,作者遵循 [27] 的方法,用 RepViT 模型替换 SAM 中的重型图像编码器,最终得到了 RepViT-SAM 模型。相关的代码已近在Github上开源了。
1. 方法论
SAM [13]由一个重量级的基于ViT的图像编码器和一个轻量级的基于提示的掩模解码器组成。它庞大的图像编码器占据了大部分推理时间开销。因此,MobileSAM [27]建议将SAM中默认的ViT-H [6]图像编码器替换为轻量级的TinyViT [24]。TinyViT由四个逐渐降低分辨率的阶段组成。TinyViT的初始阶段由利用反向残差块[20]的卷积块组成。为了在模型开始时降低分辨率,采用了步幅为2的两个卷积块。类似地,相邻阶段之间也采用了步幅为2的卷积块进行空间降采样。为了使TinyViT的最终分辨率与原始SAM中ViT-H图像编码器的分辨率保持一致,MobileSAM将TinyViT中最后一个降采样卷积的步幅设置为1。此外,MobileSAM提出了解耦蒸馏策略,以有效训练轻量级图像编码器,其中TinyViT模型直接从原始SAM中的ViT-H蒸馏而来,而无需提示引导的掩模解码器。尽管MobileSAM显著降低了分割任何物体的计算需求,但在移动设备上部署仍然面临着相当大的挑战。如表1所示,由于其巨大的内存占用,MobileSAM无法在iPhone 12上运行。此外,在Macbook上,其处理单张图像的推理时间为494毫秒,表明有很大的改进空间。
最近,RepViT [21]通过重新审视ViT视角下CNN的高效设计,展示了在移动设备上性能和延迟的权衡方面的最新成果。RepViT采用了早期卷积[25]作为起始模块,即两个步幅为2的卷积进行4×降采样。它采用了RepViT块,由结构重参数化的深度卷积[4, 5]和前馈模块组成。相邻阶段之间采用了深度降采样模块,利用步幅为2的深度卷积和逐点卷积进行空间降采样和通道维度调制。此外,在所有阶段中都采用了交叉块方式的squeeze-and-excitation [8]层。RepViT在高分辨率视觉任务的延迟方面显示出了显著优势[21],这归功于其纯卷积架构。如表1所示,将ViT-H图像编码器替换为RepViT-M2.3模型后,RepViT-SAM在延迟方面与其他模型相比显著减少。在iPhone 12上,RepViT-SAM可以顺利进行模型推理。此外,在Macbook上,RepViT-SAM比MobileSAM快近10倍。
根据[27],我们通过直接蒸馏图像编码器RepViT-M2.3,即从原始SAM [13]中的ViT-H蒸馏,利用简单的均方误差损失来训练RepViT-SAM。与[27]类似,RepViT中最后一个降采样深度卷积的步幅设置为1,以使输出分辨率与原始SAM [13]中的提示引导掩模解码器兼容。
表1. RepViT-SAM与其他模型在延迟方面的比较。延迟(毫秒)是使用iPhone 12和Macbook M1 Pro上的Core ML工具,以标准分辨率[7] 1024×1024进行测量的。OOM表示内存不足。
表2. 零样本边缘检测的比较结果。粗体表示最佳,下划线表示次佳。
2. 实验
2.1 实施细节
RepViT-SAM在与[27]相同的设置下进行了8个时期的训练。与MobileSAM [27]一样,我们只使用SAM-1B数据集[13]中的1%数据。为了加快训练过程,我们在蒸馏阶段之前预先计算并保存了来自ViT-H图像编码器的图像嵌入,这消除了在蒸馏过程中运行ViT-H的前向过程的需要,就像[27]一样。我们在BSDS500 [1, 17]上评估了RepViT-SAM在零样本边缘检测、使用COCO [14]进行零样本实例分割、在野外基准[29](SegInW)上进行分割、使用DAVIS 2017 [18]/UVO v1.0 [23]进行零样本视频对象/实例分割、使用DUTS [22]进行零样本显著对象分割,以及使用MVTec-AD [2]进行零样本异常检测的性能,遵循[3, 9, 12, 13, 19]。
点击经典文献阅读之--RepViT-SAM(利用语义分割提高NDT地图压缩和描述能力的框架) - 古月居可查看全文