update

examples/conv_tasnet/run.sh

@@ -3,7 +3,7 @@
 : <<'END'
 
 
-sh run.sh --stage 2 --stop_stage 2 --system_version centos --file_folder_name file_dir --final_model_name conv-tasnet-dns3-20250319 \
+sh run.sh --stage 1 --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

@@ -48,11 +48,10 @@ def get_args():
     parser.add_argument("--serialization_dir", default="serialization_dir", type=str)
     parser.add_argument("--seed", default=1234, type=int)
 
-    parser.add_argument("--eval_steps", default=5000, type=int)
+    parser.add_argument("--eval_steps", default=25000, type=int)
 
     parser.add_argument("--config_file", default="config.yaml", type=str)
 
-
     args = parser.parse_args()
     return args
 
@@ -240,7 +239,7 @@ def main():
         total_mr_stft_loss = 0.
         total_batches = 0.
         total_steps = 0
-        progress_bar = tqdm(
+        progress_bar_train = tqdm(
             desc="Training; epoch-{}".format(idx_epoch),
         )
         for train_batch in train_data_loader:
@@ -286,8 +285,8 @@ def main():
             average_neg_stoi_loss = round(total_neg_stoi_loss / total_batches, 4)
             average_mr_stft_loss = round(total_mr_stft_loss / total_batches, 4)
 
-            progress_bar.update(1)
-            progress_bar.set_postfix({
+            progress_bar_train.update(1)
+            progress_bar_train.set_postfix({
                 "pesq_score": average_pesq_score,
                 "loss": average_loss,
                 "ae_loss": average_ae_loss,
@@ -309,7 +308,8 @@ def main():
                 total_neg_stoi_loss = 0.
                 total_batches = 0.
 
-                progress_bar = tqdm(
+                progress_bar_train.close()
+                progress_bar_eval = tqdm(
                     desc="Evaluation; step-{}".format(total_steps),
                 )
                 with torch.no_grad():
@@ -348,16 +348,21 @@ def main():
                         average_neg_stoi_loss = round(total_neg_stoi_loss / total_batches, 4)
                         average_mr_stft_loss = round(total_mr_stft_loss / total_batches, 4)
 
-                        progress_bar.update(1)
-                        progress_bar.set_postfix({
+                        progress_bar_eval.update(1)
+                        progress_bar_eval.set_postfix({
                             "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,
-
                         })
+                progress_bar_eval.close()
+                progress_bar_train = tqdm(
+                    initial=progress_bar_train.n,
+                    postfix=progress_bar_train.postfix,
+                    desc=progress_bar_train.desc,
+                )
 
                 # save path
                 save_dir = serialization_dir / "steps-{}".format(total_steps)