| 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 | |