From b1e03bee38e0d28c0cb420e5cb4e653e4643506c Mon Sep 17 00:00:00 2001 From: ZDisket <30500847+ZDisket@users.noreply.github.com> Date: Sun, 25 Jun 2023 19:18:56 -0300 Subject: [PATCH] Add Bin loss, which increases naturalness --- configs/ljs_li44_tmbert_ft_s1_arpa.json | 6 +++-- configs/ljs_li44_tmbert_ft_small_s1_arpa.json | 2 ++ configs/ljs_li44_tmbert_s1_arpa_mid.json | 2 ++ losses.py | 25 +++++++++++++++++++ models.py | 2 +- train.py | 18 ++++++++++--- 6 files changed, 48 insertions(+), 7 deletions(-) diff --git a/configs/ljs_li44_tmbert_ft_s1_arpa.json b/configs/ljs_li44_tmbert_ft_s1_arpa.json index 757b054..fb72990 100644 --- a/configs/ljs_li44_tmbert_ft_s1_arpa.json +++ b/configs/ljs_li44_tmbert_ft_s1_arpa.json @@ -2,7 +2,7 @@ "train": { "log_interval": 200, "eval_interval": 1000, - "save_interval": 2000, + "save_interval": 5000, "seed": 1234, "epochs": 20000, "learning_rate": 2e-4, @@ -17,7 +17,9 @@ "warmup_epochs": 0, "c_mel": 45, "c_kl": 1.0, - "use_8bit": false + "use_8bit": false, + "bin_loss_warmup_epochs": 25, + "bin_loss_start_epoch": 10 }, "data": { "training_files":"filelists/train_filelist.txt.cleaned", diff --git a/configs/ljs_li44_tmbert_ft_small_s1_arpa.json b/configs/ljs_li44_tmbert_ft_small_s1_arpa.json index fc7d05b..b142c10 100644 --- a/configs/ljs_li44_tmbert_ft_small_s1_arpa.json +++ b/configs/ljs_li44_tmbert_ft_small_s1_arpa.json @@ -18,6 +18,8 @@ "c_mel": 45, "c_kl": 1.0, "use_8bit": true + "bin_loss_warmup_epochs": 25, + "bin_loss_start_epoch": 10 }, "data": { "training_files":"filelists/train_filelist.txt.cleaned", diff --git a/configs/ljs_li44_tmbert_s1_arpa_mid.json b/configs/ljs_li44_tmbert_s1_arpa_mid.json index bb9a2ea..ef8a43d 100644 --- a/configs/ljs_li44_tmbert_s1_arpa_mid.json +++ b/configs/ljs_li44_tmbert_s1_arpa_mid.json @@ -18,6 +18,8 @@ "c_mel": 45, "c_kl": 1.0, "use_8bit": true + "bin_loss_warmup_epochs": 150, + "bin_loss_start_epoch": 100 }, "data": { "training_files":"filelists/train_filelist.txt.cleaned", diff --git a/losses.py b/losses.py index 189bbeb..bf5f065 100644 --- a/losses.py +++ b/losses.py @@ -48,6 +48,31 @@ class ForwardSumLoss(torch.nn.modules.loss._Loss): total_loss /= attn_logprob.shape[0] return total_loss + + +class BinLoss(torch.nn.modules.loss._Loss): + def __init__(self): + super().__init__() + + @property + def input_types(self): + return { + "hard_attention": NeuralType(('B', 'S', 'T_spec', 'T_text'), ProbsType()), + "soft_attention": NeuralType(('B', 'S', 'T_spec', 'T_text'), ProbsType()), + } + + @property + def output_types(self): + return { + "bin_loss": NeuralType(elements_type=LossType()), + } + + + def forward(self, hard_attention, soft_attention): + log_sum = torch.log(torch.clamp(soft_attention[hard_attention == 1], min=1e-12)).sum() + return -log_sum / hard_attention.sum() + + def feature_loss(fmap_r, fmap_g): loss = 0 for dr, dg in zip(fmap_r, fmap_g): diff --git a/models.py b/models.py index 66bdccb..4cd0cb7 100644 --- a/models.py +++ b/models.py @@ -734,7 +734,7 @@ class SynthesizerTrn(nn.Module): spec, phase, x1 = self.dec(z_slice, g=g) o = self.stft.inverse(spec, phase) - return o, x1, l_length, attn, attn_logprob, ids_slice, x_mask, y_mask, (z, z_p, m_p, logs_p, m_q, logs_q) + return o, x1, l_length, attn, attn_logprob, attn_soft, ids_slice, x_mask, y_mask, (z, z_p, m_p, logs_p, m_q, logs_q) def infer(self, x, x_lengths, tm_hidden, bert, bert_lens, sid=None, noise_scale=1, length_scale=1, noise_scale_w=1., max_len=None): tm_encoded = self.tm_enc(tm_hidden) # [b, tm_last_size] diff --git a/train.py b/train.py index 37e76e5..10fd948 100644 --- a/train.py +++ b/train.py @@ -34,7 +34,8 @@ from losses import ( discriminator_loss, feature_loss, kl_loss, - ForwardSumLoss + ForwardSumLoss, + BinLoss ) from mel_processing import mel_spectrogram_torch, spec_to_mel_torch from text.symbols import symbols @@ -177,6 +178,7 @@ def train_and_evaluate(rank, epoch, hps, nets, optims, schedulers, scaler, loade train_loader.batch_sampler.set_epoch(epoch) forward_sum = ForwardSumLoss() + bin_loss = BinLoss() global global_step @@ -205,7 +207,7 @@ def train_and_evaluate(rank, epoch, hps, nets, optims, schedulers, scaler, loade hps.data.mel_fmax) - y_hat, x1, l_length, attn, attn_logprob, ids_slice, x_mask, z_mask,\ + y_hat, x1, l_length, attn, attn_logprob, attn_soft, ids_slice, x_mask, z_mask,\ (z, z_p, m_p, logs_p, m_q, logs_q) = net_g(x, x_lengths, spec, spec_lengths, mel, tm_hidden, bert, bert_lens) @@ -257,6 +259,14 @@ def train_and_evaluate(rank, epoch, hps, nets, optims, schedulers, scaler, loade with autocast(enabled=False): loss_dur = torch.sum(l_length.float()) ctc_loss = forward_sum(attn_logprob=attn_logprob, in_lens=x_lengths, out_lens=spec_lengths) + loss_align = ctc_loss + + if epoch > hps.train.bin_loss_start_epoch: + bin_loss_scale = min((epoch - hps.train.bin_loss_start_epoch) / hps.train.bin_loss_warmup_epochs, 1.0) + al_match_loss = bin_loss(hard_attention=attn, soft_attention=attn_soft) * bin_loss_scale + loss_align += al_match_loss + + loss_mel = F.l1_loss(y_mel, y_hat_mel) * hps.train.c_mel loss_kl = kl_loss(z_p, logs_q, m_p, logs_p, z_mask) * hps.train.c_kl @@ -271,7 +281,7 @@ def train_and_evaluate(rank, epoch, hps, nets, optims, schedulers, scaler, loade loss_gen = loss_gen_f + loss_gen_s losses_gen = losses_gen_f + losses_gen_s - loss_gen_all = loss_gen + loss_fm + loss_mel + loss_dur + loss_kl + ctc_loss + loss_gen_all = loss_gen + loss_fm + loss_mel + loss_dur + loss_kl + loss_align loss_gen_all /= n_acc_steps if ((batch_idx + 1) % n_acc_steps == 0) or (batch_idx + 1 == len(train_loader)): @@ -296,7 +306,7 @@ def train_and_evaluate(rank, epoch, hps, nets, optims, schedulers, scaler, loade logger.info([x.item() for x in losses] + [global_step, lr]) scalar_dict = {"loss/g/total": loss_gen_all, "loss/d/total": loss_disc_all, "learning_rate": lr, "grad_norm_d": grad_norm_d, "grad_norm_g": grad_norm_g} - scalar_dict.update({"loss/g/fm": loss_fm, "loss/g/mel": loss_mel, "loss/g/dur": loss_dur, "loss/g/kl": loss_kl, "loss/g/ctc" : ctc_loss}) + scalar_dict.update({"loss/g/fm": loss_fm, "loss/g/mel": loss_mel, "loss/g/dur": loss_dur, "loss/g/kl": loss_kl, "loss/g/align" : loss_align}) scalar_dict.update({"loss/g/{}".format(i): v for i, v in enumerate(losses_gen)}) scalar_dict.update({"loss/d_r/{}".format(i): v for i, v in enumerate(losses_disc_r)})