update

examples/mpnet_aishell/run.sh

@@ -109,7 +109,7 @@ if [ ${stage} -le 2 ] && [ ${stop_stage} -ge 2 ]; then
   $verbose && echo "stage 2: train model"
   cd "${work_dir}" || exit 1
   python3 step_2_train_model.py \
-  --train_dataset "${train_dataset}" \
+  --train_dataset "${valid_dataset}" \
   --valid_dataset "${valid_dataset}" \
   --serialization_dir "${file_dir}" \
   --config_file "${config_file}" \

examples/mpnet_aishell/step_2_train_model.py

@@ -166,6 +166,19 @@ def main():
     generator = MPNetPretrainedModel(config).to(device)
     discriminator = MetricDiscriminatorPretrainedModel(config).to(device)
 
+    # optimizer
+    logger.info("prepare optimizer, lr_scheduler")
+    num_params = 0
+    for p in generator.parameters():
+        num_params += p.numel()
+    logger.info("total parameters (generator): {:.3f}M".format(num_params/1e6))
+
+    optim_g = torch.optim.AdamW(generator.parameters(), config.learning_rate, betas=[config.adam_b1, config.adam_b2])
+    optim_d = torch.optim.AdamW(discriminator.parameters(), config.learning_rate, betas=[config.adam_b1, config.adam_b2])
+
+    scheduler_g = torch.optim.lr_scheduler.ExponentialLR(optim_g, gamma=config.lr_decay, last_epoch=last_epoch)
+    scheduler_d = torch.optim.lr_scheduler.ExponentialLR(optim_d, gamma=config.lr_decay, last_epoch=last_epoch)
+
     # resume training
     last_epoch = -1
     for epoch_i in serialization_dir.glob("epoch-*"):
@@ -179,6 +192,8 @@ def main():
         logger.info(f"resume from epoch-{last_epoch}.")
         generator_pt = serialization_dir / f"epoch-{last_epoch}/generator.pt"
         discriminator_pt = serialization_dir / f"epoch-{last_epoch}/discriminator.pt"
+        optim_g_pth = serialization_dir / f"epoch-{last_epoch}/optim_g.pth"
+        optim_d_pth = serialization_dir / f"epoch-{last_epoch}/optim_d.pth"
 
         logger.info(f"load state dict for generator.")
         with open(generator_pt.as_posix(), "rb") as f:
@@ -189,18 +204,14 @@ def main():
             state_dict = torch.load(f, map_location="cpu", weights_only=True)
         discriminator.load_state_dict(state_dict, strict=True)
 
-    # optimizer
-    logger.info("prepare optimizer, lr_scheduler")
-    num_params = 0
-    for p in generator.parameters():
-        num_params += p.numel()
-    logger.info("total parameters (generator): {:.3f}M".format(num_params/1e6))
-
-    optim_g = torch.optim.AdamW(generator.parameters(), config.learning_rate, betas=[config.adam_b1, config.adam_b2])
-    optim_d = torch.optim.AdamW(discriminator.parameters(), config.learning_rate, betas=[config.adam_b1, config.adam_b2])
-
-    scheduler_g = torch.optim.lr_scheduler.ExponentialLR(optim_g, gamma=config.lr_decay, last_epoch=last_epoch)
-    scheduler_d = torch.optim.lr_scheduler.ExponentialLR(optim_d, gamma=config.lr_decay, last_epoch=last_epoch)
+        logger.info(f"load state dict for optim_g.")
+        with open(optim_g_pth.as_posix(), "rb") as f:
+            state_dict = torch.load(f, map_location="cpu", weights_only=True)
+        optim_g.load_state_dict(state_dict, strict=True)
+        logger.info(f"load state dict for optim_d.")
+        with open(optim_d_pth.as_posix(), "rb") as f:
+            state_dict = torch.load(f, map_location="cpu", weights_only=True)
+        optim_d.load_state_dict(state_dict, strict=True)
 
     # training loop
 
@@ -369,6 +380,10 @@ def main():
         generator.save_pretrained(epoch_dir.as_posix())
         discriminator.save_pretrained(epoch_dir.as_posix())
 
+        # save optim
+        torch.save(optim_d.state_dict(), (epoch_dir / "optim_d.pth").as_posix())
+        torch.save(optim_g.state_dict(), (epoch_dir / "optim_g.pth").as_posix())
+
         model_list.append(epoch_dir)
         if len(model_list) >= args.num_serialized_models_to_keep:
             model_to_delete: Path = model_list.pop(0)