Spaces:
Sleeping
Sleeping
import torch | |
import torch.nn.functional as F | |
from torch import nn | |
# adapted from https://github.com/cvqluu/GE2E-Loss | |
class GE2ELoss(nn.Module): | |
def __init__(self, init_w=10.0, init_b=-5.0, loss_method="softmax"): | |
""" | |
Implementation of the Generalized End-to-End loss defined in https://arxiv.org/abs/1710.10467 [1] | |
Accepts an input of size (N, M, D) | |
where N is the number of speakers in the batch, | |
M is the number of utterances per speaker, | |
and D is the dimensionality of the embedding vector (e.g. d-vector) | |
Args: | |
- init_w (float): defines the initial value of w in Equation (5) of [1] | |
- init_b (float): definies the initial value of b in Equation (5) of [1] | |
""" | |
super().__init__() | |
# pylint: disable=E1102 | |
self.w = nn.Parameter(torch.tensor(init_w)) | |
# pylint: disable=E1102 | |
self.b = nn.Parameter(torch.tensor(init_b)) | |
self.loss_method = loss_method | |
print(" > Initialized Generalized End-to-End loss") | |
assert self.loss_method in ["softmax", "contrast"] | |
if self.loss_method == "softmax": | |
self.embed_loss = self.embed_loss_softmax | |
if self.loss_method == "contrast": | |
self.embed_loss = self.embed_loss_contrast | |
# pylint: disable=R0201 | |
def calc_new_centroids(self, dvecs, centroids, spkr, utt): | |
""" | |
Calculates the new centroids excluding the reference utterance | |
""" | |
excl = torch.cat((dvecs[spkr, :utt], dvecs[spkr, utt + 1 :])) | |
excl = torch.mean(excl, 0) | |
new_centroids = [] | |
for i, centroid in enumerate(centroids): | |
if i == spkr: | |
new_centroids.append(excl) | |
else: | |
new_centroids.append(centroid) | |
return torch.stack(new_centroids) | |
def calc_cosine_sim(self, dvecs, centroids): | |
""" | |
Make the cosine similarity matrix with dims (N,M,N) | |
""" | |
cos_sim_matrix = [] | |
for spkr_idx, speaker in enumerate(dvecs): | |
cs_row = [] | |
for utt_idx, utterance in enumerate(speaker): | |
new_centroids = self.calc_new_centroids(dvecs, centroids, spkr_idx, utt_idx) | |
# vector based cosine similarity for speed | |
cs_row.append( | |
torch.clamp( | |
torch.mm( | |
utterance.unsqueeze(1).transpose(0, 1), | |
new_centroids.transpose(0, 1), | |
) | |
/ (torch.norm(utterance) * torch.norm(new_centroids, dim=1)), | |
1e-6, | |
) | |
) | |
cs_row = torch.cat(cs_row, dim=0) | |
cos_sim_matrix.append(cs_row) | |
return torch.stack(cos_sim_matrix) | |
# pylint: disable=R0201 | |
def embed_loss_softmax(self, dvecs, cos_sim_matrix): | |
""" | |
Calculates the loss on each embedding $L(e_{ji})$ by taking softmax | |
""" | |
N, M, _ = dvecs.shape | |
L = [] | |
for j in range(N): | |
L_row = [] | |
for i in range(M): | |
L_row.append(-F.log_softmax(cos_sim_matrix[j, i], 0)[j]) | |
L_row = torch.stack(L_row) | |
L.append(L_row) | |
return torch.stack(L) | |
# pylint: disable=R0201 | |
def embed_loss_contrast(self, dvecs, cos_sim_matrix): | |
""" | |
Calculates the loss on each embedding $L(e_{ji})$ by contrast loss with closest centroid | |
""" | |
N, M, _ = dvecs.shape | |
L = [] | |
for j in range(N): | |
L_row = [] | |
for i in range(M): | |
centroids_sigmoids = torch.sigmoid(cos_sim_matrix[j, i]) | |
excl_centroids_sigmoids = torch.cat((centroids_sigmoids[:j], centroids_sigmoids[j + 1 :])) | |
L_row.append(1.0 - torch.sigmoid(cos_sim_matrix[j, i, j]) + torch.max(excl_centroids_sigmoids)) | |
L_row = torch.stack(L_row) | |
L.append(L_row) | |
return torch.stack(L) | |
def forward(self, x, _label=None): | |
""" | |
Calculates the GE2E loss for an input of dimensions (num_speakers, num_utts_per_speaker, dvec_feats) | |
""" | |
assert x.size()[1] >= 2 | |
centroids = torch.mean(x, 1) | |
cos_sim_matrix = self.calc_cosine_sim(x, centroids) | |
torch.clamp(self.w, 1e-6) | |
cos_sim_matrix = self.w * cos_sim_matrix + self.b | |
L = self.embed_loss(x, cos_sim_matrix) | |
return L.mean() | |
# adapted from https://github.com/clovaai/voxceleb_trainer/blob/master/loss/angleproto.py | |
class AngleProtoLoss(nn.Module): | |
""" | |
Implementation of the Angular Prototypical loss defined in https://arxiv.org/abs/2003.11982 | |
Accepts an input of size (N, M, D) | |
where N is the number of speakers in the batch, | |
M is the number of utterances per speaker, | |
and D is the dimensionality of the embedding vector | |
Args: | |
- init_w (float): defines the initial value of w | |
- init_b (float): definies the initial value of b | |
""" | |
def __init__(self, init_w=10.0, init_b=-5.0): | |
super().__init__() | |
# pylint: disable=E1102 | |
self.w = nn.Parameter(torch.tensor(init_w)) | |
# pylint: disable=E1102 | |
self.b = nn.Parameter(torch.tensor(init_b)) | |
self.criterion = torch.nn.CrossEntropyLoss() | |
print(" > Initialized Angular Prototypical loss") | |
def forward(self, x, _label=None): | |
""" | |
Calculates the AngleProto loss for an input of dimensions (num_speakers, num_utts_per_speaker, dvec_feats) | |
""" | |
assert x.size()[1] >= 2 | |
out_anchor = torch.mean(x[:, 1:, :], 1) | |
out_positive = x[:, 0, :] | |
num_speakers = out_anchor.size()[0] | |
cos_sim_matrix = F.cosine_similarity( | |
out_positive.unsqueeze(-1).expand(-1, -1, num_speakers), | |
out_anchor.unsqueeze(-1).expand(-1, -1, num_speakers).transpose(0, 2), | |
) | |
torch.clamp(self.w, 1e-6) | |
cos_sim_matrix = cos_sim_matrix * self.w + self.b | |
label = torch.arange(num_speakers).to(cos_sim_matrix.device) | |
L = self.criterion(cos_sim_matrix, label) | |
return L | |
class SoftmaxLoss(nn.Module): | |
""" | |
Implementation of the Softmax loss as defined in https://arxiv.org/abs/2003.11982 | |
Args: | |
- embedding_dim (float): speaker embedding dim | |
- n_speakers (float): number of speakers | |
""" | |
def __init__(self, embedding_dim, n_speakers): | |
super().__init__() | |
self.criterion = torch.nn.CrossEntropyLoss() | |
self.fc = nn.Linear(embedding_dim, n_speakers) | |
print("Initialised Softmax Loss") | |
def forward(self, x, label=None): | |
# reshape for compatibility | |
x = x.reshape(-1, x.size()[-1]) | |
label = label.reshape(-1) | |
x = self.fc(x) | |
L = self.criterion(x, label) | |
return L | |
def inference(self, embedding): | |
x = self.fc(embedding) | |
activations = torch.nn.functional.softmax(x, dim=1).squeeze(0) | |
class_id = torch.argmax(activations) | |
return class_id | |
class SoftmaxAngleProtoLoss(nn.Module): | |
""" | |
Implementation of the Softmax AnglePrototypical loss as defined in https://arxiv.org/abs/2009.14153 | |
Args: | |
- embedding_dim (float): speaker embedding dim | |
- n_speakers (float): number of speakers | |
- init_w (float): defines the initial value of w | |
- init_b (float): definies the initial value of b | |
""" | |
def __init__(self, embedding_dim, n_speakers, init_w=10.0, init_b=-5.0): | |
super().__init__() | |
self.softmax = SoftmaxLoss(embedding_dim, n_speakers) | |
self.angleproto = AngleProtoLoss(init_w, init_b) | |
print("Initialised SoftmaxAnglePrototypical Loss") | |
def forward(self, x, label=None): | |
""" | |
Calculates the SoftmaxAnglePrototypical loss for an input of dimensions (num_speakers, num_utts_per_speaker, dvec_feats) | |
""" | |
Lp = self.angleproto(x) | |
Ls = self.softmax(x, label) | |
return Ls + Lp | |