wehub-resource-sync
98 lines
4.0 KiB
Python
98 lines
4.0 KiB
Python
import torch
|
|
import torchmetrics
|
|
from torch import Tensor
|
|
|
|
from modules.fastspeech.tts_modules import RhythmRegulator
|
|
|
|
|
|
def linguistic_checks(pred, target, ph2word, mask=None):
|
|
if mask is None:
|
|
assert pred.shape == target.shape == ph2word.shape, \
|
|
f'shapes of pred, target and ph2word mismatch: {pred.shape}, {target.shape}, {ph2word.shape}'
|
|
else:
|
|
assert pred.shape == target.shape == ph2word.shape == mask.shape, \
|
|
f'shapes of pred, target and mask mismatch: {pred.shape}, {target.shape}, {ph2word.shape}, {mask.shape}'
|
|
assert pred.ndim == 2, f'all inputs should be 2D, but got {pred.shape}'
|
|
assert torch.any(ph2word > 0), 'empty word sequence'
|
|
assert torch.all(ph2word >= 0), 'unexpected negative word index'
|
|
assert ph2word.max() <= pred.shape[1], f'word index out of range: {ph2word.max()} > {pred.shape[1]}'
|
|
assert torch.all(pred >= 0.), f'unexpected negative ph_dur prediction'
|
|
assert torch.all(target >= 0.), f'unexpected negative ph_dur target'
|
|
|
|
|
|
class RhythmCorrectness(torchmetrics.Metric):
|
|
def __init__(self, *, tolerance, **kwargs):
|
|
super().__init__(**kwargs)
|
|
assert 0. < tolerance < 1., 'tolerance should be within (0, 1)'
|
|
self.tolerance = tolerance
|
|
self.add_state('correct', default=torch.tensor(0, dtype=torch.int), dist_reduce_fx='sum')
|
|
self.add_state('total', default=torch.tensor(0, dtype=torch.int), dist_reduce_fx='sum')
|
|
|
|
def update(self, pdur_pred: Tensor, pdur_target: Tensor, ph2word: Tensor, mask=None) -> None:
|
|
"""
|
|
|
|
:param pdur_pred: predicted ph_dur
|
|
:param pdur_target: reference ph_dur
|
|
:param ph2word: word division sequence
|
|
:param mask: valid or non-padding mask
|
|
"""
|
|
linguistic_checks(pdur_pred, pdur_target, ph2word, mask=mask)
|
|
|
|
shape = pdur_pred.shape[0], ph2word.max() + 1
|
|
wdur_pred = pdur_pred.new_zeros(*shape).scatter_add(
|
|
1, ph2word, pdur_pred
|
|
)[:, 1:] # [B, T_ph] => [B, T_w]
|
|
wdur_target = pdur_target.new_zeros(*shape).scatter_add(
|
|
1, ph2word, pdur_target
|
|
)[:, 1:] # [B, T_ph] => [B, T_w]
|
|
if mask is None:
|
|
wdur_mask = torch.ones_like(wdur_pred, dtype=torch.bool)
|
|
else:
|
|
wdur_mask = mask.new_zeros(*shape).scatter_add(
|
|
1, ph2word, mask
|
|
)[:, 1:].bool() # [B, T_ph] => [B, T_w]
|
|
|
|
correct = torch.abs(wdur_pred - wdur_target) <= wdur_target * self.tolerance
|
|
correct &= wdur_mask
|
|
|
|
self.correct += correct.sum()
|
|
self.total += wdur_mask.sum()
|
|
|
|
def compute(self) -> Tensor:
|
|
return self.correct / self.total
|
|
|
|
|
|
class PhonemeDurationAccuracy(torchmetrics.Metric):
|
|
def __init__(self, *, tolerance, **kwargs):
|
|
super().__init__(**kwargs)
|
|
self.tolerance = tolerance
|
|
self.rr = RhythmRegulator()
|
|
self.add_state('accurate', default=torch.tensor(0, dtype=torch.int), dist_reduce_fx='sum')
|
|
self.add_state('total', default=torch.tensor(0, dtype=torch.int), dist_reduce_fx='sum')
|
|
|
|
def update(self, pdur_pred: Tensor, pdur_target: Tensor, ph2word: Tensor, mask=None) -> None:
|
|
"""
|
|
|
|
:param pdur_pred: predicted ph_dur
|
|
:param pdur_target: reference ph_dur
|
|
:param ph2word: word division sequence
|
|
:param mask: valid or non-padding mask
|
|
"""
|
|
linguistic_checks(pdur_pred, pdur_target, ph2word, mask=mask)
|
|
|
|
shape = pdur_pred.shape[0], ph2word.max() + 1
|
|
wdur_target = pdur_target.new_zeros(*shape).scatter_add(
|
|
1, ph2word, pdur_target
|
|
)[:, 1:] # [B, T_ph] => [B, T_w]
|
|
pdur_align = self.rr(pdur_pred, ph2word=ph2word, word_dur=wdur_target)
|
|
|
|
accurate = torch.abs(pdur_align - pdur_target) <= pdur_target * self.tolerance
|
|
if mask is not None:
|
|
accurate &= mask
|
|
|
|
self.accurate += accurate.sum()
|
|
self.total += pdur_pred.numel() if mask is None else mask.sum()
|
|
|
|
def compute(self) -> Tensor:
|
|
return self.accurate / self.total
|