update

examples/conv_tasnet/run.sh

@@ -3,7 +3,7 @@
 : <<'END'
 
 
-sh run.sh --stage 1 --stop_stage 2 --system_version centos --file_folder_name file_dir --final_model_name conv-tasnet-dns3-20250319 \
+sh run.sh --stage 2 --stop_stage 2 --system_version centos --file_folder_name file_dir --final_model_name conv-tasnet-dns3-20250319 \
 --noise_dir "/data/tianxing/HuggingDatasets/nx_noise/data/noise/dns3-noise" \
 --speech_dir "/data/tianxing/HuggingDatasets/nx_noise/data/speech/dns3-speech" \
 --max_epochs 400

examples/conv_tasnet/step_2_train_model.py

@@ -171,18 +171,18 @@ def main():
     optimizer = torch.optim.AdamW(model.parameters(), config.lr)
 
     # 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"
+    last_steps = -1
+    for step_i in serialization_dir.glob("steps-*"):
+        step_i = Path(step_i)
+        step_idx = step_i.stem.split("-")[1]
+        step_idx = int(step_idx)
+        if step_idx > last_steps:
+            last_steps = step_idx
+
+    if last_steps != -1:
+        logger.info(f"resume from steps-{last_steps}.")
+        model_pt = serialization_dir / f"steps-{last_steps}/model.pt"
+        optimizer_pth = serialization_dir / f"steps-{last_steps}/optimizer.pth"
 
         logger.info(f"load state dict for model.")
         with open(model_pt.as_posix(), "rb") as f:
@@ -197,7 +197,7 @@ def main():
     if config.lr_scheduler == "CosineAnnealingLR":
         lr_scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(
             optimizer,
-            last_epoch=last_epoch,
+            last_epoch=-1,
             # T_max=10 * config.eval_steps,
             # eta_min=0.01 * config.lr,
             **config.lr_scheduler_kwargs,