mirror of
https://github.com/storytold/storyteller-ml.git
synced 2026-10-15 00:03:06 +00:00
Brandon Thomas
238 lines
9.4 KiB
Python
238 lines
9.4 KiB
Python
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
|