add data_splitor, modify dataset and encoder

modules/pose_encoder.py

@@ -7,12 +7,13 @@ class PoseEncoder(nn.Module):
         super(PoseEncoder, self).__init__()
         self.config = config
         pose_dim = config["pose_dim"]
+        out_dim = config["out_dim"]
         self.act = nn.ReLU(True)
 
         self.pose_encoder = nn.Sequential(
-            nn.Linear(pose_dim, 256),
+            nn.Linear(pose_dim, out_dim),
             self.act,
-            nn.Linear(256, 256),
+            nn.Linear(out_dim, out_dim),
             self.act,
         )
 

modules/transformer_seq_encoder.py

@@ -1,62 +1,78 @@
 import torch
 from torch import nn
-
+from torch.nn.utils.rnn import pad_sequence
 import PytorchBoot.stereotype as stereotype
 
+
 @stereotype.module("transformer_seq_encoder")
 class TransformerSequenceEncoder(nn.Module):
     def __init__(self, config):
         super(TransformerSequenceEncoder, self).__init__()
         self.config = config
-        embed_dim = config['pts_embed_dim'] + config['pose_embed_dim']
-        self.positional_encoding = nn.Parameter(torch.zeros(1, config['max_seq_len'], embed_dim))
-        encoder_layer = nn.TransformerEncoderLayer(d_model=embed_dim, nhead=config['num_heads'], dim_feedforward=config['ffn_dim'])
-        self.transformer_encoder = nn.TransformerEncoder(encoder_layer, num_layers=config['num_layers'])
-        self.fc = nn.Linear(embed_dim, config['output_dim'])
+        embed_dim = config["pts_embed_dim"] + config["pose_embed_dim"]
+        encoder_layer = nn.TransformerEncoderLayer(
+            d_model=embed_dim,
+            nhead=config["num_heads"],
+            dim_feedforward=config["ffn_dim"],
+            batch_first=True,
+        )
+        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_batch, pose_embedding_list_batch):
-        batch_size = len(pts_embedding_list_batch)
+        # Combine features and pad sequences
         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
+        lengths = []
 
+        for pts_embedding_list, pose_embedding_list in zip(pts_embedding_list_batch, pose_embedding_list_batch):
+            combined_features = [
+                torch.cat((pts_embed, pose_embed), dim=-1)
+                for pts_embed, pose_embed in zip(pts_embedding_list, pose_embedding_list)
+            ]
+            combined_features_batch.append(torch.stack(combined_features))
+            lengths.append(len(combined_features))
+
+        combined_tensor = pad_sequence(combined_features_batch, batch_first=True)  # Shape: [batch_size, max_seq_len, embed_dim]
+
+        # Prepare mask for padding
+        max_len = max(lengths)
+        padding_mask = torch.tensor([([0] * length + [1] * (max_len - length)) for length in lengths], dtype=torch.bool)
         # Transformer encoding
-        transformer_output = self.transformer_encoder(combined_tensor)
-        
+        transformer_output = self.transformer_encoder(combined_tensor, src_key_padding_mask=padding_mask)
+
         # 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, 
+        "pose_embed_dim": 256,
+        "num_heads": 4,  
+        "ffn_dim": 256,
+        "num_layers": 3, 
+        "output_dim": 2048, 
     }
 
     encoder = TransformerSequenceEncoder(config)
-    seq_len = 5
+    seq_len = [5, 8, 9, 4]
     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)
+
+    pts_embedding_list_batch = [
+        torch.randn(seq_len[idx], config["pts_embed_dim"]) for idx in range(batch_size)
+    ]
+    pose_embedding_list_batch = [
+        torch.randn(seq_len[idx], config["pose_embed_dim"]) for idx 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)