【SAM轻量化】EfficientSAM: Leveraged Masked Image Pretraining for Efficient Segment Anything

EfficientSAM: Leveraged Masked Image Pretraining for Efficient Segment Anything

论文链接：http://arxiv.org/abs/2408.01835
代码链接：https://github.com/yformer/EfficientSAM

一、摘要

 关键字： 预训练方法SAMI

 Segment Anything Model (SAM) 已经作为一种强大的工具出现在多种视觉应用中。驱动其在零样本迁移和高度通用性方面表现出色的关键组件是一个在广泛的高质量SA-1B数据集上训练的超大型Transformer模型。尽管有益，但SAM模型的巨大计算成本限制了其在更广泛现实世界应用中的应用。为了解决这一限制，提出了EfficientSAMs，这是一种轻量级的SAM模型，它具有显著的性能，且复杂度大大降低。作者的想法是基于利用掩码图像预训练SAMI，它学习从SAM图像编码器重建特征，以进行有效的视觉表征学习。此外，采用SAMI预训练的轻量级图像编码器和掩码解码器来构建EfficientSAMs，并在SA-1B上微调这些模型以进行分割任何任务。

  在多个视觉任务上进行了评估，包括图像分类、对象检测、实例分割和语义对象检测，并发现我们提出的预训练方法SAMI，一致地优于其他掩码图像预训练方法。在分割任何任务，如零样本实例分割方面，EfficientSAMs与SAMI预训练的轻量级图像编码器相比，表现出色，且相对于其他快速SAM模型有显著的增益（例如，在COCO/LVIS上约4个AP）。

二、创新点

• 提出了一种利用SAM进行图像预训练的框架，称为SAMI，该框架训练模型从SAM ViT-H图像编码器重建特征。这可以显著提高图像掩码预训练方法的表现。

• 证明了SAMI预训练的骨干网络能够在许多任务上很好地泛化，包括图像分类、目标检测和语义分割。

• 提供了EfficientSAMs，这是一种轻量级的SAM模型，具有最先进的质量-效率权衡（图1），它与SAM互补，适用于实际部署。
  

三、原理

3.1. Preliminary

  Masked Autoencoders.掩码自编码器（MAE）模型包含两个组件，一个编码器和一个解码器。编码器和解码器都建立在Transformer层[55]之上。MAE将图像token（即输入图像中的非重叠块）作为输入。这些输入token被分为未掩码标记和掩码标记，并给出一个掩码比率。未掩码标记将保留给编码器提取特征，而掩码标记将被设置为MAE解码器的学习目标，在自监督学习（MIM）期间需要被重建。MAE[26]采用高掩码比率（例如，75%），这防止了在预训练阶段的信息泄露（例如，简单地根据近邻外推掩码像素）。

3.2. SAM-Leveraged Masked Image Pretraining

  将MAE框架适应于获取用于“分割任何东西”模型的有效图像编码器。受到SAM [31]高通用性的启发，探索SAM图像编码器的潜在特征作为重建目标以利用MAE。所提方法强调了转移SAM中嵌入的知识。图2顶部展示了提出的利用SAM的掩码图像预训练（SAMI）的概览。编码器将未掩码的token转换为潜在特征表征，解码器利用编码器的输出特征嵌入来重建掩码token的表征。重建的表征学习由SAM的潜在特征指导。

  交叉注意解码器。在SAM特征的监督下，只有掩码token需要通过解码器重建，而编码器的输出可以在重建期间作为锚点。在交叉注意解码器中，查询来自掩码token，而键和值来源于编码器的未掩码特征和掩码特征。将掩码token的输出特征与编码器的未掩码标记的输出特征合并，作为MAE的输出嵌入。然后，这个合并的特征将被重新排序到输入图像标记的原始位置，以进行最终的MAE输出。

  线性投影头。通过编码器和交叉注意解码器获得图像输出。然后将这样的特征输入到一个小的投影头中，以对齐SAM图像编码器的特征。为了简化，只使用一个线性投影头来解决SAM图像编码器和MAE之间的特征维度不匹配问题。

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

  SAM图像编码器为$f^{sam}$，MAE的编码器和解码器为$g^e$(权重为$W_e$)和$g^d$(权重为$W_d$)，线性投影头为$h^{\theta }$，权重为$W_{\theta }$。假设输入token表示为${x_i}_{i=1}^N$，其中$N$是token的数量。输入token被随机分组为未掩码token { x i } i ∈ U \{x_{i}\}_{i \in U} {xi​}i∈U​，掩码token { x i } i ∈ M \{x_{i}\}_{i \in M} {xi​}i∈M​，以及给定的掩码比率。令特征重新排序运算符为$\varphi$，合并运算符为$\oplus$。

 目标特征来自SAM图像编码器，可以写成 f s a m ( x ) = f s a m ( { x i } i = 1 N ) f^{sam}(x) = f^{sam}(\{x_{i}\}^{N}_{i=1}) fsam(x)=fsam({xi​}i=1N​)，MAE编码器的输出是 g e ( { x i } i ∈ U ) g^{e}(\{x_{i}\}_{i \in U}) ge({xi​}i∈U​)，解码器的输出是 g d ( { x i } i ∈ M ) g^{d}(\{x_{i}\}_{i \in M}) gd({xi​}i∈M​)。线性投影头的输出是 f h ( x ) = h θ ( ϕ ( g e { x i } i ∈ U ⊕ g d { x i } i ∈ M ) ) f^{h}(x) = h^{θ}(\phi(g^{e}\{x_{i}\}_{i \in U} ⊕g^{d}\{x_{i}\}_{i \in M})) fh(x)=hθ(ϕ(ge{xi​}i∈U​⊕gd{xi​}i∈M​))。因此，目标重建损失可以公式化为：

$$L_{W_e,W_d,W_{\theta }}=\frac{1}{N}\cdot \sum _{j=1}^N||f^{sam}(x)-f^h(x)|{|}^2,\text{\hspace{1em}(1)}$$

其中，$N$ 是输入token的数量，$||\cdot ||$表示范数。在实验中，使用L2范数作为重建损失。通过最小化重建损失$L_{W_e,W_d,W_{\theta }}$，编码器$ g^{e}$ 被优化为图像backbone，用于提取特征，作为 SAM 图像编码器的特征。编码器、解码器和线性投影头被优化为从SAM 图像编码器学习上下文建模能力。在所有token上优化重建损失，将嵌入在SAM 中的知识转移。

  SAMI 用于EfficientSAM。预训练后，编码器提取特征表示以供各种视觉任务使用，而解码器被丢弃。特别是，为了构建用于分割任何任务的EfficientSAM 模型，使用SAMI 预训练的轻量级编码器，如 ViT-Tiny 和 ViT-Small 作为图像编码器，以及SAM 的默认掩码解码器，用于EfficientSAM，如图2（底部）所示。在SA-1B 数据集上微调EfficientSAM 模型，用于分割任何任务。EfficientSAM 框架的概述如图2 所示。

四、实验

4.1. Experimental Settings

 预训练数据集。掩码图像预训练方法SAMI在ImageNet-1K训练集上进行，包含120万张图像。遵循掩码图像预训练[26]，作者未使用标签信息。使用[31]中的SAM ViT-H图像编码器来生成重建特征，预训练ViT模型包括ViT-Tiny、ViT-Small和ViT-Base。

 预训练实现细节。ViT模型使用均方误差（MSE）损失进行预训练以进行重建。使用批大小为4096，AdamW优化器[40]，学习率为$2.4e-3$，${\beta }_1=0.9$，${\beta }_2=0.95$，权重衰减$0.05$，线性学习率在前40个轮次上预热，使用余弦学习率衰减来更新所提模型。仅采用随机调整大小裁剪到$224x224$分辨率，随机水平翻转，以及归一化作为数据增强。掩码比率设置为75%，解码器包含8个Transformer块，512维，与[26]中相同。在V100机器上使用PyTorch框架对SAMI进行400个轮次的预训练。作为参考，MAE[26]需要1600个轮次的预训练。

  下游任务/数据集/模型。任务和数据集。首先考虑三个基准数据集和几个代表性的视觉任务，以展示所提出的SAMI的优越性，包括在ImageNet数据集[16]上的图像分类，其中包含120万张训练和5万张验证图像；在COCO数据集[36]上的对象检测和实例分割，包含11.8万张训练和5千张验证图像；在ADE20K数据集[72]上的语义分割，包含20000/2000/3000张图像用于训练、验证和测试。然后，考虑分割任何任务，以进一步展示提出的SAMI的优势。在[31]中的SA-1B数据集上微调预训练轻量级图像编码器，该数据集包含超过10亿个掩码，来自1100万张高分辨率图像，并在COCO和LVIS[23]上测试EfficientSAMs的交互式实例分割和零样本实例分割能力。**模型。**丢弃SAMI的解码器，保留编码器作为backbone来为不同任务提取特征，与MAE[26]中相同。为不同任务应用经过良好预训练的ViT主干，包括ViT用于分类，ViTDet[34]用于检测和实例分割，Mask2former[13]用于语义分割任务，以及SAM用于分割任何任务。

  微调设置。分类任务使用AdamW优化器，其中${\beta }_1=0.9$，${\beta }_2=0.999$，权重衰减为0.05，在32个V100 GPU上微调ViTs 100轮次，每个GPU的批量为32。初始学习率为$1e-3$，前5轮次进行线性预热，并使用余弦学习率调度器衰减至0。对于ViT-Small和ViT-Base，将层学习率衰减因子设置为0.75。

  对于ViT-Tiny，不应用层学习率衰减。对于数据增强，采用RandAugment[15]，并将标签平滑设置为0.1，mixup设置为0.8。检测和实例分割任务遵循ViTDet [34]框架，通过将ViT后端适应为一个简单的特征金字塔，用于目标检测和实例分割。在COCO上采用AdamW优化器，动量${\beta }_1=0.9$，${\beta }_2=0.999$，权重衰减0.1，在64个V100 GPU上训练模型100轮次，每个GPU的批量为1。初始学习率为2e−4，前10轮次线性预热，并使用余弦学习率计划衰减至0。模型训练100轮次。对于分割任务,预训练ViT模型作为Mask2former[13]的后端，与分割层一起在ADE20K上微调。采用AdamW优化器，其中${\beta }_1=0.9$，${\beta }_2=0.999$，小批量大小为16，权重衰减为0.05，初始学习率为2e-4。学习率通过多项式学习率计划衰减至0。学习率乘数设置为0.1用于后端。输入图像分辨率为$512\times 512$。模型在8个V100 GPU上训练160K次迭代。

  对于分割任何任务,遵循[31]，将预训练的轻量级ViT模型，ViT-Tiny和ViT-Small，作为SAM框架的图像编码器，并将EfficientSAM的编码器和解码器一起在SA-1B数据集上微调5轮次。使用AdamW优化器，动量(${\beta }_1=0.9$，${\beta }_2=0.999$)，小批量大小为128，初始lrearning率设置为$4e-4$。学习率通过线性学习率计划衰减至0。将权重衰减设置为0.1。不应用数据增强。输入图像分辨率为$1024\times 1024$。EfficientSAMs在64个A100 GPU上，每个GPU内存40GB上训练。

  基准线和评估指标。基准线。对于分类任务，比较了不同预训练/蒸馏方法的ViT主干性能，包括MAE[26], SSTA[60], DMAE[2], BEiT[3], CAE[12], DINO[6], iBOT[73], DeiT[53]等。对于检测和实例语义任务，以及语义分割任务，还与ViTDet[34]和Mask2former[13]的几个预训练ViT主干进行了比较。对于分割一切任务，与SAM[31]、FastSAM[71]和MobileSAM[68]进行了比较。评估指标。根据准确性评估所提方法和所有基准。具体来说，准确性指标指的是分类任务的top-1准确性；$A{P}^{box}$, $A{P}^{mask}$, 对于检测和实例分割任务（AP：平均精度）；mIoU, 对于语义分割任务（mIoU：均值交并比）；mIoU, $AP$, $A{P}^S$, $A{P}^M$, $A{P}^L$对于分割任何东西的任务。对于效率指标，比较模型参数的数量或推理吞吐量。

4.2. Main Results

  图像分类。为了评估提出的技巧在图像分类任务上的有效性，将提出的SAMI思想应用到ViT模型上，并在ImageNet-1K上与基线模型进行性能比较。如表1所示，SAMI与预训练方法MAE、iBOT、CAE和BEiT，以及蒸馏方法DeiT和SSTA进行了比较。SAMI-B实现了84.8%的top-1精度，分别比预训练基线MAE、DMAE、iBOT、CAE和BEiT高出1.2%、0.8%、1.1%、0.9%和0.4%。与DeiT和SSTA等蒸馏方法相比，SAMI也显示出较大的改进。对于轻量级模型，如ViT-Tiny和ViT-Small，与DeiT、SSTA、DMAE和MAE相比，SAMI报告了显著的增益。目标检测与实例分割。还扩展了在SAMI上预训练的ViT主干网络，将其应用于下游目标检测和实例分割任务，并与COCO数据集上的先前预训练基线进行比较，以评估其有效性。具体来说，采用预训练的ViT主干网络，并将它们适配到Mask RCNN框架中的简单特征金字塔[25]中，以构建检测器ViTDet[34]。

  表2显示了SAMI与其他基线的总体比较。可以看到，SAMI在所有基线上都取得了更好的性能。与MAE-B相比，SAMI-B获得了0.9 $A{P}^{bbox}$和0.6$A{P}^{mask}$的增益。对于轻量级主干网络，SAMI-S和SAMI-Ti相对于MAE-Ti和MAE-S报告了显著的增益。此外，SAMI-S显著优于DeiT-S，在$A{P}^{bbox}$和$A{P}^{mask}$上分别提高了2.6和2.3。对于其他预训练基线，SAMI仍然优于DINO和iBOT。这一系列的实验验证了所提出的SAMI在目标检测和实例分割任务中提供预训练检测器主干网络的有效性。

  语义分割。进一步将预训练的骨干网络扩展到语义分割任务，以评估其有效性。具体来说，在Mask2former [13]框架中使用ViT模型作为后端，在ADE20K数据集上进行基准测试。如表3所示，与在ImageNet-1K上使用MAE预训练的后端相比，Mask2former使用SAMI预训练的后端在mIoU上实现了更好的性能，即分别提高了2.5、4.7和3.7。这一系列的实验验证了所提方法可以很好地推广到各种下游任务。

4.3. EfficientSAMs for Segment Anything Task

  Segment Anything任务是一种可提示分割的过程，它根据任何形式的提示生成分割掩码，包括点集、粗略框或掩码、自由形式的文本。作者遵循SAM [31]，并专注于COCO/LVIS上的基于点和基于框的提示分割。现在测试所提模型在分割任何任务上的泛化能力，包括零样本单点有效掩码评估和零样本实例分割。采用SAMI预训练的轻量级骨干网作为SAM的图像编码器，以构建高效的SAMs，即EfficientSAMs。然后在SA-1B数据集上微调EfficientSAMs，并在零样本单点有效掩码评估和零样本实例分割上报告性能。

  零样本单点有效掩码评估。类似于SAM[31]，评估从单个前景点分割对象。对于一般的交互式分割，也考虑了从单个框和多个点分割对象，如[31]中介绍的那样。为了实现这一点，在真实值掩码内均匀采样随机点用于点击，并为真实值掩码计算对应的紧框。由于所提模型能够预测多个掩码，只评估置信度最高的掩码，就像SAM[31]一样。

  结果。表4将EfficientSAMs与SAM、MobileSAM和SAM-MAE-Ti进行了比较。在COCO上，EfficientSAM-Ti比MobileSAM高出1.9mIoU在1点击和1个框上，具有相当的复杂性。EfficientSAM-Ti与SAMI预训练权重在COCO/LVIS交互式分割上也表现更好。EfficientSAM-S仅在COCO框上落后SAM 1.5 mIoU，在LVIS框上落后3.5 mIoU，参数减少了20倍。还发现EfficientSAMs在多个点击上与MobileSAM和SAM-MAE-Ti相比也显示出好表现。

  零样本实例分割。遵循SAM [31]，实例分割任务通过采用ViTDet[34]生成的边界框（bbox）作为提示来
执行。与bbox具有最高交并比（IoU）的掩码作为预测掩码。

  结果。表5报告了零样本实例分割的$AP$，$A{P}^S$，$A{P}^M$，$A{P}^L$。比较了EfficientSAM与MobileSAM和FastSAM。可以看到，EfficientSAM-S在COCO上获得的AP超过了FastSAM的6.5%，在LVIS上超过了7.8%。对于EffidientSAM-Ti，它 仍 然以 很 大的优势超过了FastSAM，COCO上的4.1 AP和LVIS上的5.3 AP，以及MobileSAM的3.6 AP在COCO上和5.5 AP在LVIS上。EfficientSAMs比FastSAM轻量得多，例如，EfficientSAM-Ti的9.8M参数，而FastSAM的参数为0.6G。EfficientSAM-S还显著缩小了SAM与0.6G参数之间的差距，仅减少了约2 AP。这些结果证明了EfficientSAMs在零样本实例分割方面的非凡好处，并验证了SAMI预训练方法的优势。

 定性评估。一些例子可以在图3，图4，和图5中看到两种类型的提示的预测掩码，即点提示和框提示，如MobileSAM [68]，以及分割一切的结果。这些结果表明EfficientSAMs在与SAM相比具有竞争能力。请注意，EfficientSAMs比SAM轻量得多，可
以有效地给出不错的分割结果。这意味着模型可以作为SAM的补充版本，用于许多实际任务。

  显著实例分割。显著目标分割[4]旨在从图像中分割出视觉上最吸引人的对象。将交互式实例分割扩展到无需手动创建点/框的显著实例分割。具体来说，采用一种最先进的显著性对象检测模型U2-net[45]来预测显著性图，并在显著性图中均匀采样3个随机点（3次点击），以使用EfficientSAM进行实例分割。图6可以看到EfficientSAM能够很好地执行显著实例分割。这种初步探索打开了帮助手部功能障碍人士在图像中分割感兴趣对象的可能性。

4.4. Ablation Studies

  重建损失。研究重建损失对SAMI在ImageNet-1K上的性能影响。将均方误差（MSE）重建损失与余弦相似度损失进行了比较。发现MSE重建损失表现更好，如表6所示。这推荐直接重建SAM特征而不是具有高角度相似性的目标。

  交叉注意解码器。为了重建SAM特征，直接使用编码器的输出token，只使用解码器通过交叉注意来转换掩码token。研究了如何通过解码器随所有token变化性能，就像MAE[26]那样。当在解码器中查询掩码token时，发现SAMI-Ti在ImageNet-1K上的性能比将所有token输入解码器进行目标重建的MAE[26]高出3%。类比于AnchorDETR[58]中的锚点，编码器的输出标记已经通过直接对齐SAM特征很好地学习到了，它们可以作为锚token，通过交叉注意解码器帮助掩码标记对齐。

  掩码比率。在MAE中建议使用较高的掩码比率，即75%。在SAMI中探索了不同掩码比率下的性能变化。表7可以看到观察结果与MAE[26]一致，即较高的掩码比率往往会产生良好的结果。

  重建目标。研究了重建目标的影响。使用与CLIP [46]不同的编码器来生成特征作为SAMI中的重建目标。与ViT-Tiny模型在ImageNet-1K上的MAE相比，对齐CLIP编码器生成的特征可以提高0.8%的性能。这表明掩码图像预训练受益于强大的引导重建。

  精调轮次对EfficientSAMs的影响。探讨了精调轮次对EfficientSAMs的影响。如图7所示，即使在0.1轮次时，EfficientSAM-Ti和EfficientSAM-S也能实现不错的性能。对于1轮次，性能提升超过2.5%mIoU。EfficientSAM-S的最终性能达到76.9 mIoU，仅比SAM低1.5 mIoU。这些结果展示了SAMI预训练图像编码器和EfficientSAMs的优势。

五、总结

  文中提出了一种掩膜图像预训练方法SAMI，以探索在SAM基础模型的指导下ViTs的潜力。SAMI通过重建SAM图像编码器的潜在特征来改进掩膜图像预训练，将视觉基础模型的知识转移到ViTs中。在图像分类、目标检测和实例分割、语义分割以及分割任何任务上的广泛实验一致验证了SAMI的优势。SAMI有助于构建具有预训练轻量级编码器的有
效SAMs。