多模態(tài)大模型輕量化探索-視覺大模型SAM的視覺編碼器 原創(chuàng)
往期，筆者基于LLava的數(shù)據(jù)對齊訓(xùn)練，搞了一個??Reyes?多模態(tài)大模型，并且看了些多模態(tài)大模型，相關(guān)開源的多模態(tài)大模型如：KimiVL、Internvl、QwenVL等，其視覺編碼器的尺寸都比較大，如：MoonViT-SO-400M、InternViT-6B-448px-V2_5 等都非常大，對于特定的垂直場景（或者是端側(cè)落地都不大友好），也許并不需要這么大視覺編碼器。如：表格場景（??【多模態(tài) & 文檔智能】一次多模態(tài)大模型表格識別解析探索小實踐記錄???），當(dāng)時筆者用了一個8B參數(shù)的模型及百萬表格數(shù)據(jù)進行訓(xùn)練達到了不錯的效果。近期，因此思考一些模型輕量化的方案，尋找一個輕量點的視覺編碼器（比如參數(shù)量小于100M），下面來看看SAM，供參考。
Segment Anything Model（SAM）是Meta AI發(fā)布的一個突破性圖像分割模型為計算機視覺領(lǐng)域提供一個通用的、靈活的基座視覺大模型。它受到自然語言處理（NLP）中基礎(chǔ)模型（如GPT、BERT）的啟發(fā)，強調(diào)零樣本遷移和提示式交互能力。在SA-1B數(shù)據(jù)集上的訓(xùn)練，該數(shù)據(jù)集包含超過11百萬張圖像和11億個高質(zhì)量分割掩碼，覆蓋了從日常場景到專業(yè)領(lǐng)域的多樣化內(nèi)容。
SAM借鑒了NLP領(lǐng)域的Prompt策略，通過給圖像分割任務(wù)提供Prompt提示來完成任意目標的快速分割。Prompt類型可以是「前景/背景點集、粗略的框或遮罩、任意形式的文本或者任何指示圖像中需要進行分割」的信息。如圖(a)所示，模型的輸入是原始的圖像和一些prompt，目標是輸出"valid"的分割，所謂valid，就是當(dāng)prompt的指向是模糊時，模型能夠輸出至少其中一個mask。
模型結(jié)構(gòu)
SAM的模型結(jié)構(gòu)由三個核心組件組成，Image Encoder、Prompt Encoder和Mask Decoder。分別負責(zé)圖像特征提取、提示編碼和掩碼生成。圖像經(jīng)過Image Encoder編碼，Prompt提示經(jīng)過Prompt Encoder編碼，兩部分Embedding再經(jīng)過一個輕量化的Mask Decoder得到融合后的特征。其中，Encoder部分使用的是已有模型，Decoder部分使用Transformer。 下表為三個組件的總結(jié)：
組件名稱 | 功能 | 關(guān)鍵特點 |
Image Encoder | 將輸入圖像轉(zhuǎn)換為密集特征表示 | 使用MAE預(yù)訓(xùn)練的Vision Transformer（ViT-H/16），輸入1024x1024x3，輸出64x64x256嵌入。 |
Prompt Encoder | 將用戶提示（點、框、文本、掩碼）編碼為嵌入 | 支持稀疏提示（點、框、文本）和密集提示（掩碼），使用CLIP處理文本，靈活適應(yīng)多種輸入。 |
Mask Decoder | 結(jié)合圖像和提示嵌入，生成最終分割掩碼 | 輕量級Transformer解碼器，通過自注意力與交叉注意力機制預(yù)測掩碼，實時高效。 |
Image Encoder
本文的目的是為了尋找一個輕量化的視覺編碼器，因此下面來詳細看下視覺編碼器部分。Image Encoder的作用是把圖像映射到特征空間，整體過程如下圖所示。
正如論文中所講，本質(zhì)上這個Encoder可以是任何網(wǎng)絡(luò)結(jié)構(gòu)，在這里使用的是微調(diào)的Detectron的ViT，當(dāng)然它也可以被改成傳統(tǒng)的卷積結(jié)構(gòu)，非常合理。
可以看到，Image Encoder就是一個ViT的結(jié)構(gòu)，由PatchEmbed、Transformer Encoder、Neck Convolution組成。
輸入圖像經(jīng)過ViT結(jié)構(gòu)的過程如下：
1.Patch Embedding
輸入圖像通過一個卷積base，將圖像劃分為16x16的patches，步長也為16，這樣feature map的尺寸就縮小了16倍，同時channel從3映射到768。Patch Embedding示意圖如下所示。
將輸入的圖像轉(zhuǎn)換為序列化的特征向量
Patch Embedding過程在Vision Transformer結(jié)構(gòu)圖中對應(yīng)下圖所示。
2.Transformer Encode
feature map通過16個Transformer Block，其中12個Block使用了基于Window Partition（就是把特征圖分成14*14的windows做局部的Attention）的注意力機制，以處理局部信息。另外4個Block是全局注意力模塊（多頭注意力），它們穿插在Window Partition模塊之間，以捕捉圖像的全局上下文。
循環(huán)疊加Transformer Encode
3.Neck Convolution
最后，通過兩層卷積（Neck）將通道數(shù)降低至256，生成最終的Image Embedding。其結(jié)構(gòu)圖如下所示。
SAM構(gòu)建與輕量化編碼器提取
通過下面代碼提取一個參數(shù)量大小僅為80幾M的視覺編碼器。
import torch
from functools import partial
from modeling import ImageEncoderViT, MaskDecoder, PromptEncoder, Sam, TwoWayTransformer
def build_sam_vit_b(checkpoint=None):
return _build_sam(
encoder_embed_dim=768,
encoder_depth=12,
encoder_num_heads=12,
encoder_global_attn_indexes=[2, 5, 8, 11],
checkpoint=checkpoint,
)
sam_model_registry = {
"vit_b": build_sam_vit_b,
}
def _build_sam(
encoder_embed_dim,
encoder_depth,
encoder_num_heads,
encoder_global_attn_indexes,
checkpoint=None,
):
prompt_embed_dim = 256
image_size = 1024
vit_patch_size = 16
image_embedding_size = image_size // vit_patch_size
sam = Sam(
image_encoder=ImageEncoderViT(
depth=encoder_depth,
embed_dim=encoder_embed_dim,
img_size=image_size,
mlp_ratio=4,
norm_layer=partial(torch.nn.LayerNorm, eps=1e-6),
num_heads=encoder_num_heads,
patch_size=vit_patch_size,
qkv_bias=True,
use_rel_pos=True,
global_attn_indexes=encoder_global_attn_indexes,
window_size=14,
out_chans=prompt_embed_dim,
),
prompt_encoder=PromptEncoder(
embed_dim=prompt_embed_dim,
image_embedding_size=(image_embedding_size, image_embedding_size),
input_image_size=(image_size, image_size),
mask_in_chans=16,
),
mask_decoder=MaskDecoder(
num_multimask_outputs=3,
transformer=TwoWayTransformer(
depth=2,
embedding_dim=prompt_embed_dim,
mlp_dim=2048,
num_heads=8,
),
transformer_dim=prompt_embed_dim,
iou_head_depth=3,
iou_head_hidden_dim=256,
),
pixel_mean=[123.675, 116.28, 103.53],
pixel_std=[58.395, 57.12, 57.375],
)
sam.eval()
if checkpoint is not None:
with open(checkpoint, "rb") as f:
state_dict = torch.load(f)
sam.load_state_dict(state_dict)
return sam
if __name__ == '__main__':
x = torch.zeros(2, 3, 1024, 1024)
net = build_sam_vit_b(checkpoint='sam_vit_b_01ec64.pth')
image_encoder = net.image_encoder
print(image_encoder)
print(image_encoder(x).shape) # 輸出：torch.Size([2, 256, 64, 64])
total_params = sum(p.numel() for p in image_encoder.parameters())
print(f"模型的參數(shù)量為: {(total_params/ 1e6):.2f}M") # 模型的參數(shù)量為: 89.67M參考文獻：
Segment Anything，https://arxiv.org/pdf/2304.02643
code：https://github.com/facebookresearch/segment-anything
公眾號大模型自然語言處理 作者：余俊暉
