update

examples/conv_tasnet/step_2_train_model.py

@@ -244,8 +244,6 @@ def main():
             desc="Training; epoch-{}".format(idx_epoch),
         )
         for batch in train_data_loader:
-            total_steps += 1
-
             clean_audios, noisy_audios = batch
             clean_audios = clean_audios.to(device)
             noisy_audios = noisy_audios.to(device)
@@ -299,6 +297,7 @@ def main():
             })
 
             # evaluation
+            total_steps += 1
             if total_steps % args.eval_steps:
                 model.eval()
                 torch.cuda.empty_cache()
@@ -311,7 +310,7 @@ def main():
                 total_batches = 0.
 
                 progress_bar = tqdm(
-                    desc="Evaluation; epoch-{}".format(idx_epoch),
+                    desc="Evaluation; step-{}".format(total_steps),
                 )
                 with torch.no_grad():
                     for batch in valid_data_loader:
@@ -361,28 +360,28 @@ def main():
                         })
 
             # save path
-            epoch_dir = serialization_dir / "epoch-{}".format(idx_epoch)
-            epoch_dir.mkdir(parents=True, exist_ok=False)
+            save_dir = serialization_dir / "steps-{}".format(total_steps)
+            save_dir.mkdir(parents=True, exist_ok=False)
 
             # save models
-            model.save_pretrained(epoch_dir.as_posix())
+            model.save_pretrained(save_dir.as_posix())
 
-            model_list.append(epoch_dir)
+            model_list.append(save_dir)
             if len(model_list) >= args.num_serialized_models_to_keep:
                 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())
+            torch.save(optimizer.state_dict(), (save_dir / "optimizer.pth").as_posix())
 
             # save metric
             if best_metric is None:
                 best_idx_epoch = idx_epoch
-                best_metric = average_loss
-            elif average_loss < best_metric:
+                best_metric = average_pesq_score
+            elif average_pesq_score > best_metric:
                 # great is better.
                 best_idx_epoch = idx_epoch
-                best_metric = average_loss
+                best_metric = average_pesq_score
             else:
                 pass
 
@@ -395,7 +394,7 @@ def main():
                 "neg_si_snr_loss": average_neg_si_snr_loss,
                 "neg_stoi_loss": average_neg_stoi_loss,
             }
-            metrics_filename = epoch_dir / "metrics_epoch.json"
+            metrics_filename = save_dir / "metrics_epoch.json"
             with open(metrics_filename, "w", encoding="utf-8") as f:
                 json.dump(metrics, f, indent=4, ensure_ascii=False)
 
@@ -404,7 +403,7 @@ def main():
             if best_idx_epoch == idx_epoch:
                 if best_dir.exists():
                     shutil.rmtree(best_dir)
-                shutil.copytree(epoch_dir, best_dir)
+                shutil.copytree(save_dir, best_dir)
 
             # early stop
             early_stop_flag = False