batchlize transfomer and add forward_train/test in pipeline

modules/seq_encoder/transformer_seq_encoder.py

@@ -16,32 +16,49 @@ class TransformerSequenceEncoder(SequenceEncoder):
         self.transformer_encoder = nn.TransformerEncoder(encoder_layer, num_layers=config['num_layers'])
         self.fc = nn.Linear(embed_dim, config['output_dim'])
 
-    def encode_sequence(self, pts_embedding_list, pose_embedding_list):
-        combined_features = [torch.cat((pts_embed, pose_embed), dim=-1) for pts_embed, pose_embed in zip(pts_embedding_list[:-1], pose_embedding_list[:-1])]
-        combined_tensor = torch.stack(combined_features)
-        pos_encoding = self.positional_encoding[:, :combined_tensor.size(0), :]
-        combined_tensor = combined_tensor.unsqueeze(0) + pos_encoding
-        transformer_output = self.transformer_encoder(combined_tensor).squeeze(0)
-        final_feature = transformer_output.mean(dim=0)
+    def encode_sequence(self, pts_embedding_list_batch, pose_embedding_list_batch):
+        batch_size = len(pts_embedding_list_batch)
+        combined_features_batch = []
+        
+        for i in range(batch_size):
+            combined_features = [torch.cat((pts_embed, pose_embed), dim=-1)
+                                 for pts_embed, pose_embed in zip(pts_embedding_list_batch[i][:-1], pose_embedding_list_batch[i][:-1])]
+            combined_features_batch.append(torch.stack(combined_features))
+        
+        combined_tensor = torch.stack(combined_features_batch)  # Shape: [batch_size, seq_len-1, embed_dim]
+        
+        # Adjust positional encoding to match batch size
+        pos_encoding = self.positional_encoding[:, :combined_tensor.size(1), :].repeat(batch_size, 1, 1)
+        combined_tensor = combined_tensor + pos_encoding
+
+        # Transformer encoding
+        transformer_output = self.transformer_encoder(combined_tensor)
+        
+        # Mean pooling
+        final_feature = transformer_output.mean(dim=1)
+        
+        # Fully connected layer
         final_output = self.fc(final_feature)
         
         return final_output
 
 if __name__ == "__main__":
     config = {
-    'pts_embed_dim': 1024,  # 每个点云embedding的维度
-    'pose_embed_dim': 256,  # 每个姿态embedding的维度
-    'num_heads': 4,   # 多头注意力机制的头数
-    'ffn_dim': 256,   # 前馈神经网络的维度
-    'num_layers': 3,  # Transformer 编码层数
-    'max_seq_len': 10, # 最大序列长度
-    'output_dim': 2048, # 输出特征维度
+        'pts_embed_dim': 1024,  # 每个点云embedding的维度
+        'pose_embed_dim': 256,  # 每个姿态embedding的维度
+        'num_heads': 4,   # 多头注意力机制的头数
+        'ffn_dim': 256,   # 前馈神经网络的维度
+        'num_layers': 3,  # Transformer 编码层数
+        'max_seq_len': 10, # 最大序列长度
+        'output_dim': 2048, # 输出特征维度
     }
 
     encoder = TransformerSequenceEncoder(config)
     seq_len = 5
-    pts_embedding_list = [torch.randn(config['pts_embed_dim']) for _ in range(seq_len)]
-    pose_embedding_list = [torch.randn(config['pose_embed_dim']) for _ in range(seq_len)]
-    output_feature = encoder.encode_sequence(pts_embedding_list, pose_embedding_list)
+    batch_size = 4
+    
+    pts_embedding_list_batch = [torch.randn(seq_len, config['pts_embed_dim']) for _ in range(batch_size)]
+    pose_embedding_list_batch = [torch.randn(seq_len, config['pose_embed_dim']) for _ in range(batch_size)]
+    output_feature = encoder.encode_sequence(pts_embedding_list_batch, pose_embedding_list_batch)
     print("Encoded Feature:", output_feature)
-    print("Feature Shape:", output_feature.shape)
+    print("Feature Shape:", output_feature.shape)