import torch import torch.nn as nn from .fagal import create_guided, get_pivot_points from torch.nn import functional as F import numpy as np from numba import jit class ForwardSumLoss(torch.nn.modules.loss._Loss): def __init__(self, blank_logprob=-1): super().__init__() self.log_softmax = torch.nn.LogSoftmax(dim=3) self.ctc_loss = torch.nn.CTCLoss(zero_infinity=True) self.blank_logprob = blank_logprob @property def input_types(self): return { "attn_logprob": NeuralType(('B', 'S', 'T_spec', 'T_text'), LogprobsType()), "in_lens": NeuralType(tuple('B'), LengthsType()), "out_lens": NeuralType(tuple('B'), LengthsType()), } @property def output_types(self): return { "forward_sum_loss": NeuralType(elements_type=LossType()), } def forward(self, attn_logprob, in_lens, out_lens): key_lens = in_lens query_lens = out_lens attn_logprob_padded = F.pad(input=attn_logprob, pad=(1, 0), value=self.blank_logprob) total_loss = 0.0 for bid in range(attn_logprob.shape[0]): target_seq = torch.arange(1, key_lens[bid] + 1).unsqueeze(0) curr_logprob = attn_logprob_padded[bid].permute(1, 0, 2)[: query_lens[bid], :, : key_lens[bid] + 1] curr_logprob = self.log_softmax(curr_logprob[None])[0] loss = self.ctc_loss( curr_logprob, target_seq, input_lengths=query_lens[bid : bid + 1], target_lengths=key_lens[bid : bid + 1], ) total_loss += loss total_loss /= attn_logprob.shape[0] return total_loss class Tacotron2Loss(nn.Module): """ Tacotron2 Loss """ def __init__(self, preprocess_config, model_config, train_config): super(Tacotron2Loss, self).__init__() self.n_frames_per_step = model_config["decoder"]["n_frames_per_step"] self.use_guided_attn_loss = train_config["optimizer"]["guided_attn"] self.mse_loss = nn.MSELoss() self.bce_loss = nn.BCEWithLogitsLoss() self.forward_sum = ForwardSumLoss() if self.use_guided_attn_loss: self.guided_attn_loss = GuidedAttentionLoss( sigma=train_config["optimizer"]["guided_sigma"], alpha=train_config["optimizer"]["guided_lambda"], ) def forward(self, inputs, predictions): mel_target, input_lengths, output_lengths, r_len_pad, gate_target \ = inputs[6], inputs[4], inputs[7], inputs[9], inputs[10] mel_out, mel_out_postnet, gate_out, alignments, attn_logprob, attn_hard = predictions attn_hard = attn_hard.squeeze() mel_target.requires_grad = False gate_target.requires_grad = False gate_target = gate_target.view(-1, 1) gate_out = gate_out.view(-1, 1) mel_loss = self.mse_loss(mel_out, mel_target) + \ self.mse_loss(mel_out_postnet, mel_target) gate_loss = self.bce_loss(gate_out, gate_target) al_forward_sum = self.forward_sum(attn_logprob=attn_logprob, in_lens=input_lengths, out_lens=output_lengths) if self.use_guided_attn_loss: #attn_loss = self.guided_attn_loss(alignments, attn_hard.squeeze(), input_lengths, \ # (output_lengths + r_len_pad)//self.n_frames_per_step) attn_loss = self.mse_loss(alignments, attn_hard) total_loss = mel_loss + gate_loss + attn_loss + al_forward_sum return total_loss, mel_loss, gate_loss, attn_loss, al_forward_sum else: total_loss = mel_loss + gate_loss + al_forward_sum return total_loss, mel_loss, gate_loss, al_forward_sum, torch.tensor([0.], device=mel_target.device) class GuidedAttentionLoss(nn.Module): """Guided attention loss function module. See https://github.com/espnet/espnet/blob/e962a3c609ad535cd7fb9649f9f9e9e0a2a27291/espnet/nets/pytorch_backend/e2e_tts_tacotron2.py#L25 This module calculates the guided attention loss described in `Efficiently Trainable Text-to-Speech System Based on Deep Convolutional Networks with Guided Attention`_, which forces the attention to be diagonal. .. _`Efficiently Trainable Text-to-Speech System Based on Deep Convolutional Networks with Guided Attention`: https://arxiv.org/abs/1710.08969 """ def __init__(self, sigma=0.4, alpha=1.0, reset_always=True): """Initialize guided attention loss module. Args: sigma (float, optional): Standard deviation to control how close attention to a diagonal. alpha (float, optional): Scaling coefficient (lambda). reset_always (bool, optional): Whether to always reset masks. """ super(GuidedAttentionLoss, self).__init__() self.sigma = sigma self.alpha = alpha self.reset_always = reset_always self.guided_attn_masks = None self.masks = None self.vec_pivoter = np.vectorize(get_pivot_points) def _reset_masks(self): self.guided_attn_masks = None self.masks = None def forward(self, att_ws,hard_atts, ilens, olens): """Calculate forward propagation. Args: att_ws (Tensor): Batch of attention weights (B, T_max_out, T_max_in). hard_atts (Tensor): Batch of hard attentions from sep (B, T_max_out, T_max_in) ilens (LongTensor): Batch of input lenghts (B,). olens (LongTensor): Batch of output lenghts (B,). Returns: Tensor: Guided attention loss value. """ if self.guided_attn_masks is None: self.guided_attn_masks = self._make_guided_attention_masks(ilens, olens, hard_atts.cpu().numpy()).to( att_ws.device ) if self.masks is None: self.masks = self._make_masks(ilens, olens).to(att_ws.device) losses = self.guided_attn_masks * att_ws loss = torch.mean(losses.masked_select(self.masks)) if self.reset_always: self._reset_masks() return self.alpha * loss def _make_guided_attention_masks(self, ilens, olens, hard_atts): n_batches = len(ilens) max_ilen = max(ilens) max_olen = max(olens) guided_attn_masks = torch.zeros((n_batches, max_olen, max_ilen)) for idx, (ilen, olen) in enumerate(zip(ilens, olens)): guided_attn_masks[idx, :olen, :ilen] = self._make_guided_attention_mask( ilen, olen, hard_atts[idx] ) return guided_attn_masks def _make_guided_attention_mask(self,ilen, olen, hard_att): """Make guided attention mask. hard Examples: """ pivots = np.argwhere(hard_att > 0.8) guided_att = torch.from_numpy(create_guided(hard_att,pivots,2.5)) return guided_att[:olen,:ilen] def _make_masks(self, ilens, olens): """Make masks indicating non-padded part. Args: ilens (LongTensor or List): Batch of lengths (B,). olens (LongTensor or List): Batch of lengths (B,). Returns: Tensor: Mask tensor indicating non-padded part. dtype=torch.uint8 in PyTorch 1.2- dtype=torch.bool in PyTorch 1.2+ (including 1.2) Examples: >>> ilens, olens = [5, 2], [8, 5] >>> _make_mask(ilens, olens) tensor([[[1, 1, 1, 1, 1], [1, 1, 1, 1, 1], [1, 1, 1, 1, 1], [1, 1, 1, 1, 1], [1, 1, 1, 1, 1], [1, 1, 1, 1, 1], [1, 1, 1, 1, 1], [1, 1, 1, 1, 1]], [[1, 1, 0, 0, 0], [1, 1, 0, 0, 0], [1, 1, 0, 0, 0], [1, 1, 0, 0, 0], [1, 1, 0, 0, 0], [0, 0, 0, 0, 0], [0, 0, 0, 0, 0], [0, 0, 0, 0, 0]]], dtype=torch.uint8) """ in_masks = self.make_non_pad_mask(ilens) # (B, T_in) out_masks = self.make_non_pad_mask(olens) # (B, T_out) return out_masks.unsqueeze(-1) & in_masks.unsqueeze(-2) # (B, T_out, T_in) def make_non_pad_mask(self, lengths, xs=None, length_dim=-1): return ~self.make_pad_mask(lengths, xs, length_dim) def make_pad_mask(self, lengths, xs=None, length_dim=-1): if length_dim == 0: raise ValueError("length_dim cannot be 0: {}".format(length_dim)) if not isinstance(lengths, list): lengths = lengths.tolist() bs = int(len(lengths)) if xs is None: maxlen = int(max(lengths)) else: maxlen = xs.size(length_dim) seq_range = torch.arange(0, maxlen, dtype=torch.int64) seq_range_expand = seq_range.unsqueeze(0).expand(bs, maxlen) seq_length_expand = seq_range_expand.new(lengths).unsqueeze(-1) mask = seq_range_expand >= seq_length_expand if xs is not None: assert xs.size(0) == bs, (xs.size(0), bs) if length_dim < 0: length_dim = xs.dim() + length_dim # ind = (:, None, ..., None, :, , None, ..., None) ind = tuple( slice(None) if i in (0, length_dim) else None for i in range(xs.dim()) ) mask = mask[ind].expand_as(xs).to(xs.device) return mask