本文首发:AIWalker
本文介绍了一种名为EfficientSAM的模型,该模型通过利用遮罩图像预训练来提高图像分割的性能。作者使用了一个名为SAMI的方法,通过将SAM图像编码器的特征作为重建目标,从SAM图像编码器中重建特征,从而实现遮罩图像预训练。作者还使用SAMI预训练的轻量级图像编码器构建了EfficientSAM模型,并在SA-1B数据集上进行了验证。实验结果表明,EfficientSAM模型在图像分类、目标检测、实例分割和语义分割等任务中均取得了比其他预训练方法更好的性能。此外,作者还讨论了与遮罩图像预训练相关的方法和应用。
本文方案
- Cross Attention Decoder 只有遮罩的标记需要通过解码器进行重构,而编码器的输出可以作为重构过程中的锚点。在交叉注意力解码器中,查询来自遮罩标记,键和值则来自编码器中的未遮罩特征和遮罩特征。然后,将来自交叉注意力解码器中遮罩标记的输出特征和来自编码器中未遮罩标记的输出特征进行合并,以生成MAE输出嵌入。最后,将合并后的特征重新排序到输入图像标记的原始位置,得到最终的MAE输出。
- Linear Projection Head 通过编码器和解码器获取图像输出,然后将特征输入到一个小型项目头(project head)以对齐来自SAM图像编码器的特征。为了简化,作者使用了线性投影头(linear projection head)来解决SAM图像编码器输出和MAE之间特征维度的 mismatch 问题.
- Reconstruction Loss 在每次训练迭代中,SAMI包括从SAM图像编码器进行的一次前馈特征提取,以及MAE的一次前馈和反向传播过程。通过比较SAM图像编码器和MAE线性投影头的输出,计算重构损失。
- SAMI for EfficientSAM.在预训练之后,我们的编码器提取各种视觉任务的特征表示,而解码器被丢弃。特别地,为了构建用于分割任何任务的 efficient SAM 模型,我们采用 SAMI 预训练的轻量级编码器(如 ViT-Tiny 和 ViT-Small)作为图像编码器,并使用 SAM 的默认遮罩解码器作为我们的 EfficientSAM 的解码器,如图 2(底部)所示。我们在 SA-1B 数据集上对 EfficientSAM 模型进行微调,以用于分割任何任务。