update

examples/clean_unet_aishell/step_2_train_model.py

@@ -171,6 +171,31 @@ def main():
     # optimizer
     logger.info("prepare optimizer, lr_scheduler, loss_fn, categorical_accuracy")
     optimizer = torch.optim.AdamW(model.parameters(), args.learning_rate)
+
+    # resume training
+    last_epoch = -1
+    for epoch_i in serialization_dir.glob("epoch-*"):
+        epoch_i = Path(epoch_i)
+        epoch_idx = epoch_i.stem.split("-")[1]
+        epoch_idx = int(epoch_idx)
+        if epoch_idx > last_epoch:
+            last_epoch = epoch_idx
+
+    if last_epoch != -1:
+        logger.info(f"resume from epoch-{last_epoch}.")
+        model_pt = serialization_dir / f"epoch-{last_epoch}/model.pt"
+        optimizer_pth = serialization_dir / f"epoch-{last_epoch}/optimizer.pth"
+
+        logger.info(f"load state dict for generator.")
+        with open(model_pt.as_posix(), "rb") as f:
+            state_dict = torch.load(f, map_location="cpu", weights_only=True)
+        model.load_state_dict(state_dict, strict=True)
+
+        logger.info(f"load state dict for optimizer.")
+        with open(optimizer_pth.as_posix(), "rb") as f:
+            state_dict = torch.load(f, map_location="cpu", weights_only=True)
+        optimizer.load_state_dict(state_dict)
+
     lr_scheduler = LinearWarmupCosineDecay(
         optimizer,
         lr_max=args.learning_rate,
@@ -180,19 +205,14 @@ def main():
         warmup_proportion=0.05,
         phase=("linear", "cosine"),
     )
+
     # ae_loss_fn = nn.MSELoss(reduction="mean")
     ae_loss_fn = nn.L1Loss(reduction="mean").to(device)
 
     mr_stft_loss_fn = MultiResolutionSTFTLoss(
-        fft_sizes=[512],
-        hop_sizes=[50],
-        win_lengths=[240],
-        # fft_sizes=[256, 512, 1024],
-        # hop_sizes=[25, 50, 120],
-        # win_lengths=[120, 240, 600],
-        # fft_sizes=[512, 1024, 2048],
-        # hop_sizes=[50, 120, 240],
-        # win_lengths=[240, 600, 1200],
+        fft_sizes=[256, 512, 1024],
+        hop_sizes=[25, 50, 120],
+        win_lengths=[120, 240, 600],
         sc_lambda=0.5,
         mag_lambda=0.5,
         band="full"
@@ -343,6 +363,9 @@ def main():
             model_to_delete: Path = model_list.pop(0)
             shutil.rmtree(model_to_delete.as_posix())
 
+        # save optim
+        torch.save(optimizer.state_dict(), (epoch_dir / "optimizer.pth").as_posix())
+
         # save metric
         if best_metric is None:
             best_idx_epoch = idx_epoch