update

examples/conv_tasnet/run.sh

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

@@ -298,21 +298,20 @@ def main():
             # evaluation
             total_steps += 1
             if total_steps % args.eval_steps == 0:
-                model.eval()
-                torch.cuda.empty_cache()
-
-                total_pesq_score = 0.
-                total_loss = 0.
-                total_ae_loss = 0.
-                total_neg_si_snr_loss = 0.
-                total_neg_stoi_loss = 0.
-                total_batches = 0.
-
-                progress_bar_train.close()
-                progress_bar_eval = tqdm(
-                    desc="Evaluation; step-{}".format(total_steps),
-                )
                 with torch.no_grad():
+                    torch.cuda.empty_cache()
+
+                    total_pesq_score = 0.
+                    total_loss = 0.
+                    total_ae_loss = 0.
+                    total_neg_si_snr_loss = 0.
+                    total_neg_stoi_loss = 0.
+                    total_batches = 0.
+
+                    progress_bar_train.close()
+                    progress_bar_eval = tqdm(
+                        desc="Evaluation; step-{}".format(total_steps),
+                    )
                     for eval_batch in valid_data_loader:
                         clean_audios, noisy_audios = eval_batch
                         clean_audios = clean_audios.to(device)
@@ -357,71 +356,71 @@ def main():
                             "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)
-                save_dir.mkdir(parents=True, exist_ok=False)
-
-                # save models
-                model.save_pretrained(save_dir.as_posix())
-
-                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(), (save_dir / "optimizer.pth").as_posix())
-
-                # save metric
-                if best_metric is None:
-                    best_idx_epoch = idx_epoch
-                    best_metric = average_pesq_score
-                elif average_pesq_score > best_metric:
-                    # great is better.
-                    best_idx_epoch = idx_epoch
-                    best_metric = average_pesq_score
-                else:
-                    pass
-
-                metrics = {
-                    "idx_epoch": idx_epoch,
-                    "best_idx_epoch": best_idx_epoch,
-                    "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,
-                }
-                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)
-
-                # save best
-                best_dir = serialization_dir / "best"
-                if best_idx_epoch == idx_epoch:
-                    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:
-                    patience_count = 0
-                else:
-                    patience_count += 1
-                if patience_count >= args.patience:
-                    early_stop_flag = True
-
-                # early stop
-                if early_stop_flag:
-                    break
+                    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)
+                    save_dir.mkdir(parents=True, exist_ok=False)
+
+                    # save models
+                    model.save_pretrained(save_dir.as_posix())
+
+                    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(), (save_dir / "optimizer.pth").as_posix())
+
+                    # save metric
+                    if best_metric is None:
+                        best_idx_epoch = idx_epoch
+                        best_metric = average_pesq_score
+                    elif average_pesq_score > best_metric:
+                        # great is better.
+                        best_idx_epoch = idx_epoch
+                        best_metric = average_pesq_score
+                    else:
+                        pass
+
+                    metrics = {
+                        "idx_epoch": idx_epoch,
+                        "best_idx_epoch": best_idx_epoch,
+                        "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,
+                    }
+                    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)
+
+                    # save best
+                    best_dir = serialization_dir / "best"
+                    if best_idx_epoch == idx_epoch:
+                        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:
+                        patience_count = 0
+                    else:
+                        patience_count += 1
+                    if patience_count >= args.patience:
+                        early_stop_flag = True
+
+                    # early stop
+                    if early_stop_flag:
+                        break
 
     return