ok

train_ms.py

@@ -316,6 +316,12 @@ def run():
         use_mel_posterior_encoder=hps.model.use_mel_posterior_encoder,
         use_duration_discriminator=hps.model.use_duration_discriminator,
         use_wavlm_discriminator=hps.model.use_wavlm_discriminator,
+        use_matcha=hps.model.use_matcha,
+        matcha_channels=hps.model.matcha_channels,
+        matcha_num_heads=hps.model.matcha_num_heads,
+        matcha_dropout=hps.model.matcha_dropout,
+        matcha_sigma_min=hps.model.matcha_sigma_min,
+        matcha_n_timesteps=hps.model.matcha_n_timesteps,
         inter_channels=hps.model.inter_channels,
         hidden_channels=hps.model.hidden_channels,
         filter_channels=hps.model.filter_channels,
@@ -658,6 +664,7 @@ def train_and_evaluate(
             (
                 y_hat,
                 l_length,
+                matcha_loss,
                 attn,
                 ids_slice,
                 x_mask,
@@ -745,10 +752,18 @@ def train_and_evaluate(
                 loss_dur = torch.sum(l_length.float())
                 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
+                loss_matcha = matcha_loss.float() * hps.train.c_matcha
 
                 loss_fm = feature_loss(fmap_r, fmap_g)
                 loss_gen, losses_gen = generator_loss(y_d_hat_g)
-                loss_gen_all = loss_gen + loss_fm + loss_mel + loss_dur + loss_kl
+                loss_gen_all = (
+                    loss_gen
+                    + loss_fm
+                    + loss_mel
+                    + loss_dur
+                    + loss_kl
+                    + loss_matcha
+                )
                 if net_dur_disc is not None:
                     loss_dur_gen, losses_dur_gen = generator_loss(y_dur_hat_g)
                     loss_gen_all += loss_dur_gen
@@ -764,7 +779,15 @@ def train_and_evaluate(
         if rank == 0:
             if global_step % hps.train.log_interval == 0 and not hps.speedup:
                 lr = optim_g.param_groups[0]["lr"]
-                losses = [loss_disc, loss_gen, loss_fm, loss_mel, loss_dur, loss_kl]
+                losses = [
+                    loss_disc,
+                    loss_gen,
+                    loss_fm,
+                    loss_mel,
+                    loss_dur,
+                    loss_kl,
+                    loss_matcha,
+                ]
                 # logger.info(
                 #     "Train Epoch: {} [{:.0f}%]".format(
                 #         epoch, 100.0 * batch_idx / len(train_loader)
@@ -774,6 +797,7 @@ def train_and_evaluate(
 
                 scalar_dict = {
                     "loss/g/total": loss_gen_all,
+                    "loss/g/matcha": loss_matcha,
                     "loss/d/total": loss_disc_all,
                     "learning_rate": lr,
                     "grad_norm_d": grad_norm_d,