File size: 2,763 Bytes
b21cfae | 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 | """Small causal Nemotron-L15 -> ECAPA-192 predictors; no ASR/teacher weights."""
from __future__ import annotations
import torch
from torch import nn
from torch.nn import functional as F
class SpeakerPredictor(nn.Module):
def __init__(self, architecture="gru", width=128, input_dim=1024, output_dim=192, centered=False):
super().__init__()
self.architecture = architecture
self.width = width
self.centered = centered
self.norm = nn.LayerNorm(input_dim)
if architecture == "mlp":
self.encoder = nn.Sequential(nn.Linear(input_dim * 2, width), nn.GELU(),
nn.Linear(width, width), nn.GELU())
elif architecture == "gru":
self.project = nn.Sequential(nn.Linear(input_dim, width), nn.GELU())
self.encoder = nn.GRU(width, width, batch_first=True)
elif architecture == "tcn":
self.project = nn.Sequential(nn.Linear(input_dim, width), nn.GELU())
self.encoder = nn.ModuleList([
nn.Sequential(nn.Conv1d(width, width, 3, groups=width, dilation=d),
nn.GELU(), nn.Conv1d(width, width, 1), nn.GELU())
for d in (1, 2, 4)])
else:
raise ValueError(architecture)
self.embedding = nn.Linear(width, output_dim)
self.activity = nn.Linear(width, 3) # speech / silence / ambiguous overlap
self.register_buffer("teacher_mean",torch.zeros(output_dim))
self.register_buffer("teacher_radius",torch.ones(()))
def forward(self, features, lengths=None):
x = self.norm(features)
if self.architecture == "mlp":
h = self.encoder(torch.cat([x.mean(1), x.std(1, unbiased=False)], -1))
elif self.architecture == "gru":
h, _ = self.encoder(self.project(x))
h = h[:, -1]
else:
h = self.project(x).transpose(1, 2)
for dilation, layer in zip((1, 2, 4), self.encoder):
h = h + layer(F.pad(h, (2 * dilation, 0)))
h = h[:, :, -1]
direction = F.normalize(self.embedding(h),dim=-1)
embedding = F.normalize(self.teacher_mean+self.teacher_radius*direction,dim=-1) if self.centered else direction
return embedding, self.activity(h)
def load_state(model, state):
"""Compatibility with the first uncentered pilot checkpoints."""
missing, unexpected = model.load_state_dict(state,strict=False)
if unexpected or set(missing)-{"teacher_mean","teacher_radius"}:
raise ValueError(f"Unexpected checkpoint mismatch: missing={missing}, unexpected={unexpected}")
def parameter_count(model):
return sum(p.numel() for p in model.parameters() if p.requires_grad)
|