update

examples/conv_tasnet/step_2_train_model.py

@@ -214,6 +214,7 @@ def main():
     elif config.lr_scheduler == "MultiStepLR":
         lr_scheduler = torch.optim.lr_scheduler.MultiStepLR(
             optimizer,
+            last_epoch=last_epoch,
             milestones=[10000, 20000, 30000, 40000, 50000], gamma=0.5
         )
     else:
@@ -345,7 +346,7 @@ def main():
 
                     progress_bar_train.close()
                     progress_bar_eval = tqdm(
-                        desc="Evaluation; step-{}k".format(int(step_idx/1000)),
+                        desc="Evaluation; steps-{}k".format(int(step_idx/1000)),
                     )
                     for eval_batch in valid_data_loader:
                         clean_audios, noisy_audios = eval_batch
@@ -418,7 +419,7 @@ def main():
                     )
 
                     # save path
-                    save_dir = serialization_dir / "steps-{}k".format(int(step_idx/1000))
+                    save_dir = serialization_dir / "steps-{}".format(step_idx)
                     save_dir.mkdir(parents=True, exist_ok=False)
 
                     # save models
@@ -455,6 +456,7 @@ def main():
                         "neg_si_snr_loss": average_neg_si_snr_loss,
                         "neg_stoi_loss": average_neg_stoi_loss,
                         "mr_stft_loss": average_mr_stft_loss,
+                        "pesq_loss": average_pesq_loss,
                     }
                     metrics_filename = save_dir / "metrics_epoch.json"
                     with open(metrics_filename, "w", encoding="utf-8") as f: