wehub-resource-sync
57 lines
1.8 KiB
Python
57 lines
1.8 KiB
Python
import torch
|
|
import torch.nn as nn
|
|
from torch import Tensor
|
|
|
|
|
|
class DurationLoss(nn.Module):
|
|
"""
|
|
Loss module as combination of phone duration loss, word duration loss and sentence duration loss.
|
|
"""
|
|
|
|
def __init__(self, offset, loss_type,
|
|
lambda_pdur=0.6, lambda_wdur=0.3, lambda_sdur=0.1):
|
|
super().__init__()
|
|
self.loss_type = loss_type
|
|
if self.loss_type == 'mse':
|
|
self.loss = nn.MSELoss()
|
|
elif self.loss_type == 'huber':
|
|
self.loss = nn.HuberLoss()
|
|
else:
|
|
raise NotImplementedError()
|
|
self.offset = offset
|
|
|
|
self.lambda_pdur = lambda_pdur
|
|
self.lambda_wdur = lambda_wdur
|
|
self.lambda_sdur = lambda_sdur
|
|
|
|
def linear2log(self, any_dur):
|
|
return torch.log(any_dur + self.offset)
|
|
|
|
def forward(self, dur_pred: Tensor, dur_gt: Tensor, ph2word: Tensor) -> Tensor:
|
|
dur_gt = dur_gt.to(dtype=dur_pred.dtype)
|
|
|
|
# pdur_loss
|
|
pdur_loss = self.lambda_pdur * self.loss(self.linear2log(dur_pred), self.linear2log(dur_gt))
|
|
|
|
dur_pred = dur_pred.clamp(min=0.) # clip to avoid NaN loss
|
|
|
|
# wdur loss
|
|
shape = dur_pred.shape[0], ph2word.max() + 1
|
|
wdur_pred = dur_pred.new_zeros(*shape).scatter_add(
|
|
1, ph2word, dur_pred
|
|
)[:, 1:] # [B, T_ph] => [B, T_w]
|
|
wdur_gt = dur_gt.new_zeros(*shape).scatter_add(
|
|
1, ph2word, dur_gt
|
|
)[:, 1:] # [B, T_ph] => [B, T_w]
|
|
wdur_loss = self.lambda_wdur * self.loss(self.linear2log(wdur_pred), self.linear2log(wdur_gt))
|
|
|
|
# sdur loss
|
|
sdur_pred = dur_pred.sum(dim=1)
|
|
sdur_gt = dur_gt.sum(dim=1)
|
|
sdur_loss = self.lambda_sdur * self.loss(self.linear2log(sdur_pred), self.linear2log(sdur_gt))
|
|
|
|
# combine
|
|
dur_loss = pdur_loss + wdur_loss + sdur_loss
|
|
|
|
return dur_loss
|