update

examples/mpnet_aishell/step_2_train_model.py

@@ -167,11 +167,26 @@ def main():
     discriminator = MetricDiscriminatorPretrainedModel(config).to(device)
 
     # resume training
+    epoch_max = -1
     for epoch_i in serialization_dir.glob("epoch-*"):
         epoch_i = Path(epoch_i)
         epoch_idx = epoch_i.stem.split("-")[1]
-        print(f"epoch_idx: {epoch_idx}")
-    exit(0)
+        if epoch_idx > epoch_max:
+            epoch_max = epoch_idx
+
+    if epoch_max != -1:
+        logger.info(f"resume from epoch-{epoch_max}.")
+        generator_pt = serialization_dir / f"epoch-{epoch_max}/generator.pt"
+        discriminator_pt = serialization_dir / f"epoch-{epoch_max}/discriminator.pt"
+
+        logger.info(f"load state dict for generator.")
+        with open(generator_pt.as_posix(), "rb") as f:
+            state_dict = torch.load(f, map_location="cpu")
+        generator.load_state_dict(state_dict, strict=True)
+        logger.info(f"load state dict for discriminator.")
+        with open(discriminator_pt.as_posix(), "rb") as f:
+            state_dict = torch.load(f, map_location="cpu")
+        discriminator.load_state_dict(state_dict, strict=True)
 
     # optimizer
     logger.info("prepare optimizer, lr_scheduler")
@@ -285,6 +300,8 @@ def main():
 
         # evaluation
         generator.eval()
+        discriminator.eval()
+
         torch.cuda.empty_cache()
         total_pesq_score = 0.
         total_mag_err = 0.
@@ -361,6 +378,7 @@ def main():
             best_idx_epoch = idx_epoch
             best_metric = pesq_metric
         elif pesq_metric > best_metric:
+            # great is better.
             best_idx_epoch = idx_epoch
             best_metric = pesq_metric
         else: