feat: run through the code and fix bugs [PDS-512]

code/TFC/augmentations.py

@@ -58,13 +58,15 @@ def DataTransform_FD(sample, config):
     return aug_F
 
 def remove_frequency(x, pertub_ratio=0.0):
-    mask = torch.cuda.FloatTensor(x.shape).uniform_() > pertub_ratio # maskout_ratio are False
+    # mask = torch.cuda.FloatTensor(x.shape).uniform_() > pertub_ratio # maskout_ratio are False
+    mask = torch.FloatTensor(x.shape).uniform_() > pertub_ratio # maskout_ratio are False
     mask = mask.to(x.device)
     return x*mask
 
 def add_frequency(x, pertub_ratio=0.0):
 
-    mask = torch.cuda.FloatTensor(x.shape).uniform_() > (1-pertub_ratio) # only pertub_ratio of all values are True
+    # mask = torch.cuda.FloatTensor(x.shape).uniform_() > (1-pertub_ratio) # only pertub_ratio of all values are True
+    mask = torch.FloatTensor(x.shape).uniform_() > (1-pertub_ratio) # only pertub_ratio of all values are True
     mask = mask.to(x.device)
     max_amplitude = x.max()
     random_am = torch.rand(mask.shape)*(max_amplitude*0.1)

code/TFC/model.py

@@ -1,60 +1,60 @@
-from torch import nn
-import torch
-from torch.nn import TransformerEncoder, TransformerEncoderLayer
-
-"""Two contrastive encoders"""
-class TFC(nn.Module):
-    def __init__(self, configs):
-        super(TFC, self).__init__()
-
-        encoder_layers_t = TransformerEncoderLayer(configs.TSlength_aligned, dim_feedforward=2*configs.TSlength_aligned, nhead=2, )
-        self.transformer_encoder_t = TransformerEncoder(encoder_layers_t, 2)
-
-        self.projector_t = nn.Sequential(
-            nn.Linear(configs.TSlength_aligned, 256),
-            nn.BatchNorm1d(256),
-            nn.ReLU(),
-            nn.Linear(256, 128)
-        )
-
-        encoder_layers_f = TransformerEncoderLayer(configs.TSlength_aligned, dim_feedforward=2*configs.TSlength_aligned,nhead=2,)
-        self.transformer_encoder_f = TransformerEncoder(encoder_layers_f, 2)
-
-        self.projector_f = nn.Sequential(
-            nn.Linear(configs.TSlength_aligned, 256),
-            nn.BatchNorm1d(256),
-            nn.ReLU(),
-            nn.Linear(256, 128)
-        )
-
-
-    def forward(self, x_in_t, x_in_f):
-        """Use Transformer"""
-        x = self.transformer_encoder_t(x_in_t)
-        h_time = x.reshape(x.shape[0], -1)
-
-        """Cross-space projector"""
-        z_time = self.projector_t(h_time)
-
-        """Frequency-based contrastive encoder"""
-        f = self.transformer_encoder_f(x_in_f)
-        h_freq = f.reshape(f.shape[0], -1)
-
-        """Cross-space projector"""
-        z_freq = self.projector_f(h_freq)
-
-        return h_time, z_time, h_freq, z_freq
-
-
-"""Downstream classifier only used in finetuning"""
-class target_classifier(nn.Module):
-    def __init__(self, configs):
-        super(target_classifier, self).__init__()
-        self.logits = nn.Linear(2*128, 64)
-        self.logits_simple = nn.Linear(64, configs.num_classes_target)
-
-    def forward(self, emb):
-        emb_flat = emb.reshape(emb.shape[0], -1)
-        emb = torch.sigmoid(self.logits(emb_flat))
-        pred = self.logits_simple(emb)
-        return pred
+from torch import nn
+import torch
+from torch.nn import TransformerEncoder, TransformerEncoderLayer
+
+"""Two contrastive encoders"""
+class TFC(nn.Module):
+    def __init__(self, configs):
+        super(TFC, self).__init__()
+
+        encoder_layers_t = TransformerEncoderLayer(configs.TSlength_aligned, dim_feedforward=2*configs.TSlength_aligned, nhead=configs.transformer_nhead, )
+        self.transformer_encoder_t = TransformerEncoder(encoder_layers_t, configs.transformer_num_layers)
+
+        self.projector_t = nn.Sequential(
+            nn.Linear(configs.TSlength_aligned*configs.input_channels, configs.embedding_len*2),
+            nn.BatchNorm1d(configs.embedding_len*2),
+            nn.ReLU(),
+            nn.Linear(configs.embedding_len*2, configs.embedding_len)
+        )
+
+        encoder_layers_f = TransformerEncoderLayer(configs.TSlength_aligned, dim_feedforward=2*configs.TSlength_aligned,nhead=configs.transformer_nhead,)
+        self.transformer_encoder_f = TransformerEncoder(encoder_layers_f, configs.transformer_num_layers)
+
+        self.projector_f = nn.Sequential(
+            nn.Linear(configs.TSlength_aligned*configs.input_channels, configs.embedding_len*2),
+            nn.BatchNorm1d(configs.embedding_len*2),
+            nn.ReLU(),
+            nn.Linear(configs.embedding_len*2, configs.embedding_len)
+        )
+
+
+    def forward(self, x_in_t, x_in_f):
+        """Use Transformer"""
+        x = self.transformer_encoder_t(x_in_t)
+        h_time = x.reshape(x.shape[0], -1)
+
+        """Cross-space projector"""
+        z_time = self.projector_t(h_time)
+
+        """Frequency-based contrastive encoder"""
+        f = self.transformer_encoder_f(x_in_f)
+        h_freq = f.reshape(f.shape[0], -1)
+
+        """Cross-space projector"""
+        z_freq = self.projector_f(h_freq)
+
+        return h_time, z_time, h_freq, z_freq
+
+
+"""Downstream classifier only used in finetuning"""
+class target_classifier(nn.Module):
+    def __init__(self, configs):
+        super(target_classifier, self).__init__()
+        self.logits = nn.Linear(2*128, 64)
+        self.logits_simple = nn.Linear(64, configs.num_classes_target)
+
+    def forward(self, emb):
+        emb_flat = emb.reshape(emb.shape[0], -1)
+        emb = torch.sigmoid(self.logits(emb_flat))
+        pred = self.logits_simple(emb)
+        return pred