Spaces:
Runtime error
Runtime error
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 | |