attention moe

.idea/workspace.xml

@@ -6,7 +6,8 @@
   <component name="ChangeListManager">
     <list default="true" id="23565123-73ab-4f40-a9ef-1086e0c9e1ec" name="Changes" comment="">
       <change beforePath="$PROJECT_DIR$/.idea/workspace.xml" beforeDir="false" afterPath="$PROJECT_DIR$/.idea/workspace.xml" afterDir="false" />
-      <change beforePath="$PROJECT_DIR$/src/lcnn_model.py" beforeDir="false" afterPath="$PROJECT_DIR$/src/lcnn_model.py" afterDir="false" />
+      <change beforePath="$PROJECT_DIR$/script.py" beforeDir="false" afterPath="$PROJECT_DIR$/script.py" afterDir="false" />
+      <change beforePath="$PROJECT_DIR$/src/moe_model.py" beforeDir="false" afterPath="$PROJECT_DIR$/src/moe_model.py" afterDir="false" />
     </list>
     <option name="SHOW_DIALOG" value="false" />
     <option name="HIGHLIGHT_CONFLICTS" value="true" />
@@ -136,7 +137,7 @@
       <workItem from="1743062628099" duration="35000" />
       <workItem from="1743063082652" duration="7000" />
       <workItem from="1743092790258" duration="1395000" />
-      <workItem from="1743151940209" duration="31000" />
+      <workItem from="1743151940209" duration="996000" />
     </task>
     <servers />
   </component>

script.py

@@ -14,7 +14,7 @@ from preprocess import preprocess
 
 from src.rawnet_model import RawNet
 from src.lcnn_model import LCNN
-from src.moe_model import UltimateMOE
+from src.moe_model import UltimateMOE, MOE_attention
 
 # os.environ['OMP_NUM_THREADS'] = '1'
 # os.environ['OPENBLAS_NUM_THREADS'] = '1'
@@ -51,22 +51,28 @@ print('Define Model')
 # model_path = './checkpoints/RAWNET_ASVSPOOF_FOR_INTHEWILD_PURDUE.pth'
 # model.load_state_dict(torch.load(model_path, map_location=device))
 
-# LCNN MODEL
-model = LCNN(return_emb=False).to(device)
-# model_path = './checkpoints/LCNN_ASVSPOOF_FOR_INTHEWILD_PURDUE.pth'
-model_path = './checkpoints/LCNN_ALL_DATA.pth'
-model.load_state_dict(torch.load(model_path, map_location=device))
+# # LCNN MODEL
+# model = LCNN(return_emb=False).to(device)
+# # model_path = './checkpoints/LCNN_ASVSPOOF_FOR_INTHEWILD_PURDUE.pth'
+# model_path = './checkpoints/LCNN_ALL_DATA.pth'
+# model.load_state_dict(torch.load(model_path, map_location=device))
+
+# MOE MODEL
+expert_1 = LCNN(return_emb=True).to(device)
+expert_2 = LCNN(return_emb=True).to(device)
+expert_3 = LCNN(return_emb=True).to(device)
+expert_4 = LCNN(return_emb=True).to(device)
+expert_5 = LCNN(return_emb=True).to(device)
+expert_6 = LCNN(return_emb=True).to(device)
 
-# # MOE MODEL
-# expert_1 = LCNN(return_emb=True).to(device)
-# expert_2 = LCNN(return_emb=True).to(device)
-# expert_3 = LCNN(return_emb=True).to(device)
-# expert_4 = LCNN(return_emb=True).to(device)
-#
 # model = UltimateMOE(experts=[expert_1, expert_2, expert_3, expert_4])
-# model = (model).to(device)
 # model_path = './checkpoints/MOE_ULTIMATE.pth'
-# model.load_state_dict(torch.load(model_path, map_location=device))
+
+model = MOE_attention(experts=[expert_1, expert_2, expert_3, expert_4, expert_5, expert_6], device=device, freezing=True)
+model_path = './checkpoints/MOE_ATTENTION.pth'
+
+model = (model).to(device)
+model.load_state_dict(torch.load(model_path, map_location=device))
 
 model.eval()
 

src/moe_model.py

@@ -54,4 +54,100 @@ class UltimateMOE(nn.Module):
 
         score = self.softmax(weighted_logits)
 
-        return score
+        return score
+
+
+
+class MOE_attention(nn.Module):
+    def __init__(self, experts, device, input_dim=128, freezing=False):
+        super(MOE_attention, self).__init__()
+
+        self.device = device
+        self.experts = nn.ModuleList(experts)
+        self.num_experts = len(experts)
+
+        self.proc_emb_1 = nn.Sequential(
+            nn.Linear(128, 128),
+            nn.BatchNorm1d(128),
+            nn.GLU(),
+            nn.Linear(64, 32)
+        )
+
+        self.proc_emb_2 = nn.Sequential(
+            nn.Linear(128, 128),
+            nn.BatchNorm1d(128),
+            nn.GLU(),
+            nn.Linear(64, 32)
+        )
+
+        self.proc_emb_3 = nn.Sequential(
+            nn.Linear(128, 128),
+            nn.BatchNorm1d(128),
+            nn.GLU(),
+            nn.Linear(64, 32)
+        )
+
+        self.proc_emb_4 = nn.Sequential(
+            nn.Linear(128, 128),
+            nn.BatchNorm1d(128),
+            nn.GLU(),
+            nn.Linear(64, 32)
+        )
+
+        self.proc_emb_5 = nn.Sequential(
+            nn.Linear(128, 128),
+            nn.BatchNorm1d(128),
+            nn.GLU(),
+            nn.Linear(64, 32)
+        )
+
+        self.proc_emb_6 = nn.Sequential(
+            nn.Linear(128, 128),
+            nn.BatchNorm1d(128),
+            nn.GLU(),
+            nn.Linear(64, 32)
+        )
+
+        self.MHead_Attn = nn.MultiheadAttention(
+            embed_dim=32, num_heads=4,
+            dropout=0.1, batch_first=True
+        )
+
+        self.query_proj = nn.Linear(input_dim * self.num_experts, 32)
+        self.softmax = nn.Softmax(dim=1)
+
+        if freezing:
+            for expert in self.experts:
+                for param in expert.parameters():
+                    param.requires_grad = False
+
+    def forward(self, x):
+
+        outputs = [expert(x)[0] for expert in self.experts]
+        embeddings = [expert(x)[1] for expert in self.experts]
+
+        proc_emb_1 = self.proc_emb_1(embeddings[0])
+        proc_emb_2 = self.proc_emb_2(embeddings[1])
+        proc_emb_3 = self.proc_emb_3(embeddings[2])
+        proc_emb_4 = self.proc_emb_4(embeddings[3])
+        proc_emb_5 = self.proc_emb_5(embeddings[4])
+        proc_emb_6 = self.proc_emb_6(embeddings[5])
+
+        processed_embs = torch.stack([proc_emb_1, proc_emb_2, proc_emb_3, proc_emb_4, proc_emb_5, proc_emb_6], dim=1)
+
+        query_input = torch.cat(embeddings, dim=1)
+        query = self.query_proj(query_input).unsqueeze(1)
+
+        attn_output, attn_weights = self.MHead_Attn(
+            query=query,
+            key=processed_embs,
+            value=processed_embs
+        )
+
+        gating_weights = attn_weights.mean(dim=1).unsqueeze(1)
+        expert_outputs = torch.stack(outputs, dim=1)
+
+        combined_output = torch.bmm(gating_weights, expert_outputs).squeeze(1)
+        score = self.softmax(combined_output)
+
+        return score