Update model.py

model.py

@@ -2,34 +2,42 @@ import torch
 import torch.nn as nn
 
 class DualStreamTransformer(nn.Module):
-    def __init__(self, n_feat1=25, n_feat2=12, d_model=32, num_classes=2):
+    def __init__(self, n_feat1=24, n_feat2=10, d_model=32, num_classes=2):
         super(DualStreamTransformer, self).__init__()
+        
         dim_ff = 128 
+        
         self.feat_tokenizers_1 = nn.ModuleList([nn.Linear(1, d_model) for _ in range(n_feat1)])
         self.cls_token_1 = nn.Parameter(torch.zeros(1, 1, d_model))
-
-        encoder_layer_1 = nn.TransformerEncoderLayer(
-            d_model=d_model, 
-            nhead=4, 
-            dim_feedforward=dim_ff, 
-            batch_first=True
-        )
-        self.encoder_1 = nn.TransformerEncoder(encoder_layer_1, num_layers=2) 
+        encoder_layer_1 = nn.TransformerEncoderLayer(d_model=d_model, nhead=4, dim_feedforward=dim_ff, batch_first=True)
+        self.encoder_1 = nn.TransformerEncoder(encoder_layer_1, num_layers=2)
+        
         self.feat_tokenizers_2 = nn.ModuleList([nn.Linear(1, d_model) for _ in range(n_feat2)])
         self.cls_token_2 = nn.Parameter(torch.zeros(1, 1, d_model))
-        encoder_layer_2 = nn.TransformerEncoderLayer(
-            d_model=d_model, 
-            nhead=4, 
-            dim_feedforward=dim_ff, 
-            batch_first=True
-        )
-        self.encoder_2 = nn.TransformerEncoder(encoder_layer_2, num_layers=2) 
-        self.mlp_head = nn.Sequential(
-            nn.Linear(d_model * 2, d_model),
-            nn.ReLU(),
-            nn.Linear(d_model, d_model)
+        encoder_layer_2 = nn.TransformerEncoderLayer(d_model=d_model, nhead=4, dim_feedforward=dim_ff, batch_first=True)
+        self.encoder_2 = nn.TransformerEncoder(encoder_layer_2, num_layers=2)
+        
+    
+        self.fusion = nn.Sequential(
+            nn.Linear(d_model * 2, d_model), 
+            nn.ReLU()                        
         )
 
+    def forward(self, x1, x2):
+        tokens1 = [layer(x1[:, i].unsqueeze(1)) for i, layer in enumerate(self.feat_tokenizers_1)]
+        x1_emb = torch.stack(tokens1, dim=1)
+        x1_emb = torch.cat((self.cls_token_1.expand(x1.size(0), -1, -1), x1_emb), dim=1)
+        feat1 = self.encoder_1(x1_emb)[:, 0, :]
+        
+        
+        tokens2 = [layer(x2[:, i].unsqueeze(1)) for i, layer in enumerate(self.feat_tokenizers_2)]
+        x2_emb = torch.stack(tokens2, dim=1)
+        x2_emb = torch.cat((self.cls_token_2.expand(x2.size(0), -1, -1), x2_emb), dim=1)
+        feat2 = self.encoder_2(x2_emb)[:, 0, :]
+        
+        combined = torch.cat((feat1, feat2), dim=1)
+        return self.fusion(combined) 
+
 class ArcMarginProduct(nn.Module):
     def __init__(self, in_features, out_features, s=30.0, m=0.5):
         super(ArcMarginProduct, self).__init__()