Spaces:
Running
Running
File size: 1,321 Bytes
29f689c |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 |
import torch
from torch import nn
class CTCLoss(nn.Module):
def __init__(self, use_focal_loss=False, zero_infinity=False, **kwargs):
super(CTCLoss, self).__init__()
self.loss_func = nn.CTCLoss(blank=0,
reduction='none',
zero_infinity=zero_infinity)
self.use_focal_loss = use_focal_loss
def forward(self, predicts, batch):
# predicts = predicts['res']
batch_size = predicts.size(0)
label, label_length = batch[1], batch[2]
predicts = predicts.log_softmax(2)
predicts = predicts.permute(1, 0, 2)
preds_lengths = torch.tensor([predicts.size(0)] * batch_size,
dtype=torch.long)
loss = self.loss_func(predicts, label, preds_lengths, label_length)
if self.use_focal_loss:
# Use torch.clamp to limit the range of loss, avoiding overflow in exponential calculation
clamped_loss = torch.clamp(loss, min=-20, max=20)
weight = 1 - torch.exp(-clamped_loss)
weight = torch.square(weight)
# Use torch.where to avoid multiplying by zero weight
loss = torch.where(weight > 0, loss * weight, loss)
loss = loss.mean()
return {'loss': loss}
|