Paper Reading：SAM & FastSAM & MobileSAM 的简单介绍

p.s. 前前后后参考了好多篇博文，（但没有记录）就不一一列出了

SAM

简介

今年4月，Meta AI公开了Segment Anything Model（SAM），分割一切模型

论文：2023年4月5日v1发布于arXiv
线上体验demo网站地址：https://segment-anything.com/
现已收录ICCV2023（影响因子IS：40.60

• 目标

建立一个分割领域的基础模型(foundation model)，一个交互式的分割框架，能通过给出的提示(prompt)返回有效分割掩码，同时具有强大的零样本泛化能力(zero-shot transfer)和可组合性(Compositionality)
，即作为一个接口来解决一系列新的数据分布上的下游分割问题

• 方法

关键在于三个组成部分:任务task、模型model和数据data。

通过引入三个相互关联的组件来构建分割的基础模型:
一个提示的分割任务，一个分割模型(SAM)，以及一个用于收集的数据引擎

1. 支持promptable的输入，根据输入提示的不同产生不同的输出
2. 模型设计。两个encoder（prompt & image）和一个decoder（lightweight mask decoder）的设计，实现根据不同prompt和image的输入，输出mask
3. 数据集的构建（data engine，数据引擎），通过标注一部分然后剩下一部分其他人标，最后加上人工校验部分达到一个区域闭环的效果
   使用了有史以来最大的分割数据集Segment Anything 1-Billion mask dataset（SA-1B），其内包含了1100万张图像，总计超过10亿张掩码图。（1M=100万，1B=10亿）

• SAM的模型概览：
  
  强大的图像编码器image encoder（使用MAE预训练的ViT）；
  提示编码器prompt encoder，分为稀疏sparse的输入：点point、框box、文本text和密集dense的输入：掩码mask
  其中，
  对于稀疏的输入，点和框通过位置编码来表示，文本使用CLIP文本编码器进行学习嵌入；
  对于密集输入，也就是掩码输入，使用卷积进行嵌入，并通过image embedding进行元素求和

作者定义了sparse和dense两种prompt，其中sparse prompt就是指demo中我们可以输入的点，目标框或者是描述目标的text，而dense prompt在目前的线上demo中体验不到，paper中也只说它对应的是mask类型的prompt，从代码里看应该是训练时候用的比较多，一般是上一次迭代预测出的一个粗分割的mask，粗略指出待分割的目标区域。

掩码解码器mask decoder（使用自注意力和交叉注意力的提示-图像双向Transformer解码器），

此外，当输入提示存在歧义时，网络将根据置信度对三个可能的掩码输出进行排序。

（训练中使用的损失函数包括焦点损失和Dice损失

task设计

设计灵感来源于NLP

输入不同类型的prompt：点、框、mask、文字，模型返回分割区域

这个任务的优势在于，带来了1、一种自然的预训练算法；2、一种通过提示将零样本泛化能力应用到下游分割任务的通用方法。

model设计

两个编码器：
prompt encoder，用来对不同的prompt进行编码，对不同类型的输入进行处理，点框统一用位置坐标表示，文本使用CLIP的文本编码器，映射到256维。
image encoder，用来获取输入图像的特征，使用MAE预训练的Vision Transformer模型，以原图作为输入，分辨率为1024x1024，输出为256x64x64。（对每张图像运行一次，在提示模型之前进行应用
一个译码器：
lightweight mask decoder：整合image encoder和prompt encoder分别输出的两个embedding，实现根据图像的特征和不同prompt，输出有效mask

image encoder

image encoder的作用是把图像映射到特征空间，如paper里所说，本质上这个encoder可以是任何网络结构，作者这里使用的是MAE预训练的ViT，当然它也可以被改成传统的卷积结构。

这里简单列出输入图像经过ViT的流程，整体有4个步骤：

• 输入图像进入网络，先经过一个卷积base的patch_embedding：取16*16为一个patch，步长也是16，这样feature map的尺寸就缩小了16倍，同时channel从3映射到768。
• patch_embed过后加positional_embedding： positional_embedding是个可学习的参数矩阵，初始化是0。
• 加了positional_embedding后的feature map过16个transformer block，其中12个transformer是基于window partition（就是把特征图分成14*14的windows做局部的attention）的attn模块，和4个全局attn，这4个全局attn是穿插在windowed attention中的。
• 最后过两层卷积（neck）把channel数降到256，这就是最终的image embedding的结果。

整体来看，这个部分的计算量是相对来说比较大的，demo体验过程中，只有这个过程的计算是在fb的服务器上做的，prompt encoder和mask decoder体积比较小，都是在浏览器内部或者说用本地的内存跑的，整体速度还比较快。

prompt encoder

prompt encoder负责prompt到prompt的特征空间，映射出的特征的channel和image embedding的channel一致，因为这两个后边要用attention进行融合。

1. prompt是point，那么它的映射由两个部分相加组成，一个是位置编码

位置编码使用的是Fourier Features Let Networks Learn High Frequency Functions in Low Dimensional Domains的编码方式

另一个部分是一个描述当前点是前景还是背景

因为demo里可以选择pos点也可以选择neg点，特征的可学习的一维向量。换句话说，如果当前选择的点是positive，那么就在位置编码的2维向量上加一个表示positive的一维向量，如果是neg，就加一个表示neg的一维向量，对于所有的positive的点，加上去的pos向量都是一样的。

2. prompt是box，box的映射也是由两个部分相加组成，第一部分是左上和右下两个点的位置编码，第二部分是一组一维向量用来描述这个点是“左上”还是“右下”。
   也就是说，对于左上的点，他的映射就是位置编码+“左上”这个特征的描述向量，右下的点，就是位置编码+“右下”这个特征的描述向量。

3. text prompt的情况，是nlp的映射处理

上面的3个都属于sparse prompt，那么对于mask这类的dense prompt，他的映射如下：
在输入prompt encoder之前，先要把mask降采样到4x，再过两个2x2，stride=2的卷积，这样尺寸又降了4x，就和降了16x的图像特征图尺寸一致了，再过一个1*1的卷积，把channel也升到256。如果没有提供mask（也就是我们实际推断时候的场景）这个结构会直接返回一个描述“没有mask”特征的特征图。

mask decoder

mask decoder是我认为整个模型中结构中比较复杂的部分

从左至右分析mask decoder的流程，

• 在prompt embedding进入decoder之前，先在它上面concat了一组可学习的output tokens，output tokens由两个部分构成：

• • 一个是iou token，它会在后面被分离出来用于预测iou的可靠性（对应结构图右侧的IoU output token），它受到模型计算出的iou与模型计算出的mask与GT实际的iou之间的MSE loss监督；

  • 另一个是mask token，它也会在后面被分离出来参与预测最终的mask（对应结构图右侧的output token per mask），mask受到focal loss和dice loss 20:1的加权组合监督。

  • 最终prompt embedding（这一步改名叫prompt token）和刚才提到这两个token concat到一起统称为tokens进入decoder。

• image embedding在进入decoder之前也要进行一步操作：dense prompt由于包含密集的空间信息，与image embedding所在的特征空间一致性更高，所以直接与image embedding相加融合。因为后面要与prompt做cross attention融合，这里还要先算一下image embedding的位置编码。

• 接下来{image embedding，image embedding的位置编码，tokens}进入一个两层transformer结构的decoder做融合。值得注意的是，在transformer结构中，为了保持位置信息始终不丢失，每做一次attention运算，不管是self-attention还是cross-attention，tokens都叠加一次初始的tokens，image embedding都叠加一次它自己的位置编码，并且每个attention后边都接一个layer_norm。

• • tokens先过一个self-attention。
  • tokens作为q，对image embedding做cross attention，更新tokens。
  • tokens再过两层的mlp做特征变换。
  • image embedding作为q，对tokens做cross attention，更新image embedding。

• 更新后的tokens作为q，再对更新后的image embedding做cross attention，产生最终的tokens。

• 更新后的image embedding过两层kernel_size=2, stride=2的转置卷积，升采样到4x大小（依然是4x降采样原图的大小），产生最终的image embedding。

• • mask token被从tokens中分离出来（因为他一开始就是concat上去的，可以直接按维度摘出来），过一个三层的mlp调整channel数与最终的image embedding一致，并且他们两个做矩阵乘法生成mask的预测。
  • iou token被从tokens中分离出来，也过一个三层的mlp生成最终的iou预测。

• 最后，如前文所述，分别对mask的预测和iou预测进行监督，反向传播，更新参数。

data engine的设计

该模型具有很好的zero-shot的迁移能力，关键在于data engine（数据引擎）的设计，通过模型不断产生数据，再通过产生的数据反过来去优化模型的性能（不断迭代提升的过程）。

已经标注好的图像分割资源少，故设计了一个特别的数据引擎来产生1B的数据大小

产生数据分为三个阶段：模型辅助手动标注阶段、半自动标注阶段和全自动阶段（Assisted-manual、 Semi-automatic、Fully automatic stage

Assisted-manual stage，辅助-手动

模型作为人工标注助手。先用一个模型预标注mask，人工根据预模型标注出的结果再进行refine。得到一个在公开数据集上训练出的初始模型，使用初始的模型对没有分割过的图片进行分割，又能得到许多新的数据，用新得到的数据返回训练初始模型，往返迭代（SAM进行了6次），该阶段从12万张图像中收集了4630万个掩码。

Semi-automatic stage，半自动

半自动化标准阶段，在第一阶段基础上增加mask的多样性。设计了一个bounding box detector去检测这些自动生成的框是否可信，保留可信的框并将其展示给人工标注，人工会在此基础上再增加额外的可信框，经过5次迭代，该阶段在18万张图像中额外收集590万个掩码（总共1020万个掩码）

Fully automatic stage，全自动

全自动标注。

这个阶段的标注是全自动的，因为模型有两个主要的增强。首先，在这一阶段的开始，收集了足够多的掩码来大大改进模型；其次，在这一阶段，已经开发了模糊感知模型，它允许在有歧义的情况下预测有效的掩码。

具体来说，用32x32的规则网络点来提示网络，并为每个点预测一组可能对应于有效对象的掩码。在模糊感知模型中，如果一个点位于某个部分或子部分上，模型将返回子部分、局部和整个对象。

该模型的IoU模块将选择高置信度且稳定掩码

具体来说，通过两个指标对生成的自动标注进行筛选，一是基于IoU prediction的module产生的confident的score，二是stable的score。

confident和stable具体的筛选标准：

confident mask：通过模型的预测IoU评分在88.0的阈值进行过滤

stable mask：（每个像素点是否属于一个物体的判断依据是通过输出分数的取logits值后判断的，大于0是，反之不是）调整阈值（+1、-1）后再取logits，若经过调整后的mask与之前的共同区域大于95%则认为这个mask是stable的

（如果阈值化概率图在0.5-σ，0.5+σ）产生相似的掩码，则认为是稳定掩码。

最后，在选择高置信度和稳定的掩码后，采用NMS对重复数据进行过滤。该阶段，在1100万张图像上全自动生成11亿个高质量掩码。

Zero-shot transfer

零样本迁移能力

在23个分割数据集上，做了5个实验：

1. Zero-Shot Single Point Valid Mask Evaluation Task (单点有效掩码评估任务)，主要与RITM进行比较
   

2. Zero-Shot Edge Detection(边缘检测)
   
   

3. Zero-Shot Object Proposals(目标建议)
   

4. Zero-Shot Instance Segmentation(实例分割)
   
   

5. Zero-Shot Text-to-Mask(文本提示生成掩码)
   

6. Ablations
   

启发

计算机视觉从对于算法、模型的具体的研究，转变为对采集数据的研究（数据引擎、链路）

FastSAM

简介

中科院自动化研究所，2023.6.21发布v1
线上demo体验：https://huggingface.co/spaces/An-619/FastSAM（也可以本地部署

本⽂将“分割⼀切”任务解耦为全实例分割和提⽰引导选择两阶段，通过引⼊⼈⼯先验结构，在提速 50 倍的情况下实现了与 SAM 相近的表现，是⾸个实时分割⼀切的基础模型。

基于此方法结构，FastSAM在SAM团队开源的SA-1B数据集上随机挑选了2%的图片进行训练就取得了和SAM可匹配的效果，并且速度相较最常用的32x32指令版本的SAM提升了50倍，实现了实时的“分割一切”

(TRT: TensorRT, 由nvidia提供的，是一种优化神经网络推理的加速器

目标

降低SAM的训练资源，可以应用于产业界。

方法

将SAM分成两个相对独立的任务：
全实例分割和提示引导选择(all-instance segmentation and prompt-guided selection）
下图是两个任务的架构，以虚线为界，

上半部分是第一阶段，基于卷积神经网络(CNN)检测器的实现，这里是直接使用了YOLOv8-seg并应用YOLACT原则进行实例分割，生成图像中所有实例的分割掩码。
下半部分是第二阶段，输出与提示符相对应的感兴趣区域

具体来说

实例分割

通过骨干网络和特征金字塔网络FPN从图像中提取特征，整合不同大小的特征。输出包括检测和分割分支。

检测分支输出类别和边界框，而分割分支输出k个原型(FastSAM中默认为32个)以及k个掩码系数。分割和检测任务并行计算。分割分支输入高分辨率特征图

这个map通过一个卷积层进行处理，放大，然后再通过两个卷积层来输出mask。

掩模系数类似于检测头的分类分支，范围在-1到1之间。将掩模系数与原型相乘并求和，得到实例分割结果

prompt指导选择

prompt选择了点，框，文本。

点和框，可以理解成点了或框了哪几个实例，就把哪几个实例的分割结果返回

文本prompt的话是将文本经过CLIP的text encoder得到text embedding，然后将实例分割的N个实例依次分别经过CLIP的image encoder，得到N个image embedding，然后看这些image embedding里哪个和text embedding相似度最高，然后就得文本对应的分割结果。

总结

从图片输入到输出的整个流程来看，FastSAM的速度提升了50倍，参数量也下降了至少2倍。

文中作者还比较了边缘检测、目标候选框生成、实例分割、文本提示目标定位，不一一展示

MobileSAM

韩国庆熙大学，2023.6.25发布v1
线上demo体验：https://huggingface.co/spaces/dhkim2810/MobileSAM（也可以本地部署

目标

本文旨在将SAM的"重量级"解码器替换为"轻量级"以使其可在移动端部署应用。

为达成该目标，本文提出一种"解耦蒸馏"方案对SAM的ViT-H解码器进行蒸馏，同时所得轻量级编码器可与SAM的解码器"无缝兼容" 。

此外，所提方案，只需一个GPU不到一天时间即可完成训练，比SAM小60倍且性能相当，所得模型称之为MobileSAM。

在推理速度方面，MobileSAM处理一张图像仅需10ms(8ms@Encoder，2ms@Decoder)，比FastSAM的处理速度快4倍，这就使得MobileSAM非常适合于移动应用

方法

原始的SAM的image encoder部分有632M，训练需要128个GPU

作者提出了decoupled distillation解耦蒸馏，本质上是：使用较小的图像编码器对原始SAM进行KD转换

难点主要在于图像编码器和组合解码器的耦合优化，作者将KD任务分为两个子任务：
图像编码器蒸馏，掩膜解码器微调

1、先蒸馏image encoder。原始的SAM里image encoder为ViT-H，将其替换为ViT-Tiny。此时不需要后面的mask decoder和prompt decoder，只需要使用ViT-H和ViT-Tiny输出的image embedding对齐即可，使用最简单MSE（均方误差）作为loss。（而不使用焦点损失focal loss和dice loss的组合来进行mask预测。）

2、蒸馏mask decoder。不过作者在论文里说其实蒸馏后的image encoder和原始的特征已经很好的对齐了，这里就没必要再做decoder的蒸馏，直接拷贝过来使用就行。

总结

参数和性能

蒸馏前和蒸馏后的比较，参数量直接下降为个位数，推理速度也降到了8ms。

训练需要的资源，蒸馏的时候只需要2个GPU就够了，而且数据量缩减到了11k，也就是说是原始数据量的千分之一

效果

两个图分别是用点和框作为prompt输入后的结果，可以看到MobileSAM和原始的SAM分割效果差不多。


MobileSAM还与FastSAM做了对比，比FastSAM分割更好且速度更快

• 最后放一下3个模型的比较表格