Upload processor

modeling_yangjian.py

@@ -279,7 +279,7 @@ class YangJianCompareVisualEncoder(nn.Module):
         # Decoder 部分：Query 与编码特征交互
         # 可学习的 Query Embeddings
         self.query_embeddings = nn.Parameter(
-            torch.randn(self.token_size, self.hidden_size) * 0.02
+            torch.empty(self.token_size, self.hidden_size)
         )
         
         # 只保留 Cross Attention for queries to attend to encoded features
@@ -314,47 +314,94 @@ class YangJianCompareVisualEncoder(nn.Module):
         self.encoder_mlp2 = self.encoder_mlp2.to(device=device, dtype=dtype)
         self.decoder_mlp = self.decoder_mlp.to(device=device, dtype=dtype)
 
-    def forward(self, images_hidden_states: list) -> list:
+    def _initialize_weights(self):
+        nn.init.normal_(self.query_embeddings.weight, mean=0.0, std=0.02)
+
+    def forward(self, images_hidden_states: list) -> torch.Tensor:
         """
         Args:
             images_hidden_states: List of tensor, each tensor has shape [seq_len, hidden_size]
         
         Returns:
-            List of compare visual embeddings, each has shape [token_size, hidden_size]
+            Tensor of shape [total_images, token_size, hidden_size]
         """
         if not images_hidden_states:
-            return []
+            return torch.empty(0, self.token_size, self.hidden_size)
         
         # 确保所有组件的设备和数据类型一致
-        self._ensure_device_dtype_consistency(images_hidden_states[0])
-        
-        compare_visual_embeds = []
+        # self._ensure_device_dtype_consistency(images_hidden_states[0])
+        
+        # 检查 query_embeddings 是否包含 NaN
+        if torch.isnan(self.query_embeddings).any():
+            print("警告：query_embeddings 包含 NaN 值，重新初始化")
+            nn.init.normal_(self.query_embeddings, mean=0.0, std=0.02)
+        
+        # 获取每个图像的序列长度
+        seq_lengths = [state.size(0) for state in images_hidden_states]
+        max_seq_len = max(seq_lengths)
+        batch_size = len(images_hidden_states)
+        device = images_hidden_states[0].device
+        dtype = images_hidden_states[0].dtype
+        
+        # 将所有图像填充到相同长度并堆叠
+        padded_states = []
+        attention_masks = []
+        for state in images_hidden_states:
+            pad_len = max_seq_len - state.size(0)
+            if pad_len > 0:
+                # 填充序列
+                padded_state = F.pad(state, (0, 0, 0, pad_len), mode='constant', value=0)
+                # 创建注意力掩码
+                attention_mask = torch.ones(max_seq_len, dtype=torch.bool, device=device)
+                attention_mask[state.size(0):] = False
+            else:
+                padded_state = state
+                attention_mask = torch.ones(max_seq_len, dtype=torch.bool, device=device)
+            padded_states.append(padded_state)
+            attention_masks.append(attention_mask)
+        
+        # [batch_size, max_seq_len, hidden_size]
+        batched_states = torch.stack(padded_states)
+        # [batch_size, max_seq_len]
+        attention_masks = torch.stack(attention_masks)
+        
+        # 创建循环移位的状态用于对比
+        # 对于第一个图像，使用自身作为previous
+        previous_states = torch.roll(batched_states, shifts=1, dims=0)
+        previous_states[0] = batched_states[0]
+        previous_masks = torch.roll(attention_masks, shifts=1, dims=0)
+        previous_masks[0] = attention_masks[0]
+        
+        # Encoder: 批量处理所有图像
+        encoded_features = self._encoder_forward(
+            batched_states,  # [batch_size, max_seq_len, hidden_size]
+            previous_states,  # [batch_size, max_seq_len, hidden_size]
+            attention_masks,  # [batch_size, max_seq_len]
+            previous_masks   # [batch_size, max_seq_len]
+        )
         
-        for i in range(len(images_hidden_states)):
-            current_hidden_state = images_hidden_states[i]  # [seq_len_current, hidden_size]
-            previous_hidden_state = images_hidden_states[i-1] if i > 0 else current_hidden_state  # [seq_len_prev, hidden_size]
-            
-            # Encoder 部分：双向图像特征交互
-            encoded_features = self._encoder_forward(current_hidden_state, previous_hidden_state)
-            
-            # Decoder 部分：Query 与编码特征交互
-            compare_visual_embed = self._decoder_forward(encoded_features)
-            
-            compare_visual_embeds.append(compare_visual_embed)
+        # Decoder: 批量处理所有图像
+        # 扩展query_embeddings到batch维度
+        batch_queries = self.query_embeddings.unsqueeze(0).expand(batch_size, -1, -1)
+        # [batch_size, token_size, hidden_size]
+        compare_visual_embeds = self._decoder_forward(
+            batch_queries,
+            encoded_features,
+            torch.ones(batch_size, self.token_size, dtype=torch.bool, device=device),  # query掩码
+            attention_masks  # encoded特征的掩码
+        )
         
-        return compare_visual_embeds
+        return compare_visual_embeds  # [batch_size, token_size, hidden_size]
     
-    def _encoder_forward(self, current_features, previous_features):
+    def _encoder_forward(self, current_features, previous_features, current_mask=None, previous_mask=None):
         """
         Encoder: 双向图像特征交互
-        1. previous attend to current
-        2. current attend to previous
+        Args:
+            current_features: [batch_size, seq_len, hidden_size]
+            previous_features: [batch_size, seq_len, hidden_size]
+            current_mask: [batch_size, seq_len]
+            previous_mask: [batch_size, seq_len]
         """
-        # 确保数据类型和设备一致
-        device = current_features.device
-        dtype = current_features.dtype
-        previous_features = previous_features.to(device=device, dtype=dtype)
-        
         # 第一步：previous attend to current
         residual = previous_features
         
@@ -365,7 +412,8 @@ class YangJianCompareVisualEncoder(nn.Module):
         # Cross attention: previous attend to current
         cross_attn_output1 = self.encoder_cross_attn1(
             query_states=previous_normed,
-            key_value_states=current_normed1
+            key_value_states=current_normed1,
+            attention_mask=current_mask.unsqueeze(1).unsqueeze(2) if current_mask is not None else None
         )
         
         # Residual connection
@@ -382,12 +430,13 @@ class YangJianCompareVisualEncoder(nn.Module):
         
         # Layer norm
         current_normed2 = self.encoder_norm3(current_features)
-        previous_normed2 = self.encoder_norm3(previous_features)  # 使用增强后的 previous features
+        previous_normed2 = self.encoder_norm3(previous_features)
         
         # Cross attention: current attend to previous
         cross_attn_output2 = self.encoder_cross_attn2(
             query_states=current_normed2,
-            key_value_states=previous_normed2
+            key_value_states=previous_normed2,
+            attention_mask=previous_mask.unsqueeze(1).unsqueeze(2) if previous_mask is not None else None
         )
         
         # Residual connection
@@ -401,17 +450,15 @@ class YangJianCompareVisualEncoder(nn.Module):
         
         return current_features
     
-    def _decoder_forward(self, encoded_features):
+    def _decoder_forward(self, queries, encoded_features, query_mask=None, encoded_mask=None):
         """
-        Decoder: Query 与编码特征交互（仅使用 cross attention）
+        Decoder: Query 与编码特征交互
+        Args:
+            queries: [batch_size, token_size, hidden_size]
+            encoded_features: [batch_size, seq_len, hidden_size]
+            query_mask: [batch_size, token_size]
+            encoded_mask: [batch_size, seq_len]
         """
-        # 获取设备和数据类型
-        device = encoded_features.device
-        dtype = encoded_features.dtype
-        
-        # 初始化 queries 并确保设备和数据类型一致
-        queries = self.query_embeddings.to(device=device, dtype=dtype)
-        
         # Cross attention: queries attend to encoded features
         residual = queries
         queries_normed = self.decoder_norm1(queries)
@@ -419,7 +466,8 @@ class YangJianCompareVisualEncoder(nn.Module):
         
         cross_attn_output = self.decoder_cross_attn(
             query_states=queries_normed,
-            key_value_states=encoded_normed
+            key_value_states=encoded_normed,
+            attention_mask=encoded_mask.unsqueeze(1).unsqueeze(2) if encoded_mask is not None else None
         )
         
         queries = residual + cross_attn_output
@@ -430,7 +478,7 @@ class YangJianCompareVisualEncoder(nn.Module):
         mlp_output = self.decoder_mlp(mlp_input)
         queries = residual + mlp_output
         
-        return queries  # [token_size, hidden_size]
+        return queries  # [batch_size, token_size, hidden_size]
 
 
 # 先把组件继承出来方便修改
@@ -497,10 +545,20 @@ class YangJianVisionTransformerPretrainedModel(Qwen2_5_VisionTransformerPretrain
 
         split_sizes = grid_thw.prod(-1).tolist()
         splited_hidden_states_before_merger = torch.split(hidden_states, split_sizes)
+        # [total_images, token_size, hidden_size]
         compare_visual_embeds = self.compare_visual_encoder(splited_hidden_states_before_merger)
-        # compare_visual_embeds = self.merger(compare_visual_embeds)
-        for i, embeds in enumerate(compare_visual_embeds):
-            compare_visual_embeds[i] = self.merger(embeds)
+        # 记录每个batch的token数量
+        batch_size = compare_visual_embeds.size(0)
+        token_size = compare_visual_embeds.size(1)
+        # 将所有batch的数据拼接在一起
+        # [batch_size * token_size, hidden_size]
+        flattened_embeds = compare_visual_embeds.view(-1, compare_visual_embeds.size(-1))
+        # 一次性进行merger操作
+        # 假设merger会将token数量���为原来的1/4
+        merged = self.merger(flattened_embeds)  # [(batch_size * token_size)/4, merged_hidden_size]
+        merged_token_size = token_size // self.spatial_merge_size**2  
+        # [batch_size, merged_token_size, merged_hidden_size]
+        compare_visual_embeds = merged.view(batch_size, merged_token_size, -1)
         
         hidden_states = self.merger(hidden_states)
         reverse_indices = torch.argsort(window_index)

tokenizer_config.json

@@ -202,12 +202,8 @@
   "eos_token": "<|im_end|>",
   "errors": "replace",
   "extra_special_tokens": {},
-  "max_length": null,
   "model_max_length": 131072,
-  "pad_to_multiple_of": null,
   "pad_token": "<|endoftext|>",
-  "pad_token_type_id": 0,
-  "padding_side": "right",
   "processor_class": "YangJianProcessor",
   "split_special_tokens": false,
   "tokenizer_class": "Qwen2Tokenizer",