update

examples/conv_tasnet/step_2_train_model.py

@@ -233,6 +233,7 @@ def main():
     average_neg_stoi_loss = 1000000000
 
     model_list = list()
+    best_idx_epoch = None
     best_steps = None
     best_metric = None
     patience_count = 0
@@ -409,10 +410,12 @@ def main():
 
                     # save metric
                     if best_metric is None:
+                        best_idx_epoch = idx_epoch
                         best_steps = total_steps
                         best_metric = average_pesq_score
                     elif average_pesq_score > best_metric:
                         # great is better.
+                        best_idx_epoch = idx_epoch
                         best_steps = total_steps
                         best_metric = average_pesq_score
                     else:
@@ -420,6 +423,7 @@ def main():
 
                     metrics = {
                         "idx_epoch": idx_epoch,
+                        "best_idx_epoch": best_idx_epoch,
                         "best_steps": best_steps,
                         "pesq_score": average_pesq_score,
                         "loss": average_loss,
@@ -433,14 +437,14 @@ def main():
 
                     # save best
                     best_dir = serialization_dir / "best"
-                    if best_idx_epoch == idx_epoch:
+                    if best_idx_epoch == idx_epoch and best_steps == total_steps:
                         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:
+                    if best_idx_epoch == idx_epoch and best_steps == total_steps:
                         patience_count = 0
                     else:
                         patience_count += 1