前言

本系列文章是博主在工作中使用SAM模型时的学习笔记，包含三部分：

SAM初步理解，简单介绍模型框架，不涉及细节和代码SAM细节理解，对各模块结合代码进一步分析SAM微调实例，原始代码涉及隐私，此部分使用公开的VOC2007数据集，Point和Box作为提示进行mask decoder微调讲解

此篇为第二部分，如果已看过第一部分的，可以跳过下文的模型总览中介绍输入输出和流程及最后的数据引擎part，和第一篇一致。本篇很多图和部分内容参考自【大模型系列】一文看懂SAM大模型，感谢原作者。

模型总览

SAM论文: https://arxiv.org/abs/2304.02643

SAM Github：https://github.com/facebookresearch/segment-anything

SAM在线demo: https://segment-anything.com/demo

SAM的一部分灵感是来源于NLP中的基座模型(Foundation Model)，Foundation Model是OpenAI提出的一个概念，它指的是在超大量数据集上预训练过的大模型（如GPT系列、BERT），这些模型具有非常强大的 zero-shot 和 few-shot能力，结合prompt engineering和fine tuning等技术可以将基座模型应用在各种下游任务中并实现惊人的效果。

SAM就是想构建一个这样的图像分割基座模型，即使是一个未见过的数据集，模型也能自动或半自动（基于prompt）地完成下游的分割任务。为了实现这个目标，SAM定义了一种可提示化的分割任务（promptable segmentation task），这个提示可以是点、框、掩码、文本（代码中未实现）等形式，基于这个提示模型就能分割出提示处所在物体的masks。同时这种提示可以是模糊的，比如以下图剪刀握手那的黄色部分点为提示，分割掩码可以是下图最右边三种情况中任意一种，从上到下分别代表whole, part, subpart三种层级的分割，这也是SAM兼容的。要达到这种效果就需要足够的高质量分割数据，SAM团队用他们提出的Data Engine策略成功使用人工加模型自动标注的方式制作除了一个有10亿个masks的分割数据集**SA-1B**，这也是他们核心的贡献之一，本文尾部会介绍相关流程。模型架构来说相对比较常规，主要是借鉴了ViT和DETR，本身创新不大。

如上图，SAM模型架构主要包括image encoder，prompt encoder和mask decoder三部分：

image encoder，使用了ViT模型将图像编码得到image embeddingprompt encoder，将point、box、mask、txt等提示信息进行编码，后续会和image embedding一起用于生成masksmask decoder，将上述两个模块得到的embeddings整合，然后结合两个可学习的tokens生成不同层级的masks和对应的置信度值

值得一提的是，prompt encoder和mask decoder都是非常轻量的，主要的计算开销都在image encoder上，这点从模型权重上也能看出来。以ViT_B为基础的SAM权重是375M，其中prompt encoder只有32.8k，mask decoder是16.3M(4.35%)，剩余则是image encoder，可想而知图像编码这块是非常耗时的。因此在实际推理中，一般单张图的image embedding只计算一次，然后将结果缓存起来，需要的时候直接调用。在image embedding已经计算好的情况下，论文中说给定一个prompt，生成mask时prompt encoder和mask decoder在浏览器中的计算耗时也仅需50ms。下面会具体介绍下各模块的输入输出和流程，均只考虑batch size为1的情况，代码讲解在下一篇。

Image encoder

输入：

默认是1024x1024的图像，如尺寸不一致会将原图按最长边resize

输出：

单张图的1x256x64x64的image embedding，即编码后的图像特征

流程

上图是ViT论文中的结构图，image encoder整体流程和ViT是一样的，区别在于不需要[class]token做分类，只输出最终的图像编码张量

输入1024的图，拆分成64x64的768维patchs经过attention block(window和global的MSA，相对位置编码)和MLP得到同样大小64x64x768embbeding特征再经过neck得到1x256x64x64的图片embedding

这块有一篇文字介绍的更详细，如果想了解更多细节可以看这篇：Image encoder模块Vision Transformer网络解析。

Image encoder主要由attention block和neck组成，下面将根据代码简单介绍

attention block

1024x1024x3的图片经过一个patch_size=16的PatchEmbed层，将原图分为一个个小块patch，每个patch会在channel维度展开成向量（16x16x3=768），即得到1x64x64x768的patchs，然后patchs经过attention模块（window attention或者global attention），再经MLP得到与输入x一致大小的1x64x64x768

window_partition

拆分窗口，相当于把原图拆成多个小图叠在一起

<code>x = x.view(B, Hp // window_size, window_size, Wp // window_size, window_size, C) ##先按窗口大小拆分

## .permute将数据按窗口划分，再通过.view展成多个大小一直的小图

windows = x.permute(0, 1, 3, 2, 4, 5).contiguous().view(-1, window_size, window_size, C)

Attention类

window attention和global attention计算都一样，区别在于输入x的B H W大小，window的B是batch小窗口的个数，H W是窗口大小；二global的则是针对全图计算Attention

qkv是直接用一个全连接层一次性得到然后拆分成3个多个注意力的q, k, v

相对位置编码和swin transformer类似，只针对query计算，直接加到attention上

def forward(self, x: torch.Tensor) -> torch.Tensor:

## window attention的B是batchx小窗口的个数，H W是窗口大小；

B, H, W, _ = x.shape

# qkv with shape (3, B, nHead, H * W, C) ## self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias)

qkv = self.qkv(x).reshape(B, H * W, 3, self.num_heads, -1).permute(2, 0, 3, 1, 4)

# q, k, v with shape (B * nHead, H * W, C)

q, k, v = qkv.reshape(3, B * self.num_heads, H * W, -1).unbind(0)

## self.scale多头注意力的根号d，self.scale = (dim_input // num_heads )**-0.5

attn = (q * self.scale) @ k.transpose(-2, -1)

if self.use_rel_pos:

##self.rel_pos_h是网络学到的参数

attn = add_decomposed_rel_pos(attn, q, self.rel_pos_h, self.rel_pos_w, (H, W), (H, W))

attn = attn.softmax(dim=-1)

x = (attn @ v).view(B, self.num_heads, H, W, -1).permute(0, 2, 3, 1, 4).reshape(B, H, W, -1)

x = self.proj(x) ## self.proj = nn.Linear(dim, dim)

return x

相对位置编码

即add_decomposed_rel_pos函数的实现，self.rel_pos_h就是类似swin transformer的相对位置偏移表（长度为2*M-1，因为可取的值是[-M+1,M-1]共2*M-1个），是网络学习到的

# 非global的input_size[0] = input_size[1] = 14，global是64

# multi-head attention中的head维度：head_dim = block输出的维度768 / head的数量12 = 64

self.rel_pos_h = nn.Parameter(torch.zeros(2 * input_size[0] - 1, head_dim))

self.rel_pos_w = nn.Parameter(torch.zeros(2 * input_size[1] - 1, head_dim))

根据窗口大小生成对应相对位置索引index，再从self.rel_pos_h中取对应索引位置的偏移向量得到Rh(14x14x64 ),再把query展成多个小窗口形式与Rh矩阵乘法得到最终相对位置编码，attn上直接加上h和w方向的相对位置编码即可

def add_decomposed_rel_pos(

attn: torch.Tensor,

q: torch.Tensor,

rel_pos_h: torch.Tensor,

rel_pos_w: torch.Tensor,

q_size: Tuple[int, int],

k_size: Tuple[int, int], ) -> torch.Tensor:

""" Calculate decomposed Relative Positional Embeddings

Args:

attn (Tensor): attention map.

q (Tensor): query q in the attention layer with shape (B, q_h * q_w, C).

rel_pos_h (Tensor): relative position embeddings (Lh, C) for height axis.

rel_pos_w (Tensor): relative position embeddings (Lw, C) for width axis.

q_size (Tuple): spatial sequence size of query q with (q_h, q_w).

k_size (Tuple): spatial sequence size of key k with (k_h, k_w).

Returns:

attn (Tensor): attention map with added relative positional embeddings.

"""

# q： 300x196x64

# atten：300x196x196

q_h, q_w = q_size

k_h, k_w = k_size

# Rh: 14x14x64

Rh = get_rel_pos(q_h, k_h, rel_pos_h)

Rw = get_rel_pos(q_w, k_w, rel_pos_w)

B, _, dim = q.shape

# r_q: 300x14x14x64

r_q = q.reshape(B, q_h, q_w, dim)

# rel_h: 300x14x14x14

# 等价于：

# rel_h = torch.matmul(r_q, Rh.transpose(1, 2))

# rel_w = torch.matmul(r_q.transpose(1, 2), Rw.transpose(1, 2)).transpose(1, 2)

rel_h = torch.einsum("bhwc,hkc->bhwk", r_q, Rh)

rel_w = torch.einsum("bhwc,wkc->bhwk", r_q, Rw)

# 将相对位置编码加在atten里面，再resize回300x196x196

attn = (attn.view(B, q_h, q_w, k_h, k_w) + rel_h[:, :, :, :, None] + rel_w[:, :, :, None, :]).view(B, q_h * q_w, k_h * k_w)

return attn

最终attention block的forward代码

window_size为0计算全局attention，否则计算局部的window attention

def forward(self, x: torch.Tensor) -> torch.Tensor:

shortcut = x

x = self.norm1(x)

# Window partition

if self.window_size > 0:

H, W = x.shape[1], x.shape[2]

x, pad_hw = window_partition(x, self.window_size)

x = self.attn(x)

# Reverse window partition

if self.window_size > 0:

x = window_unpartition(x, self.window_size, pad_hw, (H, W))

x = shortcut + x

x = x + self.mlp(self.norm2(x))

return x

neck降低embedding维度

Sequential(

(0): Conv2d(768, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)

(1): LayerNorm2d()

(2): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)

(3): LayerNorm2d()

)

neck部分由两个卷积层组成，分别是256x768x1x1和256x256x3x3，最后输出的image imbedding的尺寸是1x256x64x64

encoder的forward

def forward(self, x: torch.Tensor) -> torch.Tensor:

x = self.patch_embed(x) ## 拆成patchs

## 网络学习到的绝对位置编码，absolute positional embedding

if self.pos_embed is not None:

x = x + self.pos_embed

### tansformer block由window attention和global attention组成

for blk in self.blocks:

x = blk(x)

## 上面得到的x是 1*64*64*768

x = self.neck(x.permute(0, 3, 1, 2))

## 最后输出的image imbedding的尺寸是1*256*64*64

return x

Prompt encoder

输入：

point、box、mask、txt（代码未实现）等prompt，格式一般如下，B为batch size

point需要包含点的x,y坐标BxNx2和label（0为前景，1位背景）BxNx1box包含框的左上和右下两个点，BxNx4，对于某个gt即单个mask，只会有1个box；如果输入的是N个box最终会生成N个masksmask一般和SAM最终输出mask的hxw(256x256)，Bx1xHxWtxt在SAM代码中未实现，这块可以参考Grounded-Segment-Anything

输出两个：

sparse_embeddings 点和框的稀疏嵌入，形状为BxNx(embed_dim)，其中N由输入点和框的数量确定，如果两者同时有则N的计算方式为（点的个数+2x框的个数）

point box 全都没有，输出大小：Bx0x256如果只有point，输出大小：Bx(N+1)x256，会补充一个[0，0]空点在最后，label为-1，表示只有点提示；如果只有box，输出大小： (B*N)x2x256piont、box都有，输出大小：BxNx256

dense_embeddings 掩码的密集嵌入，形状为Bx(embed_dim)x(embed_H)x(embed_W)，默认大小为Bx256x64x64，没有提示时会返回一个网络学习到的no mask默认嵌入

流程

网络已自动学会了针对不通过类型提示的编码信息，输入的point、box、mask等提示加上位置编码后，再加上网络学会的综合编码信息，最终对point、box这种稀疏的提示会返回sparse embedding， 对mask会返回dense embeddings（没有mask提示时是网络学习到的embeddings）。这部分就相当于把各种提示转换为decoder能理解的格式。

point embedding

输入points是个tuple，包含(point_coords, point_labels)，point_coords一般是BxNx2的一些列点的xy坐标，point_labels是BXN个点对应的label，0代表是一个背景点，1代表是前景点（需要分割出mask的部分）,以下都以1个点和1个box为提示分析

step1：首先生成一组可学习的向量point embedding，大小为：4x1x256，即前景/背景和框的左上右下两个点：

## ModuleList((0-3): 4 x Embedding(1, 256)); 4个点代表 pos/neg point + 2 box corners

self.point_embeddings = nn.ModuleList([nn.Embedding(1, 256) for i in range(4)])

step2：再生成一组可学习的向量not_a_point_embed，大小为1x256，用于表示该位置不是一个点

self.not_a_point_embed = nn.Embedding(1, embed_dim)

step3：点的padding

N为传入点的个数，点point表示为BxNx2的tensor，如果prompt里面没有bbox只有点，则补充一个[0，0]点到points后面，其对应的label为-1，此时point大小为Bx(N+1)x2，label为Bx(N+1)；

如果传入的还有bbox，此时的point大小为BxNx2，label为BxN（不加pad）

## 没有bbox，补充【0，0】点到每个point后面，其对应的label为-1

if pad:

padding_point = torch.zeros((points.shape[0], 1, 2), device=points.device, dtype=points.dtype)

padding_label = -torch.ones((labels.shape[0], 1), device=labels.device, dtype=points.dtype)

points = torch.cat([points, padding_point], dim=1)

labels = torch.cat([labels, padding_label], dim=1)

step4：计算Positionally encode：

点的横纵坐标除以输入尺寸w,h(1024,1024)到[0,1]之间，再和随机高斯矩阵(2x128, positional_encoding_gaussian_matrix)相乘后，算sin cos值后h和w拼在一起得到点的位置编码（BxNx256），N=2表示只有point，N=1表示还有box没有pad一个空点

def _pe_encoding(self, coords: torch.Tensor) -> torch.Tensor:

"""Positionally encode points that are normalized to [0,1]."""

# assuming coords are in [0, 1]^2 square and have d_1 x ... x d_n x 2 shape

coords = 2 * coords - 1

coords = coords.to(self.positional_encoding_gaussian_matrix.dtype)

coords = coords @ self.positional_encoding_gaussian_matrix ##随机高斯矩阵 2x128

coords = 2 * np.pi * coords

# outputs d_1 x ... x d_n x C shape

return torch.cat([torch.sin(coords), torch.cos(coords)], dim=-1)

step5：计算最终的point embedding

上一步得到的256维向量，叠加学习到的embedding向量（非点、背景点、前景点，各自分别学习了）最终为BxNx256的向量

point_embedding = self.pe_layer.forward_with_coords(points, self.input_image_size)

point_embedding[labels == -1] = 0.0

point_embedding[labels == -1] += self.not_a_point_embed.weight # 对应label为-1的padding点，加上not_a_point_embed

point_embedding[labels == 0] += self.point_embeddings[0].weight # neg点加上point_embeddings[0]

point_embedding[labels == 1] += self.point_embeddings[1].weight # pos点加上point_embeddings[1]

完整的embedding流程

def _embed_points(

self,

points: torch.Tensor,

labels: torch.Tensor,

pad: bool, ) -> torch.Tensor:

"""Embeds point prompts."""

points = points + 0.5 # Shift to center of pixel

# 如果没有输入的box的话，会将points的长度用0补充形成Bx(N+1)x2，label用【-1】补充成Bx(N+1)

if pad:

padding_point = torch.zeros((points.shape[0], 1, 2), device=points.device)

padding_label = -torch.ones((labels.shape[0], 1), device=labels.device)

points = torch.cat([points, padding_point], dim=1)

labels = torch.cat([labels, padding_label], dim=1)

# 将points与一个2x128的随机高斯矩阵相乘再通过进行sin、cos运算，两者的运算结果拼接得到

# point_embedding: BxNx256 或者 Bx(N+1)x256

point_embedding = self.pe_layer.forward_with_coords(points, self.input_image_size)

point_embedding[labels == -1] = 0.0

point_embedding[labels == -1] += self.not_a_point_embed.weight

point_embedding[labels == 0] += self.point_embeddings[0].weight

point_embedding[labels == 1] += self.point_embeddings[1].weight

return point_embedding

最后输出的结果是sparse_embeddings BxNx256，如果prompt只有点则sparse_embeddings就是Bx(N+1)x256的point_embedding，如果还有box，那最终结果是BxNx256的point_embedding和叠加后面Bx1x256的box embedding拼在一起当作最终的sparse_embeddings

box embedding

对于某个gt即单个mask，只会有1个box；如果输入的是N个box最终会生成N个masks。一个box有两个点，左上和右下角点，embedding步骤如下

step1：先resize为Nx2x2；N代表多个框step2：再使用point embedding一样的位置编码方式，得到corner_embedding位置编码后的Nx2x256向量step3：再加上之前网络学习到的box角点embeding向量（0 1是前景背景，2 3是box的左上右下）

def _embed_boxes(self, boxes: torch.Tensor) -> torch.Tensor:

"""Embeds box prompts."""

boxes = boxes + 0.5 # Shift to center of pixel

coords = boxes.reshape(-1, 2, 2) ## 操作与points类似，将2个点resize成Nx2x2

## 使用point embedding编码的方式，得到corner_embedding Nx2x256

corner_embedding = self.pe_layer.forward_with_coords(coords, self.input_image_size)

## 再加上学习到的点的embedding

corner_embedding[:, 0, :] += self.point_embeddings[2].weight #2 3是box的左上右下

corner_embedding[:, 1, :] += self.point_embeddings[3].weight #2 3是box的左上右下

return corner_embedding ##Nx2x256,N代表多个框

最后输出的box的embedding的尺寸是Nx2x256,N代表多个框

sparse embedding

point和box的prompt最终结果都是sparse embedding，因此代码里会将两者合并起来得到最终结果，

如果只有point，输出大小为Bx(N+1)x256，因为padding了1个空点，比如一个点时为Bx2x256如果只有box，输出大小为 (B*N)x2x256，比如一个box时为Bx2x256point、box都有，输出大小：BxNx256，N的计算方式为（点的个数+2x框的个数），比如一个point和一个box时为Bx3x256

dense_embeddings

dense_embeddings也是mask embedding，即针对全图的编码，分两种情况，输入的prompt有无mask

有mask提示，则经过self.mask_downscaling卷积网络不断下采样，简单粗暴得到Nx256x64x64维的embedding(N是输入mask的个数)

self.mask_downscaling:

Sequential(

(0): Conv2d(1, 4, kernel_size=(2, 2), stride=(2, 2))

(1): LayerNorm2d()

(2): GELU(approximate='none')code>

(3): Conv2d(4, 16, kernel_size=(2, 2), stride=(2, 2))

(4): LayerNorm2d()

(5): GELU(approximate='none')code>

(6): Conv2d(16, 256, kernel_size=(1, 1), stride=(1, 1))

无mask提示，网络有一个1x256的向量self.no_mask_embed表示无mask提示时的特征向量，直接将其复制expand成Nx256x64x64的embedding

self.no_mask_embed = nn.Embedding(1, embed_dim)

dense_embeddings = self.no_mask_embed.weight.reshape(1, -1, 1, 1).expand(

bs, -1, self.image_embedding_size[0], self.image_embedding_size[1] )

Mask decoder