本文正在参加「金石计划」
论文 Segment Anything Task 章节精读
(1)带有提示的分割任务要在给定提示下返回一个有效的掩码,这意味着就算提示有歧义可能有多个输出,也要返回至少一个合理的掩码对象。他们把图像分割任务设计成了带有提示的分割任务,是因为这可以使用常见的预训练算法,并且可以通过提示,分割任务可以以零样本的方式迁移到下游的任务中去。下图中每一列都有一个歧义的“提示点”,但是可能返回的每一个结果也都是合理的,可能是局部的掩码,也可能是全部的掩码。
The promptable segmentation task, ... ,We choose this task because it leads to a natural pre-training algorithm and a general method for zero-shot transfer to downstream segmentation tasks via prompting.
(2)这个模型也借鉴了交互式分割中的方法,但交互式分割的目的在于给定足够的用户输入之后最终预测出一个有效的掩模,这与 SAM 不同,SAM 的目的是始终可以为任何提示预测出有效的掩模,即使这个提示是模糊的。
We adapt this method from interactive seg- mentation [109, 70], although unlike interactive segmenta- tion whose aim is to eventually predict a valid mask after enough user input, our aim is to always predict a valid mask for any prompt even when the prompt is ambiguous.
(3)这里给出了一个对于零样本概念的简单解释,也就是这个预训练任务赋予了模型在推断时对任何提示做出适当响应的能力,因此可以通过设计适当的提示来解决下游任务。比如我们将选中图片中的猫的框作为提示给模型,让它来帮我们检测出猫的边界,其实一般来说常见的分割任务都能够转化出对应的提示方式。
Zero-shot transfer. Intuitively, our pre-training task en- dows the model with the ability to respond appropriately to any prompt at inference time, and thus downstream tasks can be solved by engineering appropriate prompts
(4)其实图像分割是一个很广的领域,比如常见的交互式分割、边缘检测、超像素化、目标生成、前景分割、语义分割、实例分割、全景分割等,但是这些任务都是一个在完成一个单一的任务,而 SAM 的目标就是要制作一个通过提示工程解决许多现有的或者新的分割任务的泛能力模型。
Segmentation is a broad field,... ,our work is that a model trained for promptable segmentation can perform a new, different task at inference time by acting as a component in a larger system
论文 Segment Anything Model 章节精读
(1)下面是 SAM 模型的架构图,首先经过一个比较重量级的图像编码器输出图像嵌入,然后使用一个提示编码器将不同输入方式(points、box、text 等)输入的提示也转为嵌入,用来查询输入图像的嵌入,最后用一个掩码解码器以实时速度生成目标的掩码结果。对于有歧义的提示,SAM 可以输出多个有效掩码及其关联的置信度。如图中输出了剪子的部分手柄、剪子的全部手柄、剪子的全貌。
Segment Anything Model (SAM) overview. A heavyweight image encoder outputs an image embedding that can then be efficiently queried by a variety of input prompts to produce object masks at amortized real-time speed. For ambiguous prompts corresponding to more than one object, SAM can output multiple valid masks and associated confidence scores.
(2)其实图像编码器理论上可以使用输出 C×H×W 图像嵌入的任意网络,但是出于可扩展性和能预训练的考虑,图像编码器使用了何恺明大佬的 MAE 预训练出来的 ViT 模型,恺明出品,必属精品!文末有链接需要的同学自取,然后进行了最小化适配去处理高分辨率的输入图像。具体来说就是将图像处理成 1024×1024 的输入分辨率大小,使用编码器图像嵌入为 64×64 。为了减少通道维度,又使用了一个 1x1 的卷积核和一个 3x3 的卷积核,将通道数变为 256 ,也就是输出是一个 256x64×64 的大小,每层卷积后都跟着有一层层归一化。
Image encoder. Motivated by scalability and powerful pretraining methods, we use an MAE [47] pre-trained Vision Transformer (ViT) [33] minimally adapted to process high resolution inputs [62]. The image encoder runs once per image and can be applied prior to prompting the model.
(3)提示编码器会将提示都编码器为 256 大小的嵌入,主要考虑了对两类提示进行处理:稀疏提示(也就是点、框、文本等)和密集提示(也就是掩码标签)。对点提示进行编码时,将位置编码与指示前景或者背景的特征嵌入进行相加即可。对于框提示进行编码时,其实是一个编码对,也就是两个点提示的编码,一个是框的左上角点,另一个是框的右下角点,左上角点是将其位置编码和一个学习到的嵌入进行相加即可,右下角点的嵌入类似。对于文本提示,这里使用了已经很成熟 CLIP 文本编码器来编码。对于密集提示(即掩码)编码时,密集提示(即掩码)与图像具有空间对应关系。将掩码以比输入图像低 4 倍的分辨率输入,然后使用两个输出通道分别为 4 和 16 的 2×2 卷积进一步缩小。最后使用 1×1 卷积将通道维映射到 256 。每层卷积之间都使用 GELU 激活和层标准化。最后将掩码和图像嵌入然后逐元素相加。如果没有掩码提示,那么将添加一个表示“没有掩码”的嵌入到每个图像嵌入位置。
Prompt encoder. We consider two sets of prompts: sparse (points, boxes, text) and dense (masks). We represent points and boxes by positional encodings [95] summed with learned embeddings for each prompt type and free-form text with an off-the-shelf text encoder from CLIP [82]. Dense prompts (i.e., masks) are embedded using convolutions and summed element-wise with the image embedding.
(4)轻量级掩码解码器可以讲图像嵌入和提示嵌入映射输出预测的掩码,这里用的是 Transformer 架构,并且做了一些简单的结构修改,是一个两层的结构,但是在进行解码之前,需要将一个 output tokens 嵌入加入到 prompt tokens 嵌入中,output tokens 嵌入是用在最后解码的时候,类似于分类中的 【class】标识,这里的每个 token 在解码的时候可以看作表示图像中某个局部部位是什么东西。为了方便这里将 output tokens 嵌入和 prompt tokens 嵌入都叫 tokens ,然后进行解码的时候需要进行四个步骤(如下图所示):
- 先在 tokens 中进行 self-attention ,其实就是为了让这些 tokens 互相分配一下任务并互相通告一下,让每个 token 一会都能在图像上都有自己关注的某个部分。假如图片中有两个部分一个猫和一个狗,我们有两个 token ,那么第一个 token 就会告诉其他 token 它的任务是关注图片中的猫,第二个 token 也会告诉其他 token 它关注的是狗,这样两个 token 就不会打架也不会摸鱼。
- 从 tokens 到图像嵌入进行 cross-attention ,也就是让每个 token 在图像嵌入各个部分进行一轮查找,确定下来每个 token 自己要关注的那部分在哪。还是上面的例子,第一个 token 的任务是找图片中的猫,所以它遍历了一遍两个图片部分找到了猫是在图片的左边,第二个 token 也遍历了一遍两个图片部分发现狗在右边,在上面一步他们只是知道任务,现在他们就都找到了自己明确的关注位置了。
- 接下来使用 MLP 来更新每个 token 的嵌入,其实就是为了整合之前的特征。
- 从图像嵌入到 token 进行 cross-attention ,这里其实就是第二步的一个逆向过程,前面做了从 token 到图像嵌入的 cross-attention ,这里再反向做一次,按我的理解应该是能加强 token 和图像嵌入之间的更多互相关注的信息。也就是每个 token 都知道自己关注的部分在图像中的哪里,同时每个图像的部分也知道关注自己的 token 是哪个,是个双向奔赴的过程,这样应该在解码分割的时候会更加精准。这一步也要更新图片嵌入和所有的 token 嵌入。下一层解码器层获取这些更新的嵌入继续重复上面的步骤。
为确保解码器可以获取关键的几何信息,要将位置编码添加到图像嵌入中,这样每当图像嵌入参与注意力机制计算时候,位置编码也都进行了计算,其实就是为了加强之后对不同位置的精细化解码能力。另外,整个原始的提示 token 嵌入(包括了位置编码)在每次参与注意力计算时,都会重新添加到更新后的 token 嵌入中。这其实还是要加强提示信息在我们模型中的注意力。
运行解码器后,我们使用两个反向卷积将更新后的图像嵌入上采样放大 4 倍。然后我们将上面第三、四步的结果进行又一次注意力计算生成 output tokens 嵌入,并传入了一个 3 层 的 MLP,最后我们使用上采样的图像嵌入和生成的 output tokens 嵌入进行点积运算,以预测掩码。
Mask decoder. The mask decoder efficiently maps the im- age embedding, prompt embeddings, and an output token to a mask. This design, inspired by [14, 20], employs a modification of a Transformer decoder block [103] followed by a dynamic mask prediction head. Our modified decoder block uses prompt self-attention and cross-attention in two directions (prompt-to-image embedding and vice-versa) to update all embeddings. After running two blocks, we up- sample the image embedding and an MLP maps the output token to a dynamic linear classifier, which then computes the mask foreground probability at each image location.
参考
- MAE 模型:zhuanlan.zhihu.com/p/432614068
- CLIP 论文:arxiv.org/pdf/2103.00…
文中不妥之处欢迎讨论指正,另外更文不易,禁止转载和抄袭。