SAM轻量级改进EfficientSAM论文解读EfficientSAM: Leveraged Masked Image Pretraining for Efficient SegmentAnyth

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

一、总结

1. 简介

发表时间：2023年12月1日

论文：[2312.00863] EfficientSAM: Leveraged Masked Image Pretraining for Efficient Segment Anything (arxiv.org)https://arxiv.org/abs/2312.00863

代码：yformer/EfficientSAM: EfficientSAM: Leveraged Masked Image Pretraining for Efficient Segment Anything (github.com)https://github.com/yformer/EfficientSAM

 2. 摘要

         SAM巨大的计算成本限制了其在实际应用中的广泛应用，为了解决这个限制，提出了轻量级的EfficientSAM，它具有良好的性能，同时大大降低了复杂性。
        EfficientSAM的想法是基于利用掩码图像预训练SAMI，SAMI通过重建SAM图像编码器的潜在特征来改进掩膜图像的预训练，将视觉基础模型中的知识转移到ViT中，以实现有效的视觉表征学习。此外，采用SAMI预训练的轻量级图像编码器和掩码解码器来构建高效SAMs，并对SA-1B上的模型进行微调以进行分割任何任务。
        对多个视觉任务进行了评估，包括图像分类、目标检测、实例分割和语义目标检测，发现预训练方法SAMI始终优于其他掩模图像预训练方法。在分割任务上，EfficientSAM相比其他快速SAM模型表现良好，具有显着增益(例如，在COCO/LVIS上约4 AP)。

3. 前言

        SAM中的ViT-H图像编码器有632M个参数，掩码解码器只有3.87M个参数。因此，在实际使用中会导致较高的计算和内存成本。为了解决这一问题，Fast-SAM开发了一种基于实时CNN的架构YOLOv8-seg来降低计算成本；MobileSAM提出了一种解耦蒸馏方法，将知识从默认的ViT-H图像编码器提取到一个微小的ViT图像编码器中，用于获得SAM的轻量级图像编码器。EfficientSAM工作重点是处理SAM的效率问题。

4. 主要贡献

（1）提出了一个利用SAM的掩码图像预训练框架SAMI，它训练模型来重建SAM的ViT-H图像编码器的特征。研究表明，这可以大大提高图像掩码预训练方法的性能。
（2）证明了SAMI预训练的主干可以很好地推广到许多任务：图像分类、目标检测和语义分割。
（3）提出一种轻量级的SAM模型EfficientSAM，具有最先进的质量效率权衡。

二、 EfficientSAM模型结构

1. 整体结构

         如下图所示，EfficientSAM包含两个阶段：在ImageNet-1K数据集上的SAMI预训练(上)和在SA-1B数据集上的SAM调优(下)。
        SAMI预训练：使用一个用于屏蔽图像预训练的迁移学习设置，模型在图像分辨率为224 × 224的ImageNet数据集上进行重建损失预训练，掩码自编码器以SAM图像编码器的特征嵌入作为重建目标，然后使用监督数据对目标任务进行微调，通过SAMI预训练，在ImageNet-1K上训练ViT-Tiny/-Small/-Base等模型，提高泛化性能。
        SAM调优：SAMI预训练完成后，丢弃解码器，将轻量级编码器作为EfficientSAM的图像编码器，在SA-1B数据集上对对象检测、实例分割和语义分割方面对预训练模型进行了微调。

2. SAMI

        在本文中，我们建议使用预训练良好的轻量级ViT图像编码器(ViT- tiny /-Small)来降低SAM的复杂性，同时保持良好的性能。SAMI为分割任务产生所需的预训练轻量级ViT骨干。通过利用MAE预训练方法和SAM模型来获得高质量的预训练ViT编码器来实现的。具体来说，SAMI使用SAM编码器ViT-H生成特征嵌入，并使用轻量级编码器训练掩码图像模型，从SAM的ViT-H中重构特征。然后用SAM解码器对预训练的轻量级编码器进行微调，以完成分割任务。抛弃了SAMI的解码器，同时保留了编码器作为主干，以便像MAE一样提取不同任务的特征。我们将我们预训练良好的ViT主干应用于不同的任务，包括用于分类的ViT，用于检测和实例分割的ViTDet，用于语义分割任务的Mask2former，以及用于任何片段的SAM。

2.1 ViT

        与CNN相比，ViT展示了其优势和泛化能力，在视觉应用中取得了令人瞩目的成绩。现引入了ViT-Small/Deit-Small和ViT-Tiny/DeiT-Tiny等更小的ViT，以补充ViT-Huge, ViT-Large和ViT-Base。由于卷积能够以更少的参数和计算成本捕获局部信息，MobileViT探索将ViT与卷积相结合，该方法以更好的任务级泛化特性和更小的内存大小和计算成本优于轻量级CNN模型，如MobileNet-v2/v3。该技巧在后续的LeViT、EfficientFormer、Next-ViT、Tiny-ViT、Castling-ViT、EfficientViT等众多研究中都得到了应用。ViT的进展与EfficientSAM的工作是正交的。

2.2 Cross-Attention Decoder

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

2.3 Linear Projection

        通过编码器和交叉注意解码器获得图像输出。然后将这些特征输入到一个小的项目头中，用于对齐来自SAM图像编码器的特征。为了简单起见，我们只使用线性投影头来解决SAM图像编码器和MAE输出之间的特征维度不匹配。

2.4 重建的损失

        在每次训练迭代中，SAMI由SAM图像编码器的前馈特征提取和MAE的前馈和反向传播过程组成。比较了SAM图像编码器和MAE线性投影头的输出，计算了重建损失。我们的编码器、解码器和线性投影头经过优化，可以从SAM图像编码器中学习上下文建模能力。优化所有令牌上的重构损失传递嵌入在SAM中的知识。

3. MAE

        掩码自动编码器(MAE)模型有两个组成部分，一个编码器和一个解码器。编码器和解码器都建立在Transformer层上。MAE将图像标记，即来自输入图像的非重叠补丁作为输入。这些输入token按给定的屏蔽比率分组为未屏蔽token和被屏蔽token。MAE采用了较高的掩码比率，防止了预训练阶段的信息泄露。本文采用MAE框架来获得分割任意模型的高效图像编码器。由于SAM的高通用性，我们探索来自SAM图像编码器的潜在特征作为利用MAE的重建目标。编码器将未掩码标记转换为潜在特征表示，解码器通过从编码器嵌入的输出特征来重建被掩码标记的表示。重构的表示学习是由SAM的潜在特征引导的。       

4. 知识蒸馏

        知识蒸馏(Knowledge distillation, KD)是一种在不改变深度学习模型架构的情况下提高其性能的技术。某作者提出将黑暗知识从更大的教师模型提炼到更小的学生模型，学生模型的学习由教师模型的硬标签和软标签来监督，这种做法之后的多个作品，旨在更好地利用软标签，以传递更多的知识。后有人将蒸馏方法表示学习和分类解耦。解耦知识蒸馏将经典KD损失分为目标类知识蒸馏和非目标类知识蒸馏两部分，提高了知识转移的有效性和灵活性。另一项工作是从中间特征转移知识。FitNet直接从教师模型的中间特征中提取语义信息。也有学者引入了一种自监督助教(SSTA)来指导基于ViT的学生模型与有监督的老师一起学习。还有学者通过对齐较大的MAE教师模型和较小的MAE学生模型之间的中间特征，研究了预训练MAE模型知识蒸馏的潜力。    

三、实验流程

1. 实验设置

1.1 数据集

       在包含120万张图像的ImageNet-1K训练集上进行掩码图像预训练方法SAMI。下游任务考虑了三个基准数据集和几个具有代表性的视觉任务来证明所提出的SAMI的优越性，在ImageNet数据集上进行图像分类，其中包含120万张训练图像和50K验证图像；在118K训练图像和5K验证图像的COCO数据集上进行目标检测与实例分割，在ADE20K数据集上使用20K/2K/3K图像进行语义分割。在SA-1B数据集上对预训练的轻量级图像编码器进行了微调，使用来自11M张高分辨率图像的超过1B个掩码，并在COCO和LVIS上测试EfficientSAM的交互式实例分割和零镜头实例分割能力。

1.2 预训练实施细节

        ViT模型使用均方误差(MSE)损失进行预训练。使用批量大小为4096的AdamW优化器，其学习率为2.4e-3， β1 = 0.9， β2 = 0.95，权重衰减0.05，线性学习率热身超过前40 epochs，使用余弦学习率衰减更新模型。采用随机调整大小裁剪为224x224分辨率，随机水平翻转和规范化数据增强。掩码比设置为75%，解码器包含8个512维的Transformer块。在V100机器上使用PyTorch框架预训练SAMI 400次。

1.3 基线

（1）分类任务：比较不同预训练/蒸馏方法(MAE、SSTA、DMAE、BEiT、CAE、DINO、iBOT、DeiT等)的ViT主干的性能。
（2）检测和实例语义任务、语义分割任务：比较了ViTDet和Mask2former的几种预训练ViT主干。
（3）分割任务：比较了SAM、FastSAM和MobileSAM。

1.4 评价指标

        根据准确性评估我们的方法和所有基线。具体来说，准确度指标是指分类任务的 top-1 准确度；APbox, APmask用于检测和实例分割任务(AP：平均精度)；mIoU用于语义分割任务(mIoU：平均交集除以并集)；mIoU, AP, APS, APM, APL用于分割任务。对于效率度量，比较模型参数的数量或推理吞吐量。

2. 主要结果

2.1 整体性能

        如下图所示， (上)通过采用经过良好预训练的轻量级图像编码器进行实例分割，大大降低了复杂性的EfficientSAM模型概述。(下)EfficientSAM、MobileSAM、FastSAM和SAM在COCO上进行零射击实例分割的吞吐量/参数/性能比较。
        在单个NVIDIA A100上对所有型号的吞吐量(每秒图像)进行基准测试，并给出一个框提示。输入图像分辨率为1024 × 1024。EfficientSAM在相当复杂的情况下，比MobileSAM和FastSAM的性能高出约4 AP。EfficientSAM-S将SAM的推理时间减少了约20倍，参数大小减少了约20倍，性能下降很小，分别为44.4 AP和46.5 AP。

2.2 图像分类

        评估EfficientSAM在图像分类任务上的有效性，将提出的SAMI思想应用于ViT模型，并在ImageNet-1K的基线上比较它们的性能。如下表所示，将SAMI与MAE、iBOT、CAE、BEiT等预训练方法以及DeiT、SSTA等蒸馏方法进行了对比。

2.3 目标检测和实例分割 

        将SAMI预训练的ViT主干扩展到下游目标检测和实例分割任务中，并将其与COCO数据集上之前的预训练基线进行比较，以评估其有效性。具体来说，采用预训练的ViT骨干，并将其适应于Mask R-CNN框架中的简单特征金字塔，用于构建检测器ViTDet。下表显示了SAMI和其他基线之间的总体比较。

2.4 语义分割 

        进一步将预训练的主干扩展到语义分割任务中，以评估其有效性。具体来说，使用ViT模型作为Mask2former框架的主干，在ADE20K数据集上进行基准测试。如下表所示，使用Mask2former对ADE20K数据集进行语义分割的结果。输入分辨率为512 × 512。

3. EfficientSAM分割任务

        分割任务是一个过程的提示分割，以产生分段掩码基于任何形式的提示，包括点集，粗糙框或掩码，自由格式的文本。遵循SAM，重点研究了COCO/LVIS上基于点和基于框的提示分割。测试了EfficientSAM在分割任务上的泛化能力，包括零射击单点有效掩码评估和零射击实例分割。将SAMI预训练的轻量级骨干作为SAM的图像编码器，用于构建EfficientSAM。然后，在SA-1B数据集上对EfficientSAM进行了微调，并报告了零射击单点有效掩码评估和零射击实例分割的性能。

3.1 零射击单点有效掩模评估

        对click的真实掩码内的随机点进行均匀采样，并计算框的真实掩码对应的最紧的边界框，由于我们的模型能够预测多个掩码，因此我们只将最自信的掩码评估为SAM。
        如下表所示，在COCO和LVIS数据集上对单个框、单个点、多个点三种分割目标进行评估。

 3.2 零镜头实例分割

        以ViTDet生成的边界框为提示符执行实例分割任务。以bbox作为预测掩码，具有最高交点的掩码(IoU)。在表5中，我们报告了零射击实例分割的AP、APS、APM、APL。我们将我们的EfficientSAM与MobileSAM和FastSAM进行比较。如下表所示， 在COCO/LVIS数据集上的零镜头实例分割结果。

3.3 可视化结果 

如以下三图所示，补充理解EfficientSAM的实例分割功能的可视化结果。
        使用SAM、FastSAM、MobileSAM和我们的EfficientSAM模型对点提示输入的可视化结果。

        使用SAM、FastSAM、MobileSAM和我们的EfficientSAM模型对框提示输入的可视化结果。

         使用SAM, FastSAM, MobileSAM和我们的EfficientSAM模型分割所有内容的可视化结果。

 3.4 显著实例分割 

        显著目标分割旨在从图像中分割出最具视觉吸引力的对象，将交互式实例分割扩展到显著实例分割，而无需手动创建点/框。具体来说，采用最先进的显著性对象检测模型U2 -net来预测显著性图，并在显著性图内均匀采样3个随机点，使用EfficientSAM生成掩码并执行自动实例分割，而无需手动创建点或框。如下图所示，为基于显著性的自动实例分割结果。

4. 消融实验

        通过ViT主干的一系列消融研究来分析SAMI和EfficientSAM。

4.1 重建的损失

        在ImageNet-1K上研究了重建损失对SAMI性能的影响。比较了均方误差(MSE)重建损失和余弦相似度损失。我们发现MSE重构损失表现更好，如下表所示。这建议直接重建SAM特征，而不是具有高角度相似性的目标。

4.2 Cross-Attention译码器 

        为了重建SAM特征，直接使用编码器的输出token，只使用解码器对交叉注意的掩码token进行转换。研究了通过解码器作为MAE的所有token的性能变化。当查询解码器中的掩码token时，发现SAMI-Ti在ImageNet-1K数据集上的性能比将所有token作为MAE馈送到解码器中进行目标重建要好3%。

4.3 掩码比率

        在MAE中，建议采用较高的掩码比，如75%。将探讨性能如何随SAMI中不同的掩码比而变化。如下表所示，可以看到与MAE的观测结果一致，高掩码比往往产生较好的结果。

4.4 重建目标 

        采用与CLIP不同的编码器来生成特征，作为SAMI中的重建目标。对于ImageNet-1K数据集上的ViT-Tiny模型，从CLIP编码器中对齐特征也可以比MAE高出0.8%。这表明掩码图像预训练受益于强大的引导重建。

4.5 调优步骤对EfficientSAM的影响

        如下图所示，即使在0.1 epoch, EfficientSAM-Ti和EfficientSAM-S也能取得不错的性能。对于1 epoch，性能增益大于2.5 mIoU。EfficientSAM-S的最终性能达到76.9 mIoU，仅比SAM低1.5 mIoU。