From 9cef923d991f57b697f749295dc242e5c745ed1e Mon Sep 17 00:00:00 2001 From: erogol Date: Wed, 28 Oct 2020 15:24:18 +0100 Subject: [PATCH] ssim loss for tacotron models --- TTS/tts/configs/config.json | 10 +++- TTS/tts/layers/losses.py | 100 ++++++++++++++++++++++++++++----- TTS/tts/utils/generic_utils.py | 11 +++- tests/test_layers.py | 71 ++++++++++++++++++++++- 4 files changed, 174 insertions(+), 18 deletions(-) diff --git a/TTS/tts/configs/config.json b/TTS/tts/configs/config.json index 1b63b037..4d3e2674 100644 --- a/TTS/tts/configs/config.json +++ b/TTS/tts/configs/config.json @@ -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, diff --git a/TTS/tts/layers/losses.py b/TTS/tts/layers/losses.py index 67503a76..10ee3905 100644 --- a/TTS/tts/layers/losses.py +++ b/TTS/tts/layers/losses.py @@ -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 diff --git a/TTS/tts/utils/generic_utils.py b/TTS/tts/utils/generic_utils.py index 2361fa85..2c82611f 100644 --- a/TTS/tts/utils/generic_utils.py +++ b/TTS/tts/utils/generic_utils.py @@ -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) diff --git a/tests/test_layers.py b/tests/test_layers.py index 57be51e5..5426e195 100644 --- a/tests/test_layers.py +++ b/tests/test_layers.py @@ -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()) +