update

examples/conv_tasnet/step_2_train_model.py

@@ -171,19 +171,19 @@ def main():
     optimizer = torch.optim.AdamW(model.parameters(), config.lr)
 
     # resume training
-    last_steps = -1
+    last_step_idx = -1
     last_epoch = -1
-    for step_i in serialization_dir.glob("steps-*"):
-        step_i = Path(step_i)
-        step_idx = step_i.stem.split("-")[1]
+    for step_idx_str in serialization_dir.glob("steps-*"):
+        step_idx_str = Path(step_idx_str)
+        step_idx = step_idx_str.stem.split("-")[1]
         step_idx = int(step_idx)
-        if step_idx > last_steps:
-            last_steps = step_idx
+        if step_idx > last_step_idx:
+            last_step_idx = 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"
+    if last_step_idx != -1:
+        logger.info(f"resume from steps-{last_step_idx}.")
+        model_pt = serialization_dir / f"steps-{last_step_idx}/model.pt"
+        optimizer_pth = serialization_dir / f"steps-{last_step_idx}/optimizer.pth"
 
         logger.info(f"load state dict for model.")
         with open(model_pt.as_posix(), "rb") as f:
@@ -233,13 +233,13 @@ def main():
     average_neg_stoi_loss = 1000000000
 
     model_list = list()
-    best_idx_epoch = None
-    best_steps = None
+    best_epoch_idx = None
+    best_step_idx = None
     best_metric = None
     patience_count = 0
 
     logger.info("training")
-    for idx_epoch in range(max(0, last_epoch+1), args.max_epochs):
+    for epoch_idx in range(max(0, last_epoch+1), args.max_epochs):
         # train
         model.train()
 
@@ -251,10 +251,10 @@ def main():
         total_mr_stft_loss = 0.
         total_batches = 0.
 
-        total_steps = 0 if last_steps == -1 else last_steps
+        step_idx = 0 if last_step_idx == -1 else last_step_idx
         progress_bar_train = tqdm(
-            initial=total_steps,
-            desc="Training; epoch-{}".format(idx_epoch),
+            initial=step_idx,
+            desc="Training; epoch-{}".format(epoch_idx),
         )
         for train_batch in train_data_loader:
             clean_audios, noisy_audios = train_batch
@@ -314,8 +314,8 @@ def main():
             })
 
             # evaluation
-            total_steps += 1
-            if total_steps % config.eval_steps == 0:
+            step_idx += 1
+            if step_idx % config.eval_steps == 0:
                 with torch.no_grad():
                     torch.cuda.empty_cache()
 
@@ -328,7 +328,7 @@ def main():
 
                     progress_bar_train.close()
                     progress_bar_eval = tqdm(
-                        desc="Evaluation; step-{}k".format(int(total_steps/1000)),
+                        desc="Evaluation; step-{}k".format(int(step_idx/1000)),
                     )
                     for eval_batch in valid_data_loader:
                         clean_audios, noisy_audios = eval_batch
@@ -394,7 +394,7 @@ def main():
                     )
 
                     # save path
-                    save_dir = serialization_dir / "steps-{}k".format(int(total_steps/1000))
+                    save_dir = serialization_dir / "steps-{}k".format(int(step_idx/1000))
                     save_dir.mkdir(parents=True, exist_ok=False)
 
                     # save models
@@ -410,26 +410,27 @@ def main():
 
                     # save metric
                     if best_metric is None:
-                        best_idx_epoch = idx_epoch
-                        best_steps = total_steps
+                        best_epoch_idx = epoch_idx
+                        best_step_idx = step_idx
                         best_metric = average_pesq_score
                     elif average_pesq_score > best_metric:
                         # great is better.
-                        best_idx_epoch = idx_epoch
-                        best_steps = total_steps
+                        best_epoch_idx = epoch_idx
+                        best_step_idx = step_idx
                         best_metric = average_pesq_score
                     else:
                         pass
 
                     metrics = {
-                        "idx_epoch": idx_epoch,
-                        "best_idx_epoch": best_idx_epoch,
-                        "best_steps": best_steps,
+                        "epoch_idx": epoch_idx,
+                        "best_epoch_idx": best_epoch_idx,
+                        "best_step_idx": best_step_idx,
                         "pesq_score": average_pesq_score,
                         "loss": average_loss,
                         "ae_loss": average_ae_loss,
                         "neg_si_snr_loss": average_neg_si_snr_loss,
                         "neg_stoi_loss": average_neg_stoi_loss,
+                        "mr_stft_loss": average_mr_stft_loss,
                     }
                     metrics_filename = save_dir / "metrics_epoch.json"
                     with open(metrics_filename, "w", encoding="utf-8") as f:
@@ -437,14 +438,14 @@ def main():
 
                     # save best
                     best_dir = serialization_dir / "best"
-                    if best_idx_epoch == idx_epoch and best_steps == total_steps:
+                    if best_epoch_idx == epoch_idx and best_step_idx == step_idx:
                         if best_dir.exists():
                             shutil.rmtree(best_dir)
                         shutil.copytree(save_dir, best_dir)
 
                     # early stop
                     early_stop_flag = False
-                    if best_idx_epoch == idx_epoch and best_steps == total_steps:
+                    if best_epoch_idx == epoch_idx and best_step_idx == step_idx:
                         patience_count = 0
                     else:
                         patience_count += 1