Spaces:
Running
Running
File size: 2,354 Bytes
2cba492 |
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 |
# Adapted from: https://github.com/microsoft/UniSpeech/blob/main/downstreams/speaker_verification/models/ecapa_tdnn.py
import torch
import torch.nn as nn
from .convnext import ConvNextBackbone
class AttentiveStatsPool(nn.Module):
def __init__(self, input_channels: int, output_channels: int, attention_channels: int = 128):
super().__init__()
self.attn = nn.Sequential(
nn.Conv1d(input_channels, attention_channels, kernel_size=1),
nn.Tanh(),
nn.Conv1d(attention_channels, input_channels, kernel_size=1),
nn.Softmax(dim=2),
)
self.proj = nn.Linear(input_channels * 2, output_channels)
self.norm = nn.LayerNorm(output_channels)
def forward(self, x):
alpha = self.attn(x)
mean = torch.sum(alpha * x, dim=2)
residuals = torch.sum(alpha * (x**2), dim=2) - mean**2
std = torch.sqrt(residuals.clamp(min=1e-4, max=1e4))
x = torch.cat([mean, std], dim=1)
return self.norm(self.proj(x))
class GlobalEncoder(nn.Module):
def __init__(
self,
input_channels: int,
output_channels: int,
dim: int,
intermediate_dim: int,
num_layers: int,
skip_embed: bool = False,
attention_channels: int = 128,
use_attn_pool: bool = True,
):
super().__init__()
self.backbone = ConvNextBackbone(
input_channels=input_channels,
dim=dim,
intermediate_dim=intermediate_dim,
num_layers=num_layers,
skip_embed=skip_embed,
)
if use_attn_pool:
self.pooling = AttentiveStatsPool(
input_channels=dim, output_channels=output_channels, attention_channels=attention_channels
)
else:
self.pooling = nn.Sequential(
nn.AdaptiveAvgPool1d(1),
nn.Flatten(1),
nn.Linear(dim, output_channels),
nn.LayerNorm(output_channels),
)
self.output_channels = output_channels
@property
def output_dim(self):
return self.output_channels
def forward(self, x):
features = self.backbone(x)
# (B, T, C) -> (B, C, T)
features = features.transpose(1, 2)
return self.pooling(features) # (B, C_out)
|