update

examples/dfnet2/step_2_train_model.py

@@ -318,10 +318,11 @@ def main():
             # evaluation
             step_idx += 1
             if step_idx % config.eval_steps == 0:
-                model.eval()
                 with torch.no_grad():
                     torch.cuda.empty_cache()
 
+                    model.eval()
+
                     total_pesq_score = 0.
                     total_loss = 0.
                     total_mr_stft_loss = 0.
@@ -384,82 +385,83 @@ def main():
                             "lsnr_loss": average_lsnr_loss,
                         })
 
-                    total_pesq_score = 0.
-                    total_loss = 0.
-                    total_mr_stft_loss = 0.
-                    total_neg_si_snr_loss = 0.
-                    total_mask_loss = 0.
-                    total_lsnr_loss = 0.
-                    total_batches = 0.
-
-                    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(step_idx)
-                    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 metric
-                    if best_metric is None:
-                        best_epoch_idx = epoch_idx
-                        best_step_idx = step_idx
-                        best_metric = average_pesq_score
-                    elif average_pesq_score >= best_metric:
-                        # great is better.
-                        best_epoch_idx = epoch_idx
-                        best_step_idx = step_idx
-                        best_metric = average_pesq_score
-                    else:
-                        pass
-
-                    metrics = {
-                        "epoch_idx": epoch_idx,
-                        "best_epoch_idx": best_epoch_idx,
-                        "best_step_idx": best_step_idx,
-                        "pesq_score": average_pesq_score,
-                        "loss": average_loss,
-                        "mr_stft_loss": average_mr_stft_loss,
-                        "neg_si_snr_loss": average_neg_si_snr_loss,
-                        "mask_loss": average_mask_loss,
-                        "lsnr_loss": average_lsnr_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_epoch_idx == epoch_idx and best_step_idx == step_idx:
-                        if best_dir.exists():
-                            shutil.rmtree(best_dir)
-                        shutil.copytree(save_dir, best_dir)
-
-                    # early stop
-                    early_stop_flag = False
-                    if best_epoch_idx == epoch_idx and best_step_idx == step_idx:
-                        patience_count = 0
-                    else:
-                        patience_count += 1
-                    if patience_count >= args.patience:
-                        early_stop_flag = True
-
-                    # early stop
-                    if early_stop_flag:
-                        break
                 model.train()
 
+                total_pesq_score = 0.
+                total_loss = 0.
+                total_mr_stft_loss = 0.
+                total_neg_si_snr_loss = 0.
+                total_mask_loss = 0.
+                total_lsnr_loss = 0.
+                total_batches = 0.
+
+                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(step_idx)
+                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 metric
+                if best_metric is None:
+                    best_epoch_idx = epoch_idx
+                    best_step_idx = step_idx
+                    best_metric = average_pesq_score
+                elif average_pesq_score >= best_metric:
+                    # great is better.
+                    best_epoch_idx = epoch_idx
+                    best_step_idx = step_idx
+                    best_metric = average_pesq_score
+                else:
+                    pass
+
+                metrics = {
+                    "epoch_idx": epoch_idx,
+                    "best_epoch_idx": best_epoch_idx,
+                    "best_step_idx": best_step_idx,
+                    "pesq_score": average_pesq_score,
+                    "loss": average_loss,
+                    "mr_stft_loss": average_mr_stft_loss,
+                    "neg_si_snr_loss": average_neg_si_snr_loss,
+                    "mask_loss": average_mask_loss,
+                    "lsnr_loss": average_lsnr_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_epoch_idx == epoch_idx and best_step_idx == step_idx:
+                    if best_dir.exists():
+                        shutil.rmtree(best_dir)
+                    shutil.copytree(save_dir, best_dir)
+
+                # early stop
+                early_stop_flag = False
+                if best_epoch_idx == epoch_idx and best_step_idx == step_idx:
+                    patience_count = 0
+                else:
+                    patience_count += 1
+                if patience_count >= args.patience:
+                    early_stop_flag = True
+
+                # early stop
+                if early_stop_flag:
+                    break
+
     return
 
 

main.py

@@ -18,9 +18,11 @@ import numpy as np
 import log
 from project_settings import environment, project_path, log_directory
 from toolbox.os.command import Command
-from toolbox.torchaudio.models.mpnet.inference_mpnet import InferenceMPNet
-from toolbox.torchaudio.models.frcrn.inference_frcrn import InferenceFRCRN
 from toolbox.torchaudio.models.dfnet.inference_dfnet import InferenceDfNet
+from toolbox.torchaudio.models.dfnet2.inference_dfnet2 import InferenceDfNet2
+from toolbox.torchaudio.models.dtln.inference_dtln import InferenceDTLN
+from toolbox.torchaudio.models.frcrn.inference_frcrn import InferenceFRCRN
+from toolbox.torchaudio.models.mpnet.inference_mpnet import InferenceMPNet
 
 
 log.setup_size_rotating(log_directory=log_directory)
@@ -66,6 +68,18 @@ def shell(cmd: str):
 
 
 denoise_engines = {
+    "dtln-nx-dns3": {
+        "infer_cls": InferenceDTLN,
+        "kwargs": {
+            "pretrained_model_path_or_zip_file": (project_path / "trained_models/dtln-nx-dns3.zip").as_posix()
+        }
+    },
+    "dfnet2-nx-dns3": {
+        "infer_cls": InferenceDfNet2,
+        "kwargs": {
+            "pretrained_model_path_or_zip_file": (project_path / "trained_models/dfnet2-nx-dns3.zip").as_posix()
+        }
+    },
     "dfnet-nx-dns3": {
         "infer_cls": InferenceDfNet,
         "kwargs": {

toolbox/torchaudio/models/dfnet2/inference_dfnet2.py

@@ -20,7 +20,7 @@ from toolbox.torchaudio.models.dfnet2.modeling_dfnet2 import DfNet2PretrainedMod
 logger = logging.getLogger("toolbox")
 
 
-class InferenceDfNet(object):
+class InferenceDfNet2(object):
     def __init__(self, pretrained_model_path_or_zip_file: str, device: str = "cpu"):
         self.pretrained_model_path_or_zip_file = pretrained_model_path_or_zip_file
         self.device = torch.device(device)
@@ -99,7 +99,7 @@ class InferenceDfNet(object):
 
 def main():
     model_zip_file = project_path / "trained_models/dfnet2-nx-dns3.zip"
-    infer_model = InferenceDfNet(model_zip_file)
+    infer_model = InferenceDfNet2(model_zip_file)
 
     sample_rate = 8000
     noisy_audio_file = project_path / "data/examples/ai_agent/chinese-3.wav"