多模態大模型輕量化探索-視覺大模型SAM的視覺編碼器 原創
往期，筆者基于LLava的數據對齊訓練，搞了一個??Reyes?多模態大模型，并且看了些多模態大模型，相關開源的多模態大模型如：KimiVL、Internvl、QwenVL等，其視覺編碼器的尺寸都比較大，如：MoonViT-SO-400M、InternViT-6B-448px-V2_5 等都非常大，對于特定的垂直場景（或者是端側落地都不大友好），也許并不需要這么大視覺編碼器。如：表格場景（??【多模態 & 文檔智能】一次多模態大模型表格識別解析探索小實踐記錄???），當時筆者用了一個8B參數的模型及百萬表格數據進行訓練達到了不錯的效果。近期，因此思考一些模型輕量化的方案，尋找一個輕量點的視覺編碼器（比如參數量小于100M），下面來看看SAM，供參考。
Segment Anything Model（SAM）是Meta AI發布的一個突破性圖像分割模型為計算機視覺領域提供一個通用的、靈活的基座視覺大模型。它受到自然語言處理（NLP）中基礎模型（如GPT、BERT）的啟發，強調零樣本遷移和提示式交互能力。在SA-1B數據集上的訓練，該數據集包含超過11百萬張圖像和11億個高質量分割掩碼，覆蓋了從日常場景到專業領域的多樣化內容。
SAM借鑒了NLP領域的Prompt策略，通過給圖像分割任務提供Prompt提示來完成任意目標的快速分割。Prompt類型可以是「前景/背景點集、粗略的框或遮罩、任意形式的文本或者任何指示圖像中需要進行分割」的信息。如圖(a)所示，模型的輸入是原始的圖像和一些prompt，目標是輸出"valid"的分割，所謂valid，就是當prompt的指向是模糊時，模型能夠輸出至少其中一個mask。
模型結構
SAM的模型結構由三個核心組件組成，Image Encoder、Prompt Encoder和Mask Decoder。分別負責圖像特征提取、提示編碼和掩碼生成。圖像經過Image Encoder編碼，Prompt提示經過Prompt Encoder編碼，兩部分Embedding再經過一個輕量化的Mask Decoder得到融合后的特征。其中，Encoder部分使用的是已有模型，Decoder部分使用Transformer。 下表為三個組件的總結：
組件名稱 | 功能 | 關鍵特點 |
Image Encoder | 將輸入圖像轉換為密集特征表示 | 使用MAE預訓練的Vision Transformer（ViT-H/16），輸入1024x1024x3，輸出64x64x256嵌入。 |
Prompt Encoder | 將用戶提示（點、框、文本、掩碼）編碼為嵌入 | 支持稀疏提示（點、框、文本）和密集提示（掩碼），使用CLIP處理文本，靈活適應多種輸入。 |
Mask Decoder | 結合圖像和提示嵌入，生成最終分割掩碼 | 輕量級Transformer解碼器，通過自注意力與交叉注意力機制預測掩碼，實時高效。 |
Image Encoder
本文的目的是為了尋找一個輕量化的視覺編碼器，因此下面來詳細看下視覺編碼器部分。Image Encoder的作用是把圖像映射到特征空間，整體過程如下圖所示。
正如論文中所講，本質上這個Encoder可以是任何網絡結構，在這里使用的是微調的Detectron的ViT，當然它也可以被改成傳統的卷積結構，非常合理。
可以看到，Image Encoder就是一個ViT的結構，由PatchEmbed、Transformer Encoder、Neck Convolution組成。
輸入圖像經過ViT結構的過程如下：
1.Patch Embedding
輸入圖像通過一個卷積base，將圖像劃分為16x16的patches，步長也為16，這樣feature map的尺寸就縮小了16倍，同時channel從3映射到768。Patch Embedding示意圖如下所示。
將輸入的圖像轉換為序列化的特征向量
Patch Embedding過程在Vision Transformer結構圖中對應下圖所示。
2.Transformer Encode
feature map通過16個Transformer Block，其中12個Block使用了基于Window Partition（就是把特征圖分成14*14的windows做局部的Attention）的注意力機制，以處理局部信息。另外4個Block是全局注意力模塊（多頭注意力），它們穿插在Window Partition模塊之間，以捕捉圖像的全局上下文。
循環疊加Transformer Encode
3.Neck Convolution
最后，通過兩層卷積（Neck）將通道數降低至256，生成最終的Image Embedding。其結構圖如下所示。
SAM構建與輕量化編碼器提取
通過下面代碼提取一個參數量大小僅為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"模型的參數量為: {(total_params/ 1e6):.2f}M") # 模型的參數量為: 89.67M參考文獻：
Segment Anything，https://arxiv.org/pdf/2304.02643
code：https://github.com/facebookresearch/segment-anything
公眾號大模型自然語言處理 作者：余俊暉
