| import torch | |
| import torch.nn as nn | |
| class SelfAttentionPooling(nn.Module): | |
| """ | |
| Implementation of SelfAttentionPooling | |
| Original Paper: Self-Attention Encoding and Pooling for Speaker Recognition | |
| https://arxiv.org/pdf/2008.01077v1.pdf | |
| """ | |
| def __init__(self, input_dim): | |
| super(SelfAttentionPooling, self).__init__() | |
| self.W = nn.Linear(input_dim, 1) | |
| def forward(self, batch_rep): | |
| """ | |
| input: | |
| batch_rep : size (N, T, H), N: batch size, T: sequence length, H: Hidden dimension | |
| attention_weight: | |
| att_w : size (N, T, 1) | |
| return: | |
| utter_rep: size (N, H) | |
| """ | |
| softmax = nn.functional.softmax | |
| att_w = softmax(self.W(batch_rep).squeeze(-1)).unsqueeze(-1) | |
| utter_rep = torch.sum(batch_rep * att_w, dim=1) | |
| return utter_rep | |
| class Model(nn.Module): | |
| def __init__(self, input_dim, clipping=False, attention_pooling=False, num_judges=5000, **kwargs): | |
| super(Model, self).__init__() | |
| self.mean_net_linear = nn.Linear(input_dim, 1) | |
| self.mean_net_clipping = clipping | |
| self.mean_net_pooling = SelfAttentionPooling(input_dim) if attention_pooling else None | |
| self.bias_net_linear = nn.Linear(input_dim, 1) | |
| self.bias_net_pooling = SelfAttentionPooling(input_dim) if attention_pooling else None | |
| self.judge_embbeding = nn.Embedding(num_embeddings = num_judges, embedding_dim=input_dim) | |
| def forward(self, features, judge_ids=None): | |
| if self.mean_net_pooling is not None: | |
| x = self.mean_net_pooling(features) | |
| segment_score = self.mean_net_linear(x) | |
| else: | |
| x = self.mean_net_linear(features) | |
| segment_score = x.squeeze(-1).mean(dim=-1) | |
| if self.mean_net_clipping: | |
| segment_score = torch.tanh(segment_score) * 2 + 3 | |
| if judge_ids is None: | |
| return segment_score.squeeze(-1) | |
| else: | |
| time = features.shape[1] | |
| judge_features = self.judge_embbeding(judge_ids) | |
| judge_features = torch.stack([judge_features for i in range(time)], dim = 1) | |
| bias_features = features + judge_features | |
| if self.bias_net_pooling is not None: | |
| y = self.bias_net_pooling(bias_features) | |
| bias_score = self.bias_net_linear(y) | |
| else: | |
| y = self.bias_net_linear(bias_features) | |
| bias_score = y.squeeze(-1).mean(dim=-1) | |
| bias_score = bias_score + segment_score | |
| return segment_score.squeeze(-1), bias_score.squeeze(-1) | |