"""
Permutation Invariant Training (PIT) loss
Authors:
* Jiatong Shi 2021
"""
from itertools import permutations
import numpy as np
import torch
__all__ = [
"pit_loss",
]
# compute mask to remove the padding positions
def create_length_mask(length, max_len, num_output, device):
batch_size = len(length)
mask = torch.zeros(batch_size, max_len, num_output)
for i in range(batch_size):
mask[i, : length[i], :] = 1
mask = mask.to(device)
return mask
# compute loss for a single permutation
def pit_loss_single_permute(output, label, length):
bce_loss = torch.nn.BCEWithLogitsLoss(reduction="none")
mask = create_length_mask(length, label.size(1), label.size(2), label.device)
loss = bce_loss(output, label)
loss = loss * mask
loss = torch.sum(torch.mean(loss, dim=2), dim=1)
loss = torch.unsqueeze(loss, dim=1)
return loss
[docs]def pit_loss(output, label, length):
"""
The Permutation Invariant Training loss
Args:
output (torch.FloatTensor): prediction in (batch_size, seq_len, num_class)
label (torch.FloatTensor): label in the same shape as :code:`output`
length (torch.LongTensor): the valid length of each instance. :code:`output` and :code:`label`
share the same valid length
Returns:
tuple:
1. loss (torch.FloatTensor)
2. min_idx (int): the id with the minimum loss
3. all the permutation
"""
num_output = label.size(2)
device = label.device
permute_list = [np.array(p) for p in permutations(range(num_output))]
loss_list = []
for p in permute_list:
label_perm = label[:, :, p]
loss_perm = pit_loss_single_permute(output, label_perm, length)
loss_list.append(loss_perm)
loss = torch.cat(loss_list, dim=1)
min_loss, min_idx = torch.min(loss, dim=1)
loss = torch.sum(min_loss) / torch.sum(length.float().to(device))
return loss, min_idx, permute_list
def get_label_perm(label, perm_idx, perm_list):
batch_size = len(perm_idx)
label_list = []
for i in range(batch_size):
label_list.append(label[i, :, perm_list[perm_idx[i]]].data.cpu().numpy())
return torch.from_numpy(np.array(label_list)).float()