Add Bin loss, which increases naturalness

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",

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",

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",

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):

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]

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)})