Upload lightning_module.py

lightning_module.py

@@ -0,0 +1,41 @@
+import pytorch_lightning as pl
+import torch
+import torch.nn as nn
+import os
+import numpy as np
+import hydra
+from model import load_ssl_model, PhonemeEncoder, DomainEmbedding, LDConditioner, Projection
+
+
+class BaselineLightningModule(pl.LightningModule):
+    def __init__(self, cfg):
+        super().__init__()
+        self.cfg = cfg
+        self.construct_model()
+        self.save_hyperparameters()
+    
+    def construct_model(self):
+        self.feature_extractors = nn.ModuleList([
+            load_ssl_model(cp_path='wav2vec_small.pt'),
+            DomainEmbedding(3,128),
+        ])
+        output_dim = sum([ feature_extractor.get_output_dim() for feature_extractor in self.feature_extractors])
+        output_layers = [
+            LDConditioner(judge_dim=128,num_judges=3000,input_dim=output_dim)
+        ]
+        output_dim = output_layers[-1].get_output_dim()
+        output_layers.append(
+            Projection(hidden_dim=2048,activation=torch.nn.ReLU(),range_clipping=False,input_dim=output_dim)
+
+        )
+
+        self.output_layers = nn.ModuleList(output_layers)
+
+    def forward(self, inputs):
+        outputs = {}
+        for feature_extractor in self.feature_extractors:
+            outputs.update(feature_extractor(inputs))
+        x = outputs
+        for output_layer in self.output_layers:
+            x = output_layer(x,inputs)
+        return x