Spaces:
Sleeping
Sleeping
File size: 4,094 Bytes
42a4544 |
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 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 |
#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Wed Sep 5 20:58:34 2018
@author: harry
"""
import torch
import torch.nn as nn
from utils.hparam import hparam as hp
from utils.utils import get_centroids, get_cossim, calc_loss
from utils.kan import KANLinear
class SpeechEmbedder(nn.Module):
def __init__(self):
super(SpeechEmbedder, self).__init__()
self.LSTM_stack = nn.LSTM(hp.data.nmels, hp.model.hidden, num_layers=hp.model.num_layer, batch_first=True)
for name, param in self.LSTM_stack.named_parameters():
if 'bias' in name:
nn.init.constant_(param, 0.0)
elif 'weight' in name:
nn.init.xavier_normal_(param)
self.projection = nn.Linear(hp.model.hidden, hp.model.proj)
def forward(self, x):
x, _ = self.LSTM_stack(x.float()) #(batch, frames, n_mels)
#only use last frame
x = x[:,x.size(1)-1]
x = self.projection(x.float())
x = x / torch.norm(x, dim=1).unsqueeze(1)
return x
class SpeechEmbedderGRU(nn.Module):
def __init__(self):
super(SpeechEmbedderGRU, self).__init__()
self.GRU_stack = nn.GRU(hp.data.nmels, hp.model.hidden, num_layers=hp.model.num_layer, batch_first=True)
for name, param in self.GRU_stack.named_parameters():
if 'bias' in name:
nn.init.constant_(param, 0.0)
elif 'weight' in name:
nn.init.xavier_normal_(param)
self.projection = nn.Linear(hp.model.hidden, hp.model.proj)
def forward(self, x):
x, _ = self.GRU_stack(x.float()) #(batch, frames, n_mels)
#only use last frame
x = x[:,x.size(1)-1]
x = self.projection(x.float())
x = x / torch.norm(x, dim=1).unsqueeze(1)
return x
class SpeechEmbedderKAN(nn.Module):
def __init__(self):
super(SpeechEmbedderKAN, self).__init__()
self.LSTM_stack = nn.LSTM(hp.data.nmels, hp.model.hidden, num_layers=hp.model.num_layer, batch_first=True)
for name, param in self.LSTM_stack.named_parameters():
if 'bias' in name:
nn.init.constant_(param, 0.0)
elif 'weight' in name:
nn.init.xavier_normal_(param)
self.projection = KANLinear(hp.model.hidden, hp.model.proj)
def forward(self, x):
x, _ = self.LSTM_stack(x.float()) #(batch, frames, n_mels)
#only use last frame
x = x[:,x.size(1)-1]
x = self.projection(x.float())
x = x / torch.norm(x, dim=1).unsqueeze(1)
return x
class SpeechEmbedderBidirectional(nn.Module):
def __init__(self):
super(SpeechEmbedderBidirectional, self).__init__()
self.LSTM_stack = nn.LSTM(hp.data.nmels, hp.model.hidden, num_layers=hp.model.num_layer, batch_first=True, bidirectional=True)
for name, param in self.LSTM_stack.named_parameters():
if 'bias' in name:
nn.init.constant_(param, 0.0)
elif 'weight' in name:
nn.init.xavier_normal_(param)
self.projection = nn.Linear(hp.model.hidden, hp.model.proj)
def forward(self, x):
x, _ = self.LSTM_stack(x.float()) #(batch, frames, n_mels)
#only use last frame
x = x[:, :, :hp.model.hidden]
x = x[:,x.size(1)-1]
x = self.projection(x.float())
x = x / torch.norm(x, dim=1).unsqueeze(1)
return x
class GE2ELoss(nn.Module):
def __init__(self, device):
super(GE2ELoss, self).__init__()
self.w = nn.Parameter(torch.tensor(10.0).to(device), requires_grad=True)
self.b = nn.Parameter(torch.tensor(-5.0).to(device), requires_grad=True)
self.device = device
def forward(self, embeddings):
torch.clamp(self.w, 1e-6)
centroids = get_centroids(embeddings)
cossim = get_cossim(embeddings, centroids)
sim_matrix = self.w*cossim.to(self.device) + self.b
loss, _ = calc_loss(sim_matrix)
return loss
|