Medix-AI/coqui-tts
Medix-AI
/
coqui-tts
mirror of https://github.com/coqui-ai/TTS.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

ssim loss for tacotron models

This commit is contained in:
erogol 2020-10-28 15:24:18 +01:00
parent 9d0ae2bfb4
commit 9cef923d99
4 changed files with 174 additions and 18 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

10
TTS/tts/configs/config.json
View File

@ -69,10 +69,14 @@
// LOSS SETTINGS
"loss_masking": true, // enable / disable loss masking against the sequence padding.
"decoder_loss_alpha": 0.5, // decoder loss weight. If > 0, it is enabled
"postnet_loss_alpha": 0.25, // postnet loss weight. If > 0, it is enabled
"decoder_loss_alpha": 0.5, // original decoder loss weight. If > 0, it is enabled
"postnet_loss_alpha": 0.25, // original postnet loss weight. If > 0, it is enabled
"postnet_diff_spec_alpha": 0.25, // differential spectral loss weight. If > 0, it is enabled
"decoder_diff_spec_alpha": 0.25, // differential spectral loss weight. If > 0, it is enabled
"decoder_ssim_alpha": 0.5, // decoder ssim loss weight. If > 0, it is enabled
"postnet_ssim_alpha": 0.25, // postnet ssim loss weight. If > 0, it is enabled
"ga_alpha": 5.0, // weight for guided attention loss. If > 0, guided attention is enabled.
"diff_spec_alpha": 0.25, // differential spectral loss weight. If > 0, it is enabled
// VALIDATION
"run_eval": true,

100
TTS/tts/layers/losses.py
View File

@ -5,6 +5,7 @@ from torch import nn
from inspect import signature
from torch.nn import functional
from TTS.tts.utils.generic_utils import sequence_mask
from TTS.tts.utils.ssim import ssim
# pylint: disable=abstract-method Method
@ -25,6 +26,10 @@ class L1LossMasked(nn.Module):
class for each corresponding step.
length: A Variable containing a LongTensor of size (batch,)
which contains the length of each data in a batch.
Shapes:
x: B x T X D
target: B x T x D
length: B
Returns:
loss: An average loss value in range [0, 1] masked by the length.
"""
@ -63,6 +68,10 @@ class MSELossMasked(nn.Module):
class for each corresponding step.
length: A Variable containing a LongTensor of size (batch,)
which contains the length of each data in a batch.
Shapes:
x: B x T X D
target: B x T x D
length: B
Returns:
loss: An average loss value in range [0, 1] masked by the length.
"""
@ -87,6 +96,33 @@ class MSELossMasked(nn.Module):
return loss
class SSIMLoss(torch.nn.Module):
"""SSIM loss as explained here https://en.wikipedia.org/wiki/Structural_similarity"""
def __init__(self):
super().__init__()
self.loss_func = ssim
def forward(self, y_hat, y, length=None):
"""
Args:
y_hat (tensor): model prediction values.
y (tensor): target values.
length (tensor): length of each sample in a batch.
Shapes:
y_hat: B x T X D
y: B x T x D
length: B
Returns:
loss: An average loss value in range [0, 1] masked by the length.
"""
if length is not None:
m = sequence_mask(sequence_length=length,
max_len=y.size(1)).unsqueeze(2).float().to(
y_hat.device)
y_hat, y = y_hat * m, y * m
return 1 - self.loss_func(y_hat.unsqueeze(1), y.unsqueeze(1))
class AttentionEntropyLoss(nn.Module):
# pylint: disable=R0201
def forward(self, align):
@ -118,6 +154,10 @@ class BCELossMasked(nn.Module):
class for each corresponding step.
length: A Variable containing a LongTensor of size (batch,)
which contains the length of each data in a batch.
Shapes:
x: B x T
target: B x T
length: B
Returns:
loss: An average loss value in range [0, 1] masked by the length.
"""
@ -142,13 +182,20 @@ class DifferentailSpectralLoss(nn.Module):
super().__init__()
self.loss_func = loss_func
def forward(self, x, target, length):
def forward(self, x, target, length=None):
"""
Shapes:
x: B x T
target: B x T
length: B
Returns:
loss: An average loss value in range [0, 1] masked by the length.
"""
x_diff = x[:, 1:] - x[:, :-1]
target_diff = target[:, 1:] - target[:, :-1]
if len(signature(self.loss_func).parameters) > 2:
return self.loss_func(x_diff, target_diff, length-1)
# if loss masking is not enabled
return self.loss_func(x_diff, target_diff)
if length is None:
return self.loss_func(x_diff, target_diff)
return self.loss_func(x_diff, target_diff, length-1)
class GuidedAttentionLoss(torch.nn.Module):
@ -188,6 +235,7 @@ class GuidedAttentionLoss(torch.nn.Module):
class TacotronLoss(torch.nn.Module):
"""Collection of Tacotron set-up based on provided config."""
def __init__(self, c, stopnet_pos_weight=10, ga_sigma=0.4):
super(TacotronLoss, self).__init__()
self.stopnet_pos_weight = stopnet_pos_weight
@ -195,6 +243,7 @@ class TacotronLoss(torch.nn.Module):
self.diff_spec_alpha = c.diff_spec_alpha
self.decoder_alpha = c.decoder_loss_alpha
self.postnet_alpha = c.postnet_loss_alpha
self.ssim_alpha = c.ssim_alpha
self.config = c
# postnet and decoder loss
@ -205,12 +254,15 @@ class TacotronLoss(torch.nn.Module):
else:
self.criterion = nn.L1Loss() if c.model in ["Tacotron"
] else nn.MSELoss()
# differential spectral loss
if c.diff_spec_alpha > 0:
self.criterion_diff_spec = DifferentailSpectralLoss(loss_func=self.criterion)
# guided attention loss
if c.ga_alpha > 0:
self.criterion_ga = GuidedAttentionLoss(sigma=ga_sigma)
# differential spectral loss
if c.postnet_diff_spec_alpha > 0 or c.decoder_diff_spec_alpha > 0:
self.criterion_diff_spec = DifferentailSpectralLoss(loss_func=self.criterion)
# ssim loss
if c.postnet_ssim_alpha > 0 or c.decoder_ssim_alpha > 0:
self.criterion_ssim = SSIMLoss()
# stopnet loss
# pylint: disable=not-callable
self.criterion_st = BCELossMasked(
@ -221,6 +273,9 @@ class TacotronLoss(torch.nn.Module):
alignments, alignment_lens, alignments_backwards, input_lens):
return_dict = {}
# remove lengths if no masking is applied
if not self.config.loss_masking:
output_lens = None
# decoder and postnet losses
if self.config.loss_masking:
if self.decoder_alpha > 0:
@ -285,11 +340,30 @@ class TacotronLoss(torch.nn.Module):
loss += ga_loss * self.ga_alpha
return_dict['ga_loss'] = ga_loss * self.ga_alpha
# differential spectral loss
if self.config.diff_spec_alpha > 0:
diff_spec_loss = self.criterion_diff_spec(postnet_output, mel_input, output_lens)
loss += diff_spec_loss * self.diff_spec_alpha
return_dict['diff_spec_loss'] = diff_spec_loss
# decoder differential spectral loss
if self.config.decoder_diff_spec_alpha > 0:
decoder_diff_spec_loss = self.criterion_diff_spec(decoder_output, mel_input, output_lens)
loss += decoder_diff_spec_loss * self.decoder_diff_spec_alpha
return_dict['decoder_diff_spec_loss'] = decoder_diff_spec_loss
# postnet differential spectral loss
if self.config.postnet_diff_spec_alpha > 0:
postnet_diff_spec_loss = self.criterion_diff_spec(postnet_output, mel_input, output_lens)
loss += postnet_diff_spec_loss * self.postnet_diff_spec_alpha
return_dict['postnet_diff_spec_loss'] = postnet_diff_spec_loss
# decoder ssim loss
if self.config.decoder_ssim_alpha > 0:
decoder_ssim_loss = self.criterion_ssim(decoder_output, mel_input, output_lens)
loss += decoder_ssim_loss * self.postnet_ssim_alpha
return_dict['decoder_ssim_loss'] = decoder_ssim_loss
# postnet ssim loss
if self.config.postnet_ssim_alpha > 0:
postnet_ssim_loss = self.criterion_ssim(postnet_output, mel_input, output_lens)
loss += postnet_ssim_loss * self.postnet_ssim_alpha
return_dict['postnet_ssim_loss'] = postnet_ssim_loss
return_dict['loss'] = loss
return return_dict

11
TTS/tts/utils/generic_utils.py
View File

@ -178,10 +178,19 @@ def check_config_tts(c):
check_argument('eval_batch_size', c, restricted=True, val_type=int, min_val=1)
check_argument('r', c, restricted=True, val_type=int, min_val=1)
check_argument('gradual_training', c, restricted=False, val_type=list)
check_argument('loss_masking', c, restricted=True, val_type=bool)
check_argument('apex_amp_level', c, restricted=False, val_type=str)
# check_argument('grad_accum', c, restricted=True, val_type=int, min_val=1, max_val=100)
# loss parameters
check_argument('loss_masking', c, restricted=True, val_type=bool)
check_argument('decoder_loss_alpha', c, restricted=True, val_type=float, min_val=0)
check_argument('postnet_loss_alpha', c, restricted=True, val_type=float, min_val=0)
check_argument('postnet_diff_spec_alpha', c, restricted=True, val_type=float, min_val=0)
check_argument('decoder_diff_spec_alpha', c, restricted=True, val_type=float, min_val=0)
check_argument('decoder_ssim_alpha', c, restricted=True, val_type=float, min_val=0)
check_argument('postnet_ssim_alpha', c, restricted=True, val_type=float, min_val=0)
check_argument('ga_alpha', c, restricted=True, val_type=float, min_val=0)
# validation parameters
check_argument('run_eval', c, restricted=True, val_type=bool)
check_argument('test_delay_epochs', c, restricted=True, val_type=int, min_val=0)

71
tests/test_layers.py
View File

@ -2,7 +2,7 @@ import unittest
import torch as T
from TTS.tts.layers.tacotron import Prenet, CBHG, Decoder, Encoder
from TTS.tts.layers.losses import L1LossMasked
from TTS.tts.layers.losses import L1LossMasked, SSIMLoss
from TTS.tts.utils.generic_utils import sequence_mask
# pylint: disable=unused-variable
@ -149,3 +149,72 @@ class L1LossMaskedTests(unittest.TestCase):
(sequence_mask(dummy_length).float() - 1.0) * 100.0).unsqueeze(2)
output = layer(dummy_input + mask, dummy_target, dummy_length)
assert output.item() == 0, "0 vs {}".format(output.item())
class SSIMLossTests(unittest.TestCase):
def test_in_out(self): #pylint: disable=no-self-use
# test input == target
layer = SSIMLoss()
dummy_input = T.ones(4, 8, 128).float()
dummy_target = T.ones(4, 8, 128).float()
dummy_length = (T.ones(4) * 8).long()
output = layer(dummy_input, dummy_target, dummy_length)
assert output.item() == 0.0
# test input != target
dummy_input = T.ones(4, 8, 128).float()
dummy_target = T.zeros(4, 8, 128).float()
dummy_length = (T.ones(4) * 8).long()
output = layer(dummy_input, dummy_target, dummy_length)
assert abs(output.item() - 1.0) < 1e-4 , "1.0 vs {}".format(output.item())
# test if padded values of input makes any difference
dummy_input = T.ones(4, 8, 128).float()
dummy_target = T.zeros(4, 8, 128).float()
dummy_length = (T.arange(5, 9)).long()
mask = (
(sequence_mask(dummy_length).float() - 1.0) * 100.0).unsqueeze(2)
output = layer(dummy_input + mask, dummy_target, dummy_length)
assert abs(output.item() - 1.0) < 1e-4, "1.0 vs {}".format(output.item())
dummy_input = T.rand(4, 8, 128).float()
dummy_target = dummy_input.detach()
dummy_length = (T.arange(5, 9)).long()
mask = (
(sequence_mask(dummy_length).float() - 1.0) * 100.0).unsqueeze(2)
output = layer(dummy_input + mask, dummy_target, dummy_length)
assert output.item() == 0, "0 vs {}".format(output.item())
# seq_len_norm = True
# test input == target
layer = L1LossMasked(seq_len_norm=True)
dummy_input = T.ones(4, 8, 128).float()
dummy_target = T.ones(4, 8, 128).float()
dummy_length = (T.ones(4) * 8).long()
output = layer(dummy_input, dummy_target, dummy_length)
assert output.item() == 0.0
# test input != target
dummy_input = T.ones(4, 8, 128).float()
dummy_target = T.zeros(4, 8, 128).float()
dummy_length = (T.ones(4) * 8).long()
output = layer(dummy_input, dummy_target, dummy_length)
assert output.item() == 1.0, "1.0 vs {}".format(output.item())
# test if padded values of input makes any difference
dummy_input = T.ones(4, 8, 128).float()
dummy_target = T.zeros(4, 8, 128).float()
dummy_length = (T.arange(5, 9)).long()
mask = (
(sequence_mask(dummy_length).float() - 1.0) * 100.0).unsqueeze(2)
output = layer(dummy_input + mask, dummy_target, dummy_length)
assert abs(output.item() - 1.0) < 1e-5, "1.0 vs {}".format(output.item())
dummy_input = T.rand(4, 8, 128).float()
dummy_target = dummy_input.detach()
dummy_length = (T.arange(5, 9)).long()
mask = (
(sequence_mask(dummy_length).float() - 1.0) * 100.0).unsqueeze(2)
output = layer(dummy_input + mask, dummy_target, dummy_length)
assert output.item() == 0, "0 vs {}".format(output.item())