From 18333e68313907a0bd3da0a7cf9098a703ec0803 Mon Sep 17 00:00:00 2001 From: hofee Date: Fri, 20 Sep 2024 06:43:19 +0000 Subject: [PATCH] change dataset.py to nbv_dataset.py --- core/{dataset.py => nbv_dataset.py} | 0 modules/transformer_seq_encoder.py | 8 +------- 2 files changed, 1 insertion(+), 7 deletions(-) rename core/{dataset.py => nbv_dataset.py} (100%) diff --git a/core/dataset.py b/core/nbv_dataset.py similarity index 100% rename from core/dataset.py rename to core/nbv_dataset.py diff --git a/modules/transformer_seq_encoder.py b/modules/transformer_seq_encoder.py index 8b22b4f..1eae505 100644 --- a/modules/transformer_seq_encoder.py +++ b/modules/transformer_seq_encoder.py @@ -22,7 +22,6 @@ class TransformerSequenceEncoder(nn.Module): self.fc = nn.Linear(embed_dim, config["output_dim"]) def encode_sequence(self, pts_embedding_list_batch, pose_embedding_list_batch): - # Combine features and pad sequences combined_features_batch = [] lengths = [] @@ -36,16 +35,11 @@ class TransformerSequenceEncoder(nn.Module): 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).to(combined_tensor.device) - # Transformer encoding + 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