fix

train_ms_jp_extra.py

@@ -28,6 +28,7 @@ from mel_processing import mel_spectrogram_torch, spec_to_mel_torch
 from style_bert_vits2.logging import logger
 from style_bert_vits2.models import commons, utils
 from style_bert_vits2.models.hyper_parameters import HyperParameters
+from style_bert_vits2.models.matcha_flow import configure_matcha_only_training
 from style_bert_vits2.models.models_jp_extra import (
     DurationDiscriminator,
     MultiPeriodDiscriminator,
@@ -44,8 +45,7 @@ torch.backends.cudnn.allow_tf32 = (
 )
 torch.set_num_threads(1)
 torch.set_float32_matmul_precision("medium")
-torch.backends.cuda.sdp_kernel("flash")
-torch.backends.cuda.enable_flash_sdp(True)
+torch.backends.cuda.enable_flash_sdp(False)
 torch.backends.cuda.enable_mem_efficient_sdp(
     True
 )  # Not available if torch version is lower than 2.0
@@ -291,7 +291,7 @@ def run():
         logger.info("Using normal MAS for VITS1")
         mas_noise_scale_initial = 0.0
         noise_scale_delta = 0.0
-    if hps.model.use_duration_discriminator is True:
+    if hps.model.use_duration_discriminator is True and not hps.train.matcha_only:
         logger.info("Using duration discriminator for VITS2")
         net_dur_disc = DurationDiscriminator(
             hps.model.hidden_channels,
@@ -302,7 +302,7 @@ def run():
         ).cuda(local_rank)
     else:
         net_dur_disc = None
-    if hps.model.use_wavlm_discriminator is True:
+    if hps.model.use_wavlm_discriminator is True and not hps.train.matcha_only:
         net_wd = WavLMDiscriminator(
             hps.model.slm.hidden, hps.model.slm.nlayers, hps.model.slm.initial_channel
         ).cuda(local_rank)
@@ -335,6 +335,7 @@ def run():
         matcha_dropout=hps.model.matcha_dropout,
         matcha_sigma_min=hps.model.matcha_sigma_min,
         matcha_n_timesteps=hps.model.matcha_n_timesteps,
+        matcha_use_diff_attention=hps.model.matcha_use_diff_attention,
         inter_channels=hps.model.inter_channels,
         hidden_channels=hps.model.hidden_channels,
         filter_channels=hps.model.filter_channels,
@@ -367,6 +368,10 @@ def run():
         for param in net_g.dec.parameters():
             param.requires_grad = False
 
+    if hps.train.matcha_only:
+        trainable = configure_matcha_only_training(net_g)
+        logger.info(f"Matcha-only training enabled ({trainable:,} trainable parameters)")
+
     net_d = MultiPeriodDiscriminator(hps.model.use_spectral_norm).cuda(local_rank)
     optim_g = torch.optim.AdamW(
         filter(lambda p: p.requires_grad, net_g.parameters()),
@@ -430,7 +435,7 @@ def run():
                     ),
                     net_dur_disc,
                     optim_dur_disc,
-                    skip_optimizer=hps.train.skip_optimizer,
+                    skip_optimizer=hps.train.skip_optimizer or hps.train.matcha_only,
                 )
                 if not optim_dur_disc.param_groups[0].get("initial_lr"):
                     optim_dur_disc.param_groups[0]["initial_lr"] = dur_resume_lr
@@ -462,13 +467,13 @@ def run():
                 utils.checkpoints.get_latest_checkpoint_path(model_dir, "G_*.pth"),
                 net_g,
                 optim_g,
-                skip_optimizer=hps.train.skip_optimizer,
+                skip_optimizer=hps.train.skip_optimizer or hps.train.matcha_only,
             )
             _, optim_d, d_resume_lr, epoch_str = utils.checkpoints.load_checkpoint(
                 utils.checkpoints.get_latest_checkpoint_path(model_dir, "D_*.pth"),
                 net_d,
                 optim_d,
-                skip_optimizer=hps.train.skip_optimizer,
+                skip_optimizer=hps.train.skip_optimizer or hps.train.matcha_only,
             )
             if not optim_g.param_groups[0].get("initial_lr"):
                 optim_g.param_groups[0]["initial_lr"] = g_resume_lr
@@ -706,12 +711,21 @@ def train_and_evaluate(
     # train_loader.batch_sampler.set_epoch(epoch)
     global global_step
 
-    net_g.train()
-    net_d.train()
-    if net_dur_disc is not None:
-        net_dur_disc.train()
-    if net_wd is not None:
-        net_wd.train()
+    if hps.train.matcha_only:
+        net_g.eval()
+        net_g.module.matcha.train()
+        net_d.eval()
+        if net_dur_disc is not None:
+            net_dur_disc.eval()
+        if net_wd is not None:
+            net_wd.eval()
+    else:
+        net_g.train()
+        net_d.train()
+        if net_dur_disc is not None:
+            net_dur_disc.train()
+        if net_wd is not None:
+            net_wd.train()
     for batch_idx, (
         x,
         x_lengths,
@@ -725,7 +739,7 @@ def train_and_evaluate(
         bert,
         style_vec,
     ) in enumerate(train_loader):
-        if net_g.module.use_noise_scaled_mas:
+        if net_g.module.use_noise_scaled_mas and not hps.train.matcha_only:
             current_mas_noise_scale = (
                 net_g.module.mas_noise_scale_initial
                 - net_g.module.noise_scale_delta * global_step
@@ -769,6 +783,63 @@ def train_and_evaluate(
                 bert,
                 style_vec,
             )
+            if hps.train.matcha_only:
+                loss_matcha = matcha_loss.float() * hps.train.c_matcha
+
+        if hps.train.matcha_only:
+            optim_g.zero_grad()
+            scaler.scale(loss_matcha).backward()
+            scaler.unscale_(optim_g)
+            grad_norm_g = commons.clip_grad_value_(
+                net_g.module.matcha.parameters(), None
+            )
+            scaler.step(optim_g)
+            scaler.update()
+
+            if rank == 0:
+                if global_step % hps.train.log_interval == 0 and not hps.speedup:
+                    utils.summarize(
+                        writer=writer,
+                        global_step=global_step,
+                        scalars={
+                            "loss/g/total": loss_matcha,
+                            "loss/g/matcha": loss_matcha,
+                            "learning_rate": optim_g.param_groups[0]["lr"],
+                            "grad_norm_g": grad_norm_g,
+                        },
+                    )
+                if (
+                    global_step % hps.train.eval_interval == 0
+                    and global_step != 0
+                    and initial_step != global_step
+                ):
+                    assert hps.model_dir is not None
+                    utils.checkpoints.save_checkpoint(
+                        net_g,
+                        optim_g,
+                        hps.train.learning_rate,
+                        epoch,
+                        os.path.join(hps.model_dir, f"G_{global_step}.pth"),
+                    )
+                    utils.safetensors.save_safetensors(
+                        net_g,
+                        epoch,
+                        os.path.join(
+                            config.out_dir,
+                            f"{config.model_name}_e{epoch}_s{global_step}.safetensors",
+                        ),
+                        for_infer=True,
+                    )
+
+            global_step += 1
+            if pbar is not None:
+                pbar.set_description(
+                    f"Matcha FT {epoch}({100.0 * batch_idx / len(train_loader):.0f}%)/{hps.train.epochs}"
+                )
+                pbar.update()
+            continue
+
+        with autocast(enabled=hps.train.bf16_run, dtype=torch.bfloat16):
             mel = spec_to_mel_torch(
                 spec,
                 hps.data.filter_length,