mirror of
https://github.com/modelscope/FunASR
synced 2025-09-15 14:48:36 +08:00
41 lines
1.4 KiB
Python
41 lines
1.4 KiB
Python
import numpy as np
|
|
import torch
|
|
import torch.nn.functional as F
|
|
from scipy.optimize import linear_sum_assignment
|
|
|
|
|
|
def standard_loss(ys, ts):
|
|
losses = [F.binary_cross_entropy(torch.sigmoid(y), t) * len(y) for y, t in zip(ys, ts)]
|
|
loss = torch.sum(torch.stack(losses))
|
|
n_frames = torch.from_numpy(np.array(np.sum([t.shape[0] for t in ts]))).to(torch.float32).to(ys[0].device)
|
|
loss = loss / n_frames
|
|
return loss
|
|
|
|
|
|
def fast_batch_pit_n_speaker_loss(ys, ts):
|
|
with torch.no_grad():
|
|
bs = len(ys)
|
|
indices = []
|
|
for b in range(bs):
|
|
y = ys[b].transpose(0, 1)
|
|
t = ts[b].transpose(0, 1)
|
|
C, _ = t.shape
|
|
y = y[:, None, :].repeat(1, C, 1)
|
|
t = t[None, :, :].repeat(C, 1, 1)
|
|
bce_loss = F.binary_cross_entropy(torch.sigmoid(y), t, reduction="none").mean(-1)
|
|
C = bce_loss.cpu()
|
|
indices.append(linear_sum_assignment(C))
|
|
labels_perm = [t[:, idx[1]] for t, idx in zip(ts, indices)]
|
|
|
|
return labels_perm
|
|
|
|
|
|
def cal_power_loss(logits, power_ts):
|
|
losses = [F.cross_entropy(input=logit, target=power_t.to(torch.long)) * len(logit) for logit, power_t in
|
|
zip(logits, power_ts)]
|
|
loss = torch.sum(torch.stack(losses))
|
|
n_frames = torch.from_numpy(np.array(np.sum([power_t.shape[0] for power_t in power_ts]))).to(torch.float32).to(
|
|
power_ts[0].device)
|
|
loss = loss / n_frames
|
|
return loss
|