Fix vision_tensors shape for bs=1 multi-image inputs

model/model_vlm.py

@@ -20,12 +20,10 @@ class VLMConfig(MiniMindConfig):
         self.image_token_len = kwargs.get("image_token_len", 64)
         super().__init__(**kwargs)
 
-# 视觉投影层，把视觉特征投影到语言特征空间
 class MMVisionProjector(nn.Module):
     def __init__(self, in_dim, out_dim, source_tokens=256, target_tokens=64):
         super().__init__()
         self.target_tokens = target_tokens
-        #融合系数，把source_tokens个视觉token压缩为target个
         self.merge = source_tokens // target_tokens
         self.mlp = nn.Sequential(
             nn.Linear(in_dim * self.merge, out_dim),
@@ -34,25 +32,20 @@ class MMVisionProjector(nn.Module):
         )
     def forward(self, x):
         b, n, d = x.shape
-        #若输入的x.shape为(8, 256, 768)，则输出为(8, 256/4=64, 768*4=3072),压缩一张图片的视觉token数，最后还要通过mlp，把3072变成768，这其实有信息损失的问题，但是llm只需要上层的语义特征即可
-        #torch的矩阵乘法只需要满足前一个矩阵最后一维和后一个矩阵的倒数第二维相同即可
         x = x.reshape(b, self.target_tokens, d * self.merge)
         return self.mlp(x)
 
 # 继承自语言模型
 class MiniMindVLM(MiniMindForCausalLM):
-    #类属性，用于指定配置类
     config_class = VLMConfig
 
     def __init__(self, config: VLMConfig = None, vision_model_path="./model/siglip2-base-p16-ve"):
-        #实例属性，用于存储配置
         self.config = config or VLMConfig()
         super().__init__(self.config)
         self.vision_encoder, self.processor = self.__class__.get_vision_model(vision_model_path)
         self.vision_proj = MMVisionProjector(self.config.image_hidden_size, self.config.hidden_size, target_tokens=self.config.image_token_len)
 
-    #被init方法调用，self.processor接到get_vision_model的返回值，即processor
+    @staticmethod
-    @staticmethod #写给用户看的，说明这个函数是静态方法，不需要实例化类就可以调用
     def get_vision_model(model_path: str):
         from transformers import logging as hf_logging
         hf_logging.set_verbosity_error()
@@ -60,55 +53,34 @@ class MiniMindVLM(MiniMindForCausalLM):
             return None, None
         model = Siglip2VisionModel.from_pretrained(model_path)
         processor = Siglip2ImageProcessor.from_pretrained(model_path)
-        # 冻结 vision_encoder 的所有参数，现成的视觉模型不做训练
+        # 冻结 vision_encoder 的所有参数
         for param in model.parameters():
             param.requires_grad = False
-        #siglip跑的是推理模式
         return model.eval(), processor
 
-    #图片通过processor转换为RGB格式的tensor
     @staticmethod
     def image2tensor(image, processor):
-        #如果图片是RGBA或LA格式，则转换为RGB格式
         if image.mode in ['RGBA', 'LA']: image = image.convert('RGB')
         inputs = processor(images=image, return_tensors="pt")
         return inputs
 
-
     @staticmethod
     def get_image_embeddings(image_inputs, vision_model):
-        #如果image_inputs有keys属性，说明是字典
         if hasattr(image_inputs, 'keys'):
-            #等价于for k, v in image_inputs.items():
-            #    if v.ndim > 2 and v.shape[1] == 1:
-            #        v = v.squeeze(1)
-            #    else:
-            #        v = v
-            #(1, 1, 3, 224, 224)→(1, 3, 224, 224)
             image_inputs = {k: v.squeeze(1) if v.ndim > 2 and v.shape[1] == 1 else v for k, v in image_inputs.items()}
         with torch.no_grad():
             outputs = vision_model(**image_inputs)
-        #只要outputs大礼包里面的last_hidden_state（图像特征）
         return outputs.last_hidden_state
 
     @torch.compiler.disable
-    #构造文本视觉融合输入
-    #h和vf都是经过embedded后的特征，只不过一个是文本的，一个是视觉的
     def count_vision_proj(self, tokens, h, vision_tensors=None, seqlen=512):
-        #self.config.image_ids是一个列表，通常只有1个元素，是视觉token的id索引
-        #这两个条件说的是一个事情
         if vision_tensors is None or not self.config.image_ids:
             return h
         marker, vf = self.config.image_ids[0], vision_tensors
-        #把vf从(8, 64, 768)变成(8, 1, 64, 768)
         if vf.dim() == 3:
             vf = vf.unsqueeze(1)
         out = []
-        #遍历每个样本，h表示文本输入
         for b in range(h.size(0)):
-            #h是text_tokens的特征(batch_size, seq_length, hidden_size)
-            #tokens是inputs_ids(batch_size, seq_length)
-            #Attention:VLM的
             hb, seq, k, i = h[b], tokens[b].tolist(), 0, 0
             while i < len(seq):
                 if seq[i] == marker:
@@ -116,14 +88,9 @@ class MiniMindVLM(MiniMindForCausalLM):
                     while i < len(seq) and seq[i] == marker:
                         i += 1
                     if k < vf.size(1):
-                        #做一个拼接，中间放视觉特征，前后放文本特征，往往i-start就是target_tokens=64
-                        #这里的k就是0，因为unsqueeze(1)，单文本只能对应单图
-                        #不过可以拓展，只要不用前面unsqueeze(1)的逻辑，把第二个维度改成num_images，每次就代表着一个样本中的第几张图片
                         hb = torch.cat((hb[:start], vf[b][k][:i - start], hb[i:]), dim=0)[:seqlen]
-                        #k是图片样本的索引
                         k += 1
                 else:
-                    #不是就继续找
                     i += 1
             out.append(hb)
         return torch.stack(out)
@@ -142,17 +109,14 @@ class MiniMindVLM(MiniMindForCausalLM):
         past_key_values = past_key_values or [None] * len(self.model.layers)
         start_pos = past_key_values[0][0].shape[1] if past_key_values[0] is not None else 0
 
-        #文本ids经过embedded得到hidden_states
         hidden_states = self.model.dropout(self.model.embed_tokens(input_ids))
 
-        #图片经过embedded得到vision_tensors
         if pixel_values is not None and start_pos == 0:
             if hasattr(pixel_values, 'keys'):
-                #这样写强调静态方法，不需要实例化类就可以调用。否则写为self.get_image_embeddings
                 img_emb = MiniMindVLM.get_image_embeddings(pixel_values, self.vision_encoder)
                 vision_tensors = self.vision_proj(img_emb)
+
             else:
-                # 一段多图：N 表示一个样本中的图片数，[B, N, 1, C, H, W] -> [B, N, C, H, W]
                 if len(pixel_values.shape) == 6:
                     pixel_values = pixel_values.squeeze(2)
                 bs, num, c, im_h, im_w = pixel_values.shape