LLaVA 系列模型结构详解

from accelerate import init_empty_weights, load_checkpoint_and_dispatch
from transformers import LlavaConfig
import sys, os

# 获取 lite_llama 目录的绝对路径并添加到 sys.path 中
sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))
from lite_llama.models.llava import LlavaLlama

hf_model_path = "/gemini/code/liuhaotian/llava-v1.5-7b"

def test_LlavaLlama_structure(hf_model_path):
    
    # 使用 init_empty_weights 初始化空模型
    with init_empty_weights():
        config = LlavaConfig.from_pretrained(hf_model_path)
        model = LlavaLlama(config)
        
        # 打印没有加载权重的 LlavaLlama 模型结构
        print(model)
        # 打印模型的简单摘要
        print(f"模型总参数量: {sum(p.numel() for p in model.parameters()) / 1e6:.2f} M")

        # 可选择打印部分参数信息
        for name, param in list(model.named_parameters())[:]:  # 打印模型参数
            print(name, param.shape)

if __name__ == "__main__":
    test_LlavaLlama_structure(hf_model_path)

class LlavaLlama(nn.Module):
    def __init__(self, llava_config: LlavaConfig):
        super().__init__()
        self.device = "cuda"  # 默认运行在 GPU 上
        self.llava_config = llava_config

        # 提取文本模型配置并转为 LlamaConfig 类型（为了自定义加载）
        text_config = self.llava_config.text_config
        self.llama_config = LlamaConfig.from_dict(text_config.to_dict())

        # 指定提取哪一层视觉特征
        self.select_layer = llava_config.vision_feature_layer
        self.select_feature = llava_config.vision_feature_select_strategy

        # 初始化视觉编码器（比如 CLIP）
        self.vision_tower = AutoModel.from_config(llava_config.vision_config)

        # 初始化多模态投影模块
        self.multi_modal_projector = LlavaMultiModalProjector(
            vision_hidden_size=llava_config.vision_config.hidden_size,
            text_hidden_size=llava_config.text_config.hidden_size,
            projector_hidden_act=llava_config.projector_hidden_act
        )

        # 初始化 LLaMA 语言模型
        self.language_model = LlamaModel(self.llama_config)

        # 设置 pad token（防止 None 类型报错）
        self.pad_token_id = self.llava_config.pad_token_id if self.llava_config.pad_token_id is not None else -1

from transformers import AutoModel, LlavaConfig

from transformers import AutoModel, LlavaConfig

def vision_encode(self, image_tensor):
    x = image_tensor.half().to(device=self.device)

    # 1. 通过视觉处理模块提取图像特征
    x = self.vision_tower(x, output_hidden_states = True)
    x = x.hidden_states[self.select_layer]
    x = self._select_image_features(x, self.select_feature)

    # 2. 通过多模态投影器将图像特征转换为多模态(llm) embedding 维度
    image_features = self.multi_modal_projector(x)

    assert not torch.isnan(image_features).any(), f"After vision_tower image_features tensor contains NaN values!"
    return image_features

def get_multi_modal_input_embeddings(
    self,
    input_ids: torch.Tensor,
    vision_embeddings = None,
) -> torch.Tensor:
    """获取输入嵌入，包括文本和视觉嵌入的合并。"""
    # torch.Size([1, 22]) --> torch.Size([1, 22, 4096])
    llm_inputs_embeds = self.language_model.get_input_embeddings(input_ids) 
    
    if vision_embeddings is not None:
        inputs_embeds, position_ids = merge_input_ids_with_image_features(
            input_ids,              # 文本 token ID
            llm_inputs_embeds,      # 文本嵌入向量 torch.Size([1, 22, 4096]) 
            vision_embeddings,      # 视觉嵌入向量 shape torch.Size([1, 576, 4096])
            self.llava_config.pad_token_id,       # pad token ID
            self.llava_config.image_token_index,  # 图像 token 的插入索引 32000
        )

    return inputs_embeds, position_ids

def forward(
    self, 
    input_ids: torch.Tensor, 
    position_ids: torch.Tensor,  
    atten_info, 
    image_tensor: Optional[torch.FloatTensor] = None,
):
    input_ids = input_ids.to(self.device) # 将 input_ids 移动到设备
    if position_ids is not None: # 如果提供了 position_ids，将其移动到设备
        position_ids = position_ids.to(self.device)
        
    if input_ids.shape[1] != 1: # 判断是不是首次 token 输出
        vision_embeddings = self.vision_encode(image_tensor) #  torch.Size([1, 3, 336, 336]) --> torch.Size([1, 576, 4096])
        inputs_embeds, position_ids = self.get_multi_modal_input_embeddings(input_ids, vision_embeddings)
    else: # 进入 decode 阶段, 无需再做视觉编码
        inputs_embeds = None
    
    hidden_states = self.language_model(input_ids = input_ids,
                                        position_ids = position_ids,
                                        atten_info = atten_info,
                                        inputs_embeds = inputs_embeds
                                        )
    
    return hidden_states

def merge_input_ids_with_image_features(
    input_ids: torch.Tensor, 
    inputs_embeds: torch.Tensor, 
    image_features: torch.Tensor,
    pad_token_id: int,
    image_token_index: int
):

def merge_input_ids_with_image_features(
    input_ids: torch.Tensor, 
    inputs_embeds: torch.Tensor, 
    image_features: torch.Tensor,
    pad_token_id: int,
    image_token_index: int
):
    """
    将 input_ids 与 image_features 合并，生成最终的嵌入和位置 ID。
    
    Args:
        input_ids (torch.Tensor): 输入的 token IDs, 形状为 (batch_size, sequence_length)
        inputs_embeds (torch.Tensor): 文本嵌入，形状为 (batch_size, sequence_length, embed_dim)
        image_features (torch.Tensor): 视觉编码后的图像特征，形状为 (num_images, num_image_patches, embed_dim)
        pad_token_id (int): 填充 token 的 ID
        image_token_index (int): 图像 token 的 ID
    
    Returns:
        final_embedding (torch.Tensor): 合并后的嵌入张量，形状为 (batch_size, max_embed_dim, embed_dim)
        position_ids (torch.Tensor): 位置 ID, 形状为 (batch_size, max_embed_dim)
    """
    # 1, 基础 shape 信息提取
    num_images, num_image_patches, embed_dim = image_features.shape # torch.Size([1, 576, 4096])
    batch_size, sequence_length = input_ids.shape # torch.Size([1, 22])

    # 2, 掩码与填充处理
    attention_mask = (input_ids != pad_token_id).long()
    left_padding = not torch.sum(input_ids[:, -1] == pad_token_id).bool().any()
    batch_image_token_mask  = input_ids == image_token_index
    
    # 3, 计算新序列长度
    num_special_image_tokens = torch.sum(batch_image_token_mask , dim=-1) # 统计每个样本（batch 中每条序列）里出现了多少个“图像占位符” token。
    max_embed_dim = (num_special_image_tokens.max() * (num_image_patches - 1)) + sequence_length
    batch_indices, non_image_indices = torch.where(input_ids != image_token_index) 

    # 4, 位置映射计算
    # 得到每个原始 token 在新序列中占据的开始位置索引。
    new_token_positions = torch.cumsum((batch_image_token_mask  * (num_image_patches - 1) + 1), -1) - 1 
    nb_image_pad = max_embed_dim - 1 - new_token_positions[:, -1] 
    if left_padding:
        new_token_positions += nb_image_pad[:, None]  # offset for left padding
    text_to_overwrite = new_token_positions[batch_indices, non_image_indices]
    
    # 5，构建融合张量
    final_embedding = torch.zeros(
        batch_size, max_embed_dim, embed_dim, dtype=inputs_embeds.dtype, device=inputs_embeds.device
    )
    final_embedding[batch_indices, text_to_overwrite] = inputs_embeds[batch_indices, non_image_indices] # 填充文本嵌入

    # 确定图像特征插入位置，通过找到 final_embedding 中所有全 0 的位置
    image_to_overwrite = torch.all(final_embedding == 0, dim=-1)  # 找出 final_embedding 中所有维度为0的位置
    image_to_overwrite &= image_to_overwrite.cumsum(-1) - 1 >= nb_image_pad[:, None].to(target_device)

    # 将 image_features 重新排列为 (num_images * num_image_patches, embed_dim)，并填充到 final_embedding 的相应位置。
    final_embedding[image_to_overwrite] = image_features.contiguous().view(-1, embed_dim).to(target_device)
    
    # 6，生成新的 position_ids
    position_ids = torch.arange(max_embed_dim, dtype=torch.long, device=inputs_embeds.device).unsqueeze(0).expand(batch_size, -1)

    # 7，处理填充位置的嵌入, 将填充位置的嵌入设为0
    batch_indices_pad, pad_indices = torch.where(input_ids == pad_token_id)
    indices_to_mask = new_token_positions[batch_indices_pad, pad_indices]

    final_embedding[batch_indices_pad, indices_to_mask] = 0

    return final_embedding, position_ids

# === 输入示例 ===
# batch_size=1, seq_len=5, embed_dim=3; num_images=1, num_patches=2
input_ids       = torch.tensor([[11, 99, 22, 0, 0]])       # 99 代表 image_token_index，0 是 pad_token_id
inputs_embeds   = torch.arange(1, 1+1*5*3).reshape(1,5,3).float()
image_features  = torch.tensor([[[9,9,9],[8,8,8]]]).float()  # shape (1,2,3)
pad_token_id    = 0
image_token_index = 99

fe, pids = merge_input_ids_with_image_features(
    input_ids,
    inputs_embeds,
    image_features,
    pad_token_id,
    image_token_index
)

print("final_embedding shape:", fe.shape)
print("position_ids      shape:", pids.shape)
print("final_embedding:\n", fe)
print("position_ids:\n", pids)

attention_mask = (input_ids != pad_token_id)        # [[True,True,True,False,False]]
left_padding   = not (input_ids[:, -1] == pad_token_id).any()  # False
is_image_token = (input_ids == image_token_index)   # [[False,True,False,False,False]]

# 统计每个样本（batch 中每条序列）里出现了多少个“图像占位符” token。
num_images = is_image_token.sum(dim=-1)  # [1]
max_len = Seq_len + num_images.max()*(num_patches-1)  
# 5 + 1*(2-1) = 6

input_ids = tensor([
    [10, 99, 20, 99,  0],   # 第一条序列里有两个“99”
    [99,  0, 99, 30, 99]    # 第二条序列里有三个“99”
])
image_token_index = 99
batch_image_token_mask  = (input_ids == image_token_index)
# → tensor([
#     [False, True, False, True, False],
#     [True,  False, True,  False, True ]
#   ])

# 统计每个样本（batch 中每条序列）里出现了多少个“图像占位符” token。
max_num_images = batch_image_token_mask.sum(dim=-1)
print("num_images per batch sample:", max_num_images) 

# num_images per batch sample: tensor([2, 3])