factorized the architecture

hungarian_net/dataset.py → hungarian_net/lightning_datamodules/hungarian_datamodule.py

@@ -2,7 +2,7 @@
 from lightning import LightningDataModule
 from torch.utils.data import DataLoader, Dataset
 
-from hungarian_net.generate_hnet_training_data import load_obj
+from generate_hnet_training_data import load_obj
 
 
 class HungarianDataset(Dataset):
@@ -101,6 +101,7 @@ def compute_weighted_accuracy(self, n1star, n0star):
 
         return WA
 
+#TODO: factorize HungarianDataset in HungarianDataModule
 
 class HungarianDataModule(LightningDataModule):
     """

hungarian_net/models.py → hungarian_net/lightning_modules/hnet_gru_lightning.py

@@ -9,6 +9,7 @@
 from sklearn.metrics import f1_score
 from torch import optim
 from torchmetrics import MetricCollection
+from hungarian_net.torch_modules.hnet_gru import HNetGRU
 
 
 class HNetGRULightning(L.LightningModule):
@@ -238,77 +239,3 @@ def common_logging(
 
         # F1-score
         self.log(f"{stage}_f1", f1, on_step=False, on_epoch=True, prog_bar=False)
-
-        # Log metrics
-        # predicted_phonemes = self.get_phonemes_from_logits(outputs["logits"])
-        # target_phonemes = batch["phonemes_str"]
-        # metrics_to_log = self.metrics(predicted_phonemes, target_phonemes)
-        # metrics_to_log = {f"{stage}/{k}": v for k, v in metrics_to_log.items()}
-
-        # self.log_dict(dictionary=metrics_to_log, sync_dist=True, prog_bar=True)
-
-
-class AttentionLayer(nn.Module):
-    def __init__(self, in_channels, out_channels, key_channels):
-        super(AttentionLayer, self).__init__()
-        self.conv_Q = nn.Conv1d(in_channels, key_channels, kernel_size=1, bias=False)
-        self.conv_K = nn.Conv1d(in_channels, key_channels, kernel_size=1, bias=False)
-        self.conv_V = nn.Conv1d(in_channels, out_channels, kernel_size=1, bias=False)
-
-    def forward(self, x):
-        Q = self.conv_Q(x)
-        K = self.conv_K(x)
-        V = self.conv_V(x)
-        A = Q.permute(0, 2, 1).matmul(K).softmax(2)
-        x = A.matmul(V.permute(0, 2, 1)).permute(0, 2, 1)
-        return x
-
-    def __repr__(self):
-        return (
-            self._get_name()
-            + "(in_channels={}, out_channels={}, key_channels={})".format(
-                self.conv_Q.in_channels,
-                self.conv_V.out_channels,
-                self.conv_K.out_channels,
-            )
-        )
-
-
-class HNetGRU(nn.Module):
-    def __init__(self, max_len=4, hidden_size=128):
-        super().__init__()
-        self.nb_gru_layers = 1
-        self.max_len = max_len
-        self.gru = nn.GRU(max_len, hidden_size, self.nb_gru_layers, batch_first=True)
-        self.attn = AttentionLayer(hidden_size, hidden_size, hidden_size)
-        self.fc1 = nn.Linear(hidden_size, max_len)
-
-    def forward(self, query):
-        # query - batch x seq x feature
-
-        out, _ = self.gru(query)
-        # out - batch x seq x hidden
-
-        out = out.permute((0, 2, 1))
-        # out - batch x hidden x seq
-
-        out = self.attn.forward(out)
-        # out - batch x hidden x seq
-
-        out = out.permute((0, 2, 1))
-        out = torch.tanh(out)
-        # out - batch x seq x hidden
-
-        out = self.fc1(out)
-        # out - batch x seq x feature
-
-        out1 = out.view(out.shape[0], -1)
-        # out1 - batch x (seq x feature)
-
-        out2, _ = torch.max(out, dim=-1)
-        # out2 - batch x seq x 1
-
-        out3, _ = torch.max(out, dim=-2)
-        # out3 - batch x 1 x feature
-
-        return out1.squeeze(), out2.squeeze(), out3.squeeze()

hungarian_net/torch_modules/attention_layer.py

@@ -0,0 +1,36 @@
+from functools import partial
+from typing import Any, Dict, Tuple
+
+import lightning as L
+import torch
+import torch.nn as nn
+import torchmetrics
+from lightning.pytorch.utilities.types import STEP_OUTPUT
+from sklearn.metrics import f1_score
+from torch import optim
+from torchmetrics import MetricCollection
+
+class AttentionLayer(nn.Module):
+    def __init__(self, in_channels, out_channels, key_channels):
+        super(AttentionLayer, self).__init__()
+        self.conv_Q = nn.Conv1d(in_channels, key_channels, kernel_size=1, bias=False)
+        self.conv_K = nn.Conv1d(in_channels, key_channels, kernel_size=1, bias=False)
+        self.conv_V = nn.Conv1d(in_channels, out_channels, kernel_size=1, bias=False)
+
+    def forward(self, x):
+        Q = self.conv_Q(x)
+        K = self.conv_K(x)
+        V = self.conv_V(x)
+        A = Q.permute(0, 2, 1).matmul(K).softmax(2)
+        x = A.matmul(V.permute(0, 2, 1)).permute(0, 2, 1)
+        return x
+
+    def __repr__(self):
+        return (
+            self._get_name()
+            + "(in_channels={}, out_channels={}, key_channels={})".format(
+                self.conv_Q.in_channels,
+                self.conv_V.out_channels,
+                self.conv_K.out_channels,
+            )
+        )

hungarian_net/torch_modules/hnet_gru.py

@@ -0,0 +1,51 @@
+from functools import partial
+from typing import Any, Dict, Tuple
+
+import lightning as L
+import torch
+import torch.nn as nn
+import torchmetrics
+from lightning.pytorch.utilities.types import STEP_OUTPUT
+from sklearn.metrics import f1_score
+from torch import optim
+from torchmetrics import MetricCollection
+from hungarian_net.torch_modules.attention_layer import AttentionLayer
+
+class HNetGRU(nn.Module):
+    def __init__(self, max_len=4, hidden_size=128):
+        super().__init__()
+        self.nb_gru_layers = 1
+        self.max_len = max_len
+        self.gru = nn.GRU(max_len, hidden_size, self.nb_gru_layers, batch_first=True)
+        self.attn = AttentionLayer(hidden_size, hidden_size, hidden_size)
+        self.fc1 = nn.Linear(hidden_size, max_len)
+
+    def forward(self, query):
+        # query - batch x seq x feature
+
+        out, _ = self.gru(query)
+        # out - batch x seq x hidden
+
+        out = out.permute((0, 2, 1))
+        # out - batch x hidden x seq
+
+        out = self.attn.forward(out)
+        # out - batch x hidden x seq
+
+        out = out.permute((0, 2, 1))
+        out = torch.tanh(out)
+        # out - batch x seq x hidden
+
+        out = self.fc1(out)
+        # out - batch x seq x feature
+
+        out1 = out.view(out.shape[0], -1)
+        # out1 - batch x (seq x feature)
+
+        out2, _ = torch.max(out, dim=-1)
+        # out2 - batch x seq x 1
+
+        out3, _ = torch.max(out, dim=-2)
+        # out3 - batch x 1 x feature
+
+        return out1.squeeze(), out2.squeeze(), out3.squeeze()

hungarian_net/train_hnet.py → run.py

@@ -16,8 +16,8 @@
 from sklearn.metrics import f1_score
 from torch.utils.data import DataLoader
 
-from hungarian_net.dataset import HungarianDataModule, HungarianDataset
-from hungarian_net.models import HNetGRULightning
+from hungarian_net.lightning_datamodules.hungarian_datamodule import HungarianDataModule, HungarianDataset
+from hungarian_net.lightning_modules.hnet_gru_lightning import HNetGRULightning
 
 
 # @hydra.main(

hungarian_net/visualize_hnet_results.py → visualize_hnet_results.py

@@ -3,7 +3,8 @@
 import matplotlib.pyplot as plot
 import torch
 from IPython import embed
-from train_hnet import HNetGRU, HungarianDataset
+from hungarian_net.lightning_datamodules import HungarianDataset
+from hungarian_net.torch_modules import HNetGRU
 
 use_cuda = False
 max_len = 2