一點(diǎn)就分享系列(理解篇5）Meta 出品 Segment Anything 通俗解讀——主打一個(gè)”Zero shot“是貢獻(xiàn)，CV依然在！
前言
Meta開源了Segment Anything項(xiàng)目和官網(wǎng)展現(xiàn)了非常令人印象深刻的分割Demo,其SAM模型給CV研究領(lǐng)域，個(gè)人簡要理解有以下貢獻(xiàn)：
1.Openset形式的分割任務(wù)的模型
2.超大的分割數(shù)據(jù)集
3.更好的應(yīng)用思路，提示和組合可以派生很多有趣的應(yīng)用
4.促進(jìn)CV及多模態(tài)以及多任務(wù)一統(tǒng)趨勢(shì)
META最近很活躍。先提出了LLAMA去對(duì)標(biāo)GPT3，這幾天又來了CV的大模型SAM給我們驚喜，今天來整理分析一波。 另外最重要的一定要致敬谷歌，沒有transformer就沒有現(xiàn)在的大模型，多模態(tài)AI領(lǐng)域的這么多研究成果。
源碼地址：https://github.com/facebookresearch/segment-anything
一、Segment Anything
SAM這波問世，很多人對(duì)CV統(tǒng)一的趨勢(shì)無限看好，我們先不論這個(gè)，因?yàn)槲矣X得CV方向的工程應(yīng)用點(diǎn)就決定了，畢竟現(xiàn)實(shí)問題都是具體情況具體分析 ，不單單只考慮研究層面，我覺得更多的意義是帶來促進(jìn)基礎(chǔ)研究，就像你用GPT3更多只是輔助你編程，而不是取代！CV依然在，該干嘛干嘛。
大模型正在沖擊AI領(lǐng)域，而其重要的特點(diǎn)是 Zero-shot(零樣本遷移性), 從OPENADI的CLIP的那一年開始我們所見的多模態(tài)研究早已層出不窮，（還是那句話我一直都看好多模態(tài)，從幾年前我就在念了），SAM加劇了這個(gè)趨勢(shì)，但還遠(yuǎn)遠(yuǎn)不夠，對(duì)從業(yè)者來說大模型的出現(xiàn)其實(shí)是利大于弊，這使得我們?cè)僖膊恍枰ㄖ苹墒斓囊曈X模型任務(wù)和解放了微調(diào)訓(xùn)練等過程成為了可能，那么下面來看一下meta是怎么設(shè)計(jì)出的。經(jīng)過代碼初步閱讀，目前text的輸入clip還不支持,但是其實(shí)也不難實(shí)現(xiàn)。
1. 大模型的前置需求——寶貴的大規(guī)模數(shù)據(jù)集
SAM是在1100萬的圖像數(shù)據(jù)集上訓(xùn)練出來的：一個(gè)強(qiáng)大的大模型是需要多樣化的海量數(shù)據(jù)去喂養(yǎng)出來的，但是目前沒有用于分割的超大規(guī)模數(shù)據(jù)集素材，因此Meta首先解決的數(shù)據(jù)集構(gòu)建問題，他們提供了一個(gè)名為SA-1B的數(shù)據(jù)集(1億掩碼、1100萬張圖像)，思路就是在線循環(huán)獲取思路，為了獲取更多的掩碼，建立“數(shù)據(jù)引擎”,就是我們常用的訓(xùn)練模型標(biāo)注數(shù)據(jù)再訓(xùn)練的“循環(huán)邏輯”,具體分為三個(gè)階段：
1.SAM模型去輔助人員標(biāo)記MASK，類似于交互分割；
2.SAM模型經(jīng)過prompt的提示可能性的去標(biāo)注部分對(duì)象自動(dòng)標(biāo)注，其余標(biāo)注對(duì)象由人工完成；
3.SAN模型最后使用前景信息去使得SAM自動(dòng)標(biāo)注。
如下圖展示，高效的數(shù)據(jù)量制作結(jié)果比對(duì)，比 COCO 完全手動(dòng)基于多邊形的掩碼注釋快 6.5 倍，比之前最大的數(shù)據(jù)注釋工作快 2 倍，后者也是模型輔助的。
2.基礎(chǔ)任務(wù)的泛化方式
SAM的簡單描述，是交互式分割和自動(dòng)分割的統(tǒng)一，如下圖展示,提供點(diǎn)坐標(biāo)（前后景）、BOX、文本使用prompt的形式去實(shí)時(shí)輸出分割掩碼，該掩碼應(yīng)該是有效高質(zhì)量的，因?yàn)檎鎸?shí)場(chǎng)景中即使通過promopt也往往存在歧義性，輸出應(yīng)該有至少一個(gè)對(duì)象的合理結(jié)果，比如衣服上的一塊像素可能描述衣服也可以是穿衣服的人，那么具備 prompt的分割是的特點(diǎn)：
1. 可以當(dāng)作預(yù)訓(xùn)練模型的，即使這個(gè)是歧義的，但是它必須是有效正確的；
2. 通過提示去改善并完成通用的下游分割任務(wù)，其實(shí)最大的特點(diǎn)那就是：Zero-shot以及其prompt的靈活性，交互性只是一個(gè)偏向應(yīng)用上的需求。如何展示其特性？具體例子如，你實(shí)際需要解決A的分割，但是你需要自動(dòng)化的流程，如果當(dāng)你有一個(gè)A的訓(xùn)練好的檢測(cè)器，你不再需要從新定制訓(xùn)練，只需要將問題轉(zhuǎn)化為制作prompt上。
由此可見，整體的模型并不那么復(fù)雜，且會(huì)更利于工程上的部署集成。
3.模型結(jié)構(gòu)設(shè)計(jì)
SAM如上圖所示，主要由一個(gè)圖像的encoder和提示的encoder通過一個(gè)輕量的掩碼deocoder去完成最終的預(yù)測(cè)，并且考慮到上述的目標(biāo)語義的歧義性問題，使得模型具備歧義性，研究人員對(duì)一個(gè)提示會(huì)預(yù)測(cè)多種掩碼，如上圖中“剪刀”的三個(gè)結(jié)果。
3.0 語義多樣性問題
如上述，作者使用了對(duì)一個(gè)PROMPT盡可能的來表示多個(gè)語義，最終總結(jié)出對(duì)于掩碼表示三個(gè)程度：整體、局部、子部分，依舊如官方圖例。
3.1 圖像編碼器
順便提一下目前使用比較多的幾個(gè)Base on transformer的CV主干，SWIN系列/VIT系列/DETR系列，這里主要以體系backbone構(gòu)建為基準(zhǔn)，當(dāng)然還有很多改進(jìn)的算子，這里不提了。
輸入圖像進(jìn)行embedding，這步很簡單，從代碼中看出，這個(gè)編碼器是源于meta的vitdet，VITDET是何凱明的工作，其實(shí)也是MAE 和ViT-based Mask R-CNN[的衍生，這里簡單普及一下，可以理解為是對(duì)標(biāo)swintransformer的另一設(shè)計(jì)思路，解決的問題:
在這項(xiàng)工作中，何愷明等研究者追求的是一個(gè)不同的方向：探索僅使用普通、非分層主干的目標(biāo)檢測(cè)器。如果這一方向取得成功，僅使用原始 ViT
主干進(jìn)行目標(biāo)檢測(cè)將成為可能。在這一方向上，預(yù)訓(xùn)練設(shè)計(jì)將與微調(diào)需求解耦，上游與下游任務(wù)的獨(dú)立性將保持，就像基于 ConvNet
的研究一樣。這一方向也在一定程度上遵循了 ViT
的理念，即在追求通用特征的過程中減少歸納偏置。由于非局部自注意力計(jì)算可以學(xué)習(xí)平移不變特征，它們也可以從某種形式的監(jiān)督或自我監(jiān)督預(yù)訓(xùn)練中學(xué)習(xí)尺度等變特征。
下面簡單概述下VITDET：
研究者表示，在這項(xiàng)研究中，他們的目標(biāo)不是開發(fā)新的組件，而是通過最小的調(diào)整克服上述挑戰(zhàn)。具體來說，他們的檢測(cè)器僅從一個(gè)普通 ViT 主干的最后一個(gè)特征圖構(gòu)建一個(gè)簡單的特征金字塔。這一方案放棄了 FPN 設(shè)計(jì)和分層主干的要求。為了有效地從高分辨率圖像中提取特征，他們的檢測(cè)器使用簡單的非重疊窗口注意力，少量的跨窗口塊來傳播信息，這些塊可以是全局注意力或卷積。這些調(diào)整只在微調(diào)過程中進(jìn)行，不會(huì)改變預(yù)訓(xùn)練，總結(jié)要點(diǎn)，基于VIT系列的設(shè)計(jì)，
- MAE特點(diǎn)：CV中transformer的應(yīng)用瓶頸，CV應(yīng)用Transformer是有監(jiān)督的方式，而NLP大部分用的是自監(jiān)督的方式，如何能在Transformer的方法中使用自監(jiān)督的方法，就是MAE提出的背景，MAE相當(dāng)于CV的“BERT”（非對(duì)稱的編解碼器結(jié)構(gòu)，非對(duì)稱的意思是編碼器只能看到可見的塊（這樣也可以極大的減少計(jì)算量），解碼器可以看到全部，比如圖像在分塊后，隨機(jī)蓋住，encoder收到的是其余未遮蓋的圖片，然后我們補(bǔ)上缺失的序列，再送到decoder中去，得到就是”補(bǔ)全“，這樣減少了圖像的冗余，因?yàn)檎Z言是單一的token序列，而圖像分塊的PATCH信息并不唯一，因此要遮蓋75%+的像素，通過插值回復(fù)）。
- 通過viT構(gòu)建，不使用窗口化的和FPN的結(jié)構(gòu)設(shè)計(jì)，它的好處可以通過由大步幅、單一尺度圖構(gòu)建的簡單金字塔來有效地獲得，只要信息能在少量的層中很好地跨窗口傳播，僅僅使用主干的最后一的特征層輸出，因?yàn)樗鼞?yīng)該具有最強(qiáng)大的特征，這樣窗口注意力就夠用了。
代碼比較冗長，就是常規(guī)的模型Block定義，比較常見，源碼在這，有興趣的自行閱讀和論文，SAM的VITDET圖像編碼器代碼定義
3.2 提示編碼器
接受sparse的輸入：point pair /boxes/text ，dense 輸入：mask掩碼。該decoder通過位置編碼和每個(gè)提示輸入的embedding相加來表示輸入的特征，同時(shí)使用CLIP的文本編碼器來表示特征文本。MASK則通過卷積才采樣后和圖像編碼器的embedding相加。其輸出作為maskdeocer的輸入。
這里對(duì)代碼做了簡化，簡單分析下代碼，
class PromptEncoder(nn.Module):
def __init__(
self,
embed_dim: int,
image_embedding_size: Tuple[int, int],
input_image_size: Tuple[int, int],
mask_in_chans: int,
activation: Type[nn.Module] = nn.GELU,
) -> None:
"""
Encodes prompts for input to SAM's mask decoder.
Arguments:
embed_dim (int): The prompts' embedding dimension
image_embedding_size (tuple(int, int)): The spatial size of the
image embedding, as (H, W).
input_image_size (int): The padded size of the image as input
to the image encoder, as (H, W).
mask_in_chans (int): The number of hidden channels used for
encoding input masks.
activation (nn.Module): The activation to use when encoding
input masks.
"""
super().__init__()
self.embed_dim = embed_dim
self.input_image_size = input_image_size
self.image_embedding_size = image_embedding_size
self.pe_layer = PositionEmbeddingRandom(embed_dim // 2) #隨機(jī)位置編碼初始化
self.num_point_embeddings: int = 4 # pos/neg point + 2 box corners 對(duì)點(diǎn)的emdding數(shù)量
point_embeddings = [nn.Embedding(1, embed_dim) for i in range(self.num_point_embeddings)]
self.point_embeddings = nn.ModuleList(point_embeddings)
self.not_a_point_embed = nn.Embedding(1, embed_dim)
self.mask_input_size = (4 * image_embedding_size[0], 4 * image_embedding_size[1])
## 對(duì)掩碼輸入設(shè)計(jì)一個(gè)CNN下采樣模塊
self.mask_downscaling = nn.Sequential(
nn.Conv2d(1, mask_in_chans // 4, kernel_size=2, stride=2),
LayerNorm2d(mask_in_chans // 4),
activation(),
nn.Conv2d(mask_in_chans // 4, mask_in_chans, kernel_size=2, stride=2),
LayerNorm2d(mask_in_chans),
activation(),
nn.Conv2d(mask_in_chans, embed_dim, kernel_size=1),
)
self.no_mask_embed = nn.Embedding(1, embed_dim) #初始化一個(gè)emdding，在沒有mask輸入到prompt的情況下
def forward(
self,
points: Optional[Tuple[torch.Tensor, torch.Tensor]],
boxes: Optional[torch.Tensor],
masks: Optional[torch.Tensor],
) -> Tuple[torch.Tensor, torch.Tensor]:
"""
Embeds different types of prompts, returning both sparse and dense
embeddings.
Arguments:
points (tuple(torch.Tensor, torch.Tensor) or none): point coordinates
and labels to embed.
boxes (torch.Tensor or none): boxes to embed
masks (torch.Tensor or none): masks to embed
Returns:
torch.Tensor: sparse embeddings for the points and boxes, with shape
BxNx(embed_dim), where N is determined by the number of input points
and boxes.
torch.Tensor: dense embeddings for the masks, in the shape
Bx(embed_dim)x(embed_H)x(embed_W)
"""
bs = self._get_batch_size(points, boxes, masks)
sparse_embeddings = torch.empty((bs, 0, self.embed_dim), device=self._get_device()) #點(diǎn)、文本的稀疏embeddings
if points is not None:
coords, labels = points
point_embeddings = self._embed_points(coords, labels, pad=(boxes is None))
sparse_embeddings = torch.cat([sparse_embeddings, point_embeddings], dim=1)
if boxes is not None: #boxes embedding
box_embeddings = self._embed_boxes(boxes)
sparse_embeddings = torch.cat([sparse_embeddings, box_embeddings], dim=1)
if masks is not None:
dense_embeddings = self._embed_masks(masks) #掩碼embedding
else:
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]
)
return sparse_embeddings, dense_embedding
最終會(huì)輸出一個(gè)密集embeds一個(gè)稀疏embeds（由存在的點(diǎn)、框拼接），然后你會(huì)發(fā)現(xiàn)并沒有官方給出的clip的文本輸入處理功能，應(yīng)該是沒釋放出來，其實(shí)大體邏輯：無非就是將CLIP作為特征embeddings，適配到掩碼解碼器中。
3.3 掩碼編碼器
仍然基于transfromer設(shè)計(jì)，只不過這個(gè)maskdecoder，作為最終的解碼器，使用了self-attention和cross-attention的組合來處理以上編碼器輸入的embedding，通過上采樣和MLP線性分類，最終計(jì)算每個(gè)像素的前景概率。這里如果對(duì)交叉注意力機(jī)制不理解的，簡單概述下其設(shè)計(jì)原理：
1.Cross attention概念 可以輸入兩個(gè)序列，必須具有相同的維度，兩個(gè)序列可以是不同的模式形態(tài)（如：文本、聲音、圖像） 一個(gè)序列作為輸入的Q，定義了輸出的序列長度，另一個(gè)序列提供輸入的K&V
2.Cross-attention將兩個(gè)相同維度的嵌入序列不對(duì)稱地組合在一起，而其中一個(gè)序列用作查詢Q輸入，而另一個(gè)序列用作鍵K和值V輸入。當(dāng)然也存在個(gè)別情況，他們的QKV是來自兩種輸入，而SELF-attention來自同一個(gè)輸入，這種交叉的注意力更適合多模態(tài)輸入。
具體而言，我們簡要看看代碼，代碼的定義很規(guī)范，首先我們看一下上述我們替的注意力組合模塊代碼，這里我還是挑最核心的說，概述下這個(gè)模塊的流程：
“TwoWayAttention”就是cross-att和self-att的組合，一個(gè)Block由四層構(gòu)成，如下注釋，
- 第一層的自注意力收到的embeds的sparse的輸入：輸入一個(gè)Query，和其sparse類位置編碼query_pe相加后得到新的Query，經(jīng)過self-attention
會(huì)判斷是否為第一層，不是的話會(huì)進(jìn)行殘差操作，再norm輸出 ； - 第二層的交叉注意力層，spare to dense—text2image.將token到image特征，上一層的輸出加上spare類的位置編碼query_pe和該層的位置編碼向量key_pe,其實(shí)該層是處理圖像的位置特征編碼，相加后經(jīng)過cross-attention輸出后還是殘差+norm.
- 第三層是MLP+殘差后Norm，沒什么好說的
- 最后一層依然是交叉注意力，計(jì)算流程是一樣的，image2text,dense圖像emd 嵌入到token中，和第二層反向，且都是2倍采樣。
class TwoWayAttentionBlock(nn.Module):
def __init__(
self,
embedding_dim: int,
num_heads: int,
mlp_dim: int = 2048,
activation: Type[nn.Module] = nn.ReLU,
attention_downsample_rate: int = 2,
skip_first_layer_pe: bool = False,
) -> None:
"""
A transformer block with four layers: (1) self-attention of sparse
inputs, (2) cross attention of sparse inputs to dense inputs, (3) mlp
block on sparse inputs, and (4) cross attention of dense inputs to sparse
inputs.
Arguments:
embedding_dim (int): the channel dimension of the embeddings
num_heads (int): the number of heads in the attention layers
mlp_dim (int): the hidden dimension of the mlp block
activation (nn.Module): the activation of the mlp block
skip_first_layer_pe (bool): skip the PE on the first layer
"""
super().__init__()
self.self_attn = Attention(embedding_dim, num_heads)
self.norm1 = nn.LayerNorm(embedding_dim)
self.cross_attn_token_to_image = Attention(
embedding_dim, num_heads, downsample_rate=attention_downsample_rate
)
self.norm2 = nn.LayerNorm(embedding_dim)
self.mlp = MLPBlock(embedding_dim, mlp_dim, activation)
self.norm3 = nn.LayerNorm(embedding_dim)
self.norm4 = nn.LayerNorm(embedding_dim)
self.cross_attn_image_to_token = Attention(
embedding_dim, num_heads, downsample_rate=attention_downsample_rate
)
self.skip_first_layer_pe = skip_first_layer_pe
def forward(
self, queries: Tensor, keys: Tensor, query_pe: Tensor, key_pe: Tensor
) -> Tuple[Tensor, Tensor]:
# Self attention block
if self.skip_first_layer_pe:
queries = self.self_attn(q=queries, k=queries, v=queries)
else:
q = queries + query_pe
attn_out = self.self_attn(q=q, k=q, v=queries)
queries = queries + attn_out
queries = self.norm1(queries)
# Cross attention block, tokens attending to image embedding
q = queries + query_pe
k = keys + key_pe
attn_out = self.cross_attn_token_to_image(q=q, k=k, v=keys)
queries = queries + attn_out
queries = self.norm2(queries)
# MLP block
mlp_out = self.mlp(queries)
queries = queries + mlp_out
queries = self.norm3(queries)
# Cross attention block, image embedding attending to tokens
q = queries + query_pe
k = keys + key_pe
attn_out = self.cross_attn_image_to_token(q=k, k=q, v=queries)
keys = keys + attn_out
keys = self.norm4(keys)
return queries, keys
下面還是文字?jǐn)⑹霭?，通過多個(gè)TwoWayAttentionBlock和一個(gè)LN層，我們可以得到輸出的Tensor：queries, keys，然后經(jīng)過一個(gè)Cross注意力層和LN層
# Apply the final attenion layer from the points to the image
q = queries + point_embedding
k = keys + image_pe
attn_out = self.final_attn_token_to_image(q=q, k=k, v=keys)
queries = queries + attn_out
queries = self.norm_final_attn(queries)
最后，我們?cè)倏碝askdecoder ，簡要看下代碼：
接受各個(gè)編碼器的輸入embeddings,輸出low_res_masks, iou_prediction,后處理取到最終masks
low_res_masks, iou_predictions = self.mask_decoder(
image_embeddings=curr_embedding.unsqueeze(0),
image_pe=self.prompt_encoder.get_dense_pe(),
sparse_prompt_embeddings=sparse_embeddings,
dense_prompt_embeddings=dense_embeddings,
multimask_output=multimask_output,
)
masks = self.postprocess_masks(
low_res_masks,
input_size=image_record["image"].shape[-2:],
original_size=image_record["original_size"],
)
masks = masks > self.mask_threshold
outputs.append(
{
"masks": masks,
"iou_predictions": iou_predictions,
"low_res_logits": low_res_masks,
}
)
4.訓(xùn)練部分和損失函數(shù)
由于官方?jīng)]有公開訓(xùn)練代碼，損失函數(shù)只是論文中提到了使用了DETR的DICE作分割LOSS和Focal LOSS。
5. 零樣本的遷移驗(yàn)證
從前景點(diǎn)、邊緣信息、BOX、文本CLIP等方式進(jìn)行實(shí)驗(yàn)驗(yàn)證，這里就多過多介紹了。
總的來說，其實(shí)全篇核心就是展示：ZERO-SHOT文章來源：http://www.zghlxwxcb.cn/news/detail-416531.html
6.應(yīng)用拓展
6.1 自動(dòng)生成prompot的方式
目前看來自動(dòng)實(shí)現(xiàn)的PROMPT主要依賴于視覺任務(wù)模型和語言模型，其CLIP部分代碼還沒release，故這里其實(shí)目前有兩種最直接的拓展：文章來源地址http://www.zghlxwxcb.cn/news/detail-416531.html
- 自己調(diào)用CLIP集成到該項(xiàng)目中
- 在開源項(xiàng)目中（https://github.com/IDEA-Research/Grounded-Segment-Anything），在CVPR2023的開源項(xiàng)目和論文Grounded-DINO中介紹了同樣作為zero-shoT泛化模型是一個(gè)多模態(tài)的融合模型（這個(gè)后續(xù)會(huì)另起文章，下篇預(yù)告），論文中模型設(shè)計(jì)如下：
可以看到底層的算子出現(xiàn)了我們之前說的cross-att, 看來基本多模態(tài)的注意力這算是一個(gè)“標(biāo)配”了，通過text提供圖像的檢測(cè)，并且作為Prompt來提供SAM作為下游任務(wù)去使用，并且也集合Stabel diffusion去作prompt類似于controlnet去控制的機(jī)制一樣。4. 以上結(jié)合GPT、SD等可玩性十足，總之只要是類似PROMPT機(jī)制的都可以進(jìn)行結(jié)合，下篇ZERO-SHOT模型預(yù)告準(zhǔn)備中！該篇如果開源項(xiàng)目中有補(bǔ)充更新同步更新！搬磚不易，希望我3個(gè)小時(shí)的奮筆疾書能換你的點(diǎn)贊關(guān)注~
到了這里，關(guān)于一點(diǎn)就分享系列(理解篇5）Meta 出品 Segment Anything 4月6號(hào)版核心極速解讀——主打一個(gè)”Zero shot“是貢獻(xiàn)和輔助，CV依然在！的文章就介紹完了。如果您還想了解更多內(nèi)容，請(qǐng)?jiān)谟疑辖撬阉鱐OY模板網(wǎng)以前的文章或繼續(xù)瀏覽下面的相關(guān)文章，希望大家以后多多支持TOY模板網(wǎng)！
