wehub-resource-sync
47 lines
1.1 KiB
Python
47 lines
1.1 KiB
Python
import torch
|
|
from torch.nn import functional as F
|
|
|
|
|
|
class DotScorer(torch.nn.Module):
|
|
def __init__(self):
|
|
"""Initialize DotScorer."""
|
|
super().__init__()
|
|
|
|
def forward(
|
|
self,
|
|
xs_pad: torch.Tensor,
|
|
spk_emb: torch.Tensor,
|
|
):
|
|
# xs_pad: B, T, D
|
|
# spk_emb: B, N, D
|
|
"""Forward pass for training.
|
|
|
|
Args:
|
|
xs_pad: TODO.
|
|
spk_emb: TODO.
|
|
"""
|
|
scores = torch.matmul(xs_pad, spk_emb.transpose(1, 2))
|
|
return scores
|
|
|
|
|
|
class CosScorer(torch.nn.Module):
|
|
def __init__(self):
|
|
"""Initialize CosScorer."""
|
|
super().__init__()
|
|
|
|
def forward(
|
|
self,
|
|
xs_pad: torch.Tensor,
|
|
spk_emb: torch.Tensor,
|
|
):
|
|
# xs_pad: B, T, D
|
|
# spk_emb: B, N, D
|
|
"""Forward pass for training.
|
|
|
|
Args:
|
|
xs_pad: TODO.
|
|
spk_emb: TODO.
|
|
"""
|
|
scores = F.cosine_similarity(xs_pad.unsqueeze(2), spk_emb.unsqueeze(1), dim=-1)
|
|
return scores
|