EfficientSAM：小模型也能万物分割！Meta改进SAM，参数仅为原版5%

作者：weixin_40725706 | 2024-06-26 04:20:48

踩

efficientsam

本文来源机器之心编辑：陈萍、蛋酱

对于 2023 年的计算机视觉领域来说，「分割一切」（Segment Anything Model）是备受关注的一项研究进展。

Meta四月份发布的「分割一切模型（SAM）」效果，它能很好地自动分割图像中的所有内容

Segment Anything 的关键特征是基于提示的视觉 Transformer（ViT）模型，该模型是在一个包含来自 1100 万张图像的超过 10 亿个掩码的视觉数据集 SA-1B 上训练的，可以分割给定图像上的任何目标。这种能力使得 SAM 成为视觉领域的基础模型，并在超出视觉之外的领域也能产生应用价值。

尽管有上述优点，但由于 SAM 中的 ViT-H 图像编码器有 632M 个参数（基于提示的解码器只需要 387M 个参数），因此实际使用 SAM 执行任何分割任务的计算和内存成本都很高，这对实时应用来说具有挑战性。后续，研究者们也提出了一些改进策略：将默认 ViT-H 图像编码器中的知识提炼到一个微小的 ViT 图像编码器中，或者使用基于 CNN 的实时架构降低用于 Segment Anything 任务的计算成本。

在最近的一项研究中，Meta 研究者提出了另外一种改进思路 —— 利用 SAM 的掩码图像预训练 (SAMI)。这是通过利用 MAE 预训练方法和 SAM 模型实现的，以获得高质量的预训练 ViT 编码器。

论文链接：https://arxiv.org/pdf/2312.00863.pdf
论文主页：https://yformer.github.io/efficient-sam/

这一方法降低了 SAM 的复杂性，同时能够保持良好的性能。具体来说，SAMI 利用 SAM 编码器 ViT-H 生成特征嵌入，并用轻量级编码器训练掩码图像模型，从而从 SAM 的 ViT-H 而不是图像补丁重建特征，产生的通用 ViT 骨干可用于下游任务，如图像分类、物体检测和分割等。然后，研究者利用 SAM 解码器对预训练的轻量级编码器进行微调，以完成任何分割任务。

为了评估该方法，研究者采用了掩码图像预训练的迁移学习设置，即首先在图像分辨率为 224 × 224 的 ImageNet 上使用重构损失对模型进行预训练，然后使用监督数据在目标任务上对模型进行微调。

通过 SAMI 预训练，可以在 ImageNet-1K 上训练 ViT-Tiny/-Small/-Base 等模型，并提高泛化性能。对于 ViT-Small 模型，研究者在 ImageNet-1K 上进行 100 次微调后，其 Top-1 准确率达到 82.7%，优于其他最先进的图像预训练基线。

研究者在目标检测、实例分割和语义分割上对预训练模型进行了微调。在所有这些任务中，本文方法都取得了比其他预训练基线更好的结果，更重要的是在小模型上获得了显著收益。

论文作者 Yunyang Xiong 表示：本文提出的 EfficientSAM 参数减少了 20 倍，但运行时间快了 20 倍，只与原始 SAM 模型的差距在 2 个百分点以内，大大优于 MobileSAM/FastSAM。

在 demo 演示中，点击图片中的动物，EfficientSAM 就能快速将物体进行分割：

EfficientSAM 还能准确标定出图片中的人：

试玩地址：https://ab348ea7942fe2af48.gradio.live/

方法

EfficientSAM 包含两个阶段：1）在 ImageNet 上对 SAMI 进行预训练（上）；2）在 SA-1B 上微调 SAM（下）。

EfficientSAM 主要包含以下组件：

交叉注意力解码器：在 SAM 特征的监督下，本文观察到只有掩码 token 需要通过解码器重建，而编码器的输出可以在重建过程中充当锚点（anchors）。在交叉注意力解码器中，查询来自于掩码 token，键和值源自编码器的未掩码特征和掩码特征。本文将来自交叉注意力解码器掩码 token 的输出特征和来自编码器的未掩码 token 的输出特征进行合并，以进行 MAE 输出嵌入。然后，这些组合特征将被重新排序到最终 MAE 输出的输入图像 token 的原始位置。

线性投影头。研究者通过编码器和交叉注意力解码器获得的图像输出，接下来将这些特征输入到一个小型项目头（project head）中，以对齐 SAM 图像编码器中的特征。为简单起见，本文仅使用线性投影头来解决 SAM 图像编码器和 MAE 输出之间的特征维度不匹配问题。

重建损失。在每次训练迭代中，SAMI 包括来自 SAM 图像编码器的前向特征提取以及 MAE 的前向和反向传播过程。来自 SAM 图像编码器和 MAE 线性投影头的输出会进行比较，从而计算重建损失。

经过预训练，编码器可以对各种视觉任务的特征表示进行提取，而且解码器也会被废弃。特别是，为了构建用于分割任何任务的高效 SAM 模型，本文采用 SAMI 预训练的轻量级编码器（例如 ViT-Tiny 和 ViT-Small）作为 EfficientSAM 的图像编码器和 SAM 的默认掩码解码器，如图所示 2（底部）。本文在 SA-1B 数据集上对 EfficientSAM 模型进行微调，以实现分割任何任务。

实验

图像分类。为了评估本文方法在图像分类任务上的有效性，研究者将 SAMI 思想应用于 ViT 模型，并比较它们在 ImageNet-1K 上的性能。

如表 1 将 SAMI 与 MAE、iBOT、CAE 和 BEiT 等预训练方法以及 DeiT 和 SSTA 等蒸馏方法进行了比较。

SAMI-B 的 top1 准确率达到 84.8%，比预训练基线、MAE、DMAE、iBOT、CAE 和 BEiT 都高。与 DeiT 和 SSTA 等蒸馏方法相比，SAMI 也显示出较大的改进。对于 ViT-Tiny 和 ViT-Small 等轻量级模型，SAMI 结果与 DeiT、SSTA、DMAE 和 MAE 相比有显著的增益。

目标检测和实例分割。本文还将经过 SAMI 预训练的 ViT 主干扩展到下游目标检测和实例分割任务上，并将其与在 COCO 数据集上经过预训练的基线进行比较。如表 2 所示， SAMI 始终优于其他基线的性能。

这些实验结果表明，SAMI 在目标检测和实例分割任务中所提供的预训练检测器主干非常有效。

语义分割。本文进一步将预训练主干扩展到语义分割任务，以评估其有效性。结果如表 3 所示，使用 SAMI 预训练主干网的 Mask2former 在 ImageNet-1K 上比使用 MAE 预训练的主干网实现了更好的 mIoU。这些实验结果验证了本文提出的技术可以很好地泛化到各种下游任务。

表 4 将 EfficientSAMs 与 SAM、MobileSAM 和 SAM-MAE-Ti 进行比较。在 COCO 上，EfficientSAM-Ti 的性能优于 MobileSAM。EfficientSAM-Ti 具有 SAMI 预训练权重，也比 MAE 预训练权重表现更好。

此外， EfficientSAM-S 在 COCO box 仅比 SAM 低 1.5 mIoU，在 LVIS box 上比 SAM 低 3.5 mIoU，参数减少了 20 倍。本文还发现，与 MobileSAM 和 SAM-MAE-Ti 相比，EfficientSAM 在多次点击（multiple click）方面也表现出了良好的性能。

表 5 展示了零样本实例分割的 AP、APS、APM 和 APL。研究者将 EfficientSAM 与 MobileSAM 和 FastSAM 进行了比较，可以看到，与 FastSAM 相比，EfficientSAM-S 在 COCO 上获得了超过 6.5 个 AP，在 LVIS 上获得了 7.8 个 AP。就 EffidientSAM-Ti 而言，仍然远远优于 FastSAM，在 COCO 上为 4.1 个 AP，在 LVIS 上为 5.3 个 AP，而 MobileSAM 在 COCO 上为 3.6 个 AP，在 LVIS 上为 5.5 个 AP。

而且，EfficientSAM 比 FastSAM 轻得多，efficientSAM-Ti 的参数为 9.8M，而 FastSAM 的参数为 68M。