Update model.py

model.py

@@ -4,43 +4,32 @@ import torch.nn as nn
 class DualStreamTransformer(nn.Module):
     def __init__(self, n_feat1=25, n_feat2=12, d_model=32, num_classes=2):
         super(DualStreamTransformer, self).__init__()
-        
-        # Stream 1: CCMQ Tokenizer & Encoder
+        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, batch_first=True)
-        self.encoder_1 = nn.TransformerEncoder(encoder_layer_1, num_layers=3)
-        
-        # Stream 2: OSDI Tokenizer & Encoder
+
+        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, batch_first=True)
-        self.encoder_2 = nn.TransformerEncoder(encoder_layer_2, num_layers=3)
-        
-        # Fusion 層
+        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) # 輸出 Embedding 給 ArcMargin
+            nn.Linear(d_model, d_model)
         )
 
-    def forward(self, x1, x2):
-        # Stream 1 推論
-        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, :] # 取 CLS token
-        
-        # Stream 2 推論
-        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, :] # 取 CLS token
-        
-        # 特徵融合
-        combined = torch.cat((feat1, feat2), dim=1)
-        return self.mlp_head(combined)
-
 class ArcMarginProduct(nn.Module):
     def __init__(self, in_features, out_features, s=30.0, m=0.5):
         super(ArcMarginProduct, self).__init__()
@@ -48,6 +37,5 @@ class ArcMarginProduct(nn.Module):
         nn.init.xavier_uniform_(self.weight)
         
     def predict(self, x):
-        # 推論時直接做線性映射或餘弦相似度
         cosine = torch.matmul(nn.functional.normalize(x), nn.functional.normalize(self.weight).t())
         return cosine