File size: 1,370 Bytes
f498134 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 |
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
|