End-to-end diarization loss

This is the implement of PIT loss, FastPIT loss and OPTM loss for end-to-end diarization training. Acknowledge @tdedecko for the implement of Hungarian algorithm at https://github.com/tdedecko/hungarian-algorithm.

2021/07/22 update:

Replace Hungurian.py with scipy.optimize.linear_sum_assignment to fix problems.

Prerequisites

Python
Pytorch
numpy
scipy

Simple example:

The loss directory is regarded as the python package. Usage:

import torch
from loss import PITLoss, FastPITLoss, OPTMLoss

B, T, n_class = 1, 4, 4     # batch_size * nframe * n_speaker
device = 'cpu'
pred = torch.rand(B, T, n_class).to(device)
label = torch.randint(0, 2, (B, T, n_class)).float().to(device)

criterion1 = PITLoss()
criterion2 = FastPITLoss()
criterion3 = OPTMLoss()
loss1, assigned_label1 = criterion1(pred, label)
loss2, assigned_label2 = criterion2(pred, label)
loss3, assigned_label3 = criterion3(pred, label)
print(loss1, loss2, loss3)
print(assigned_label1)
print(assigned_label2)
print(assigned_label3)

Citation

Q. Lin et al.: Optimal Mapping Loss: A Faster Loss for End-to-End Speaker Diarization, Odyssey 2020.