76 lines
2.9 KiB
Python
76 lines
2.9 KiB
Python
|
import torch
|
|||
|
|
import torch.nn as nn
|
|||
|
|
|
|||
|
|
|
|||
|
|
class TemporalAveragePooling(nn.Module):
|
|||
|
|
def __init__(self):
|
|||
|
|
"""TAP
|
|||
|
|
Paper: Multi-Task Learning with High-Order Statistics for X-vector based Text-Independent Speaker Verification
|
|||
|
|
Link: https://arxiv.org/pdf/1903.12058.pdf
|
|||
|
|
"""
|
|||
|
|
super(TemporalAveragePooling, self).__init__()
|
|||
|
|
|
|||
|
|
def forward(self, x):
|
|||
|
|
"""Computes Temporal Average Pooling Module
|
|||
|
|
Args:
|
|||
|
|
x (torch.Tensor): Input tensor (#batch, channels, frames).
|
|||
|
|
Returns:
|
|||
|
|
torch.Tensor: Output tensor (#batch, channels)
|
|||
|
|
"""
|
|||
|
|
x = torch.mean(x, dim=2)
|
|||
|
|
return x
|
|||
|
|
|
|||
|
|
|
|||
|
|
class TemporalStatisticsPooling(nn.Module):
|
|||
|
|
def __init__(self):
|
|||
|
|
"""TSP
|
|||
|
|
Paper: X-vectors: Robust DNN Embeddings for Speaker Recognition
|
|||
|
|
Link： http://www.danielpovey.com/files/2018_icassp_xvectors.pdf
|
|||
|
|
"""
|
|||
|
|
super(TemporalStatisticsPooling, self).__init__()
|
|||
|
|
|
|||
|
|
def forward(self, x):
|
|||
|
|
"""Computes Temporal Statistics Pooling Module
|
|||
|
|
Args:
|
|||
|
|
x (torch.Tensor): Input tensor (#batch, channels, frames).
|
|||
|
|
Returns:
|
|||
|
|
torch.Tensor: Output tensor (#batch, channels*2)
|
|||
|
|
"""
|
|||
|
|
mean = torch.mean(x, dim=2)
|
|||
|
|
var = torch.var(x, dim=2)
|
|||
|
|
x = torch.cat((mean, var), dim=1)
|
|||
|
|
return x
|
|||
|
|
|
|||
|
|
|
|||
|
|
class SelfAttentivePooling(nn.Module):
|
|||
|
|
def __init__(self, in_dim, bottleneck_dim=128):
|
|||
|
|
# Use Conv1d with stride == 1 rather than Linear, then we don't need to transpose inputs.
|
|||
|
|
# attention dim = 128
|
|||
|
|
super(SelfAttentivePooling, self).__init__()
|
|||
|
|
self.linear1 = nn.Conv1d(in_dim, bottleneck_dim, kernel_size=1) # equals W and b in the paper
|
|||
|
|
self.linear2 = nn.Conv1d(bottleneck_dim, in_dim, kernel_size=1) # equals V and k in the paper
|
|||
|
|
|
|||
|
|
def forward(self, x):
|
|||
|
|
# DON'T use ReLU here! In experiments, I find ReLU hard to converge.
|
|||
|
|
alpha = torch.tanh(self.linear1(x))
|
|||
|
|
alpha = torch.softmax(self.linear2(alpha), dim=2)
|
|||
|
|
mean = torch.sum(alpha * x, dim=2)
|
|||
|
|
return mean
|
|||
|
|
|
|||
|
|
|
|||
|
|
class AttentiveStatsPool(nn.Module):
|
|||
|
|
def __init__(self, in_dim, bottleneck_dim=128):
|
|||
|
|
super().__init__()
|
|||
|
|
# Use Conv1d with stride == 1 rather than Linear, then we don't need to transpose inputs.
|
|||
|
|
self.linear1 = nn.Conv1d(in_dim, bottleneck_dim, kernel_size=1) # equals W and b in the paper
|
|||
|
|
self.linear2 = nn.Conv1d(bottleneck_dim, in_dim, kernel_size=1) # equals V and k in the paper
|
|||
|
|
|
|||
|
|
def forward(self, x):
|
|||
|
|
# DON'T use ReLU here! In experiments, I find ReLU hard to converge.
|
|||
|
|
alpha = torch.tanh(self.linear1(x))
|
|||
|
|
alpha = torch.softmax(self.linear2(alpha), dim=2)
|
|||
|
|
mean = torch.sum(alpha * x, dim=2)
|
|||
|
|
residuals = torch.sum(alpha * x ** 2, dim=2) - mean ** 2
|
|||
|
|
std = torch.sqrt(residuals.clamp(min=1e-9))
|
|||
|
|
return torch.cat([mean, std], dim=1)
|