update

examples/dfnet/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 frcrn \
+sh run.sh --stage 2 --stop_stage 2 --system_version centos --file_folder_name file_dir --final_model_name dfnet-dns3 \
 --noise_dir "/data/tianxing/HuggingDatasets/nx_noise/data/noise/dns3-noise" \
 --speech_dir "/data/tianxing/HuggingDatasets/nx_noise/data/speech/dns3-speech"
 

examples/dfnet/yaml/config.yaml

@@ -51,3 +51,21 @@ df_lookahead: 2
 
 # runtime
 use_post_filter: true
+
+# train
+lr: 0.001
+lr_scheduler: "CosineAnnealingLR"
+lr_scheduler_kwargs:
+  T_max: 250000
+  eta_min: 0.0001
+
+max_epochs: 100
+clip_grad_norm: 10.0
+seed: 1234
+
+min_snr_db: -10
+max_snr_db: 20
+
+num_workers: 8
+batch_size: 32
+eval_steps: 10000

examples/frcrn/step_2_train_model.py

@@ -1,6 +1,8 @@
 #!/usr/bin/python3
 # -*- coding: utf-8 -*-
 """
+https://arxiv.org/abs/2206.07293
+
 FRCRN 论文中:
 在 WSJ0 数据集上训练了 120 个 epoch 得到 pesq 3.62, stoi 98.24, si-snr 21.33
 
@@ -188,17 +190,17 @@ def main():
     if last_step_idx != -1:
         logger.info(f"resume from steps-{last_step_idx}.")
         model_pt = serialization_dir / f"steps-{last_step_idx}/model.pt"
-        optimizer_pth = serialization_dir / f"steps-{last_step_idx}/optimizer.pth"
+        # optimizer_pth = serialization_dir / f"steps-{last_step_idx}/optimizer.pth"
 
         logger.info(f"load state dict for model.")
         with open(model_pt.as_posix(), "rb") as f:
             state_dict = torch.load(f, map_location="cpu", weights_only=True)
         model.load_state_dict(state_dict, strict=True)
 
-        logger.info(f"load state dict for optimizer.")
-        with open(optimizer_pth.as_posix(), "rb") as f:
-            state_dict = torch.load(f, map_location="cpu", weights_only=True)
-        optimizer.load_state_dict(state_dict)
+        # logger.info(f"load state dict for optimizer.")
+        # with open(optimizer_pth.as_posix(), "rb") as f:
+        #     state_dict = torch.load(f, map_location="cpu", weights_only=True)
+        # optimizer.load_state_dict(state_dict)
 
     if config.lr_scheduler == "CosineAnnealingLR":
         lr_scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(
@@ -377,15 +379,12 @@ def main():
                         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_epoch_idx = epoch_idx
                         best_step_idx = step_idx
                         best_metric = average_pesq_score
-                    elif average_pesq_score > best_metric:
+                    elif average_pesq_score >= best_metric:
                         # great is better.
                         best_epoch_idx = epoch_idx
                         best_step_idx = step_idx

examples/mpnet/yaml/config.yaml

@@ -25,3 +25,6 @@ dist_config:
   dist_backend: nccl
   dist_url: tcp://localhost:54321
   world_size: 1
+
+discriminator_dim: 32
+discriminator_in_channel: 2

main.py

@@ -16,6 +16,7 @@ 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
 
 log.setup_size_rotating(log_directory=log_directory)
 
@@ -93,6 +94,12 @@ denoise_engines = {
             "pretrained_model_path_or_zip_file": (project_path / "trained_models/mpnet-aishell-11-epoch.zip").as_posix()
         }
     },
+    "frcrn-dns3": {
+        "infer_cls": InferenceFRCRN,
+        "kwargs": {
+            "pretrained_model_path_or_zip_file": (project_path / "trained_models/frcrn-dns3-220k-steps.zip").as_posix()
+        }
+    },
 }
 
 

toolbox/torchaudio/models/dfnet/configuration_dfnet.py

@@ -50,6 +50,22 @@ class DfNetConfig(PretrainedConfig):
                  df_lookahead: int = 2,
 
                  use_post_filter: bool = False,
+
+                 lr: float = 0.001,
+                 lr_scheduler: str = "CosineAnnealingLR",
+                 lr_scheduler_kwargs: dict = None,
+
+                 max_epochs: int = 100,
+                 clip_grad_norm: float = 10.,
+                 seed: int = 1234,
+
+                 min_snr_db: float = -10,
+                 max_snr_db: float = 20,
+
+                 num_workers: int = 4,
+                 batch_size: int = 4,
+                 eval_steps: int = 25000,
+
                  **kwargs
                  ):
         super(DfNetConfig, self).__init__(**kwargs)
@@ -104,6 +120,22 @@ class DfNetConfig(PretrainedConfig):
         # runtime
         self.use_post_filter = use_post_filter
 
+        #
+        self.lr = lr
+        self.lr_scheduler = lr_scheduler
+        self.lr_scheduler_kwargs = lr_scheduler_kwargs or dict()
+
+        self.max_epochs = max_epochs
+        self.clip_grad_norm = clip_grad_norm
+        self.seed = seed
+
+        self.min_snr_db = min_snr_db
+        self.max_snr_db = max_snr_db
+
+        self.num_workers = num_workers
+        self.batch_size = batch_size
+        self.eval_steps = eval_steps
+
 
 if __name__ == "__main__":
     pass

toolbox/torchaudio/models/frcrn/inference_frcrn.py

@@ -0,0 +1,114 @@
+#!/usr/bin/python3
+# -*- coding: utf-8 -*-
+import logging
+from pathlib import Path
+import shutil
+import tempfile, time
+import zipfile
+
+import librosa
+import numpy as np
+import torch
+import torchaudio
+
+torch.set_num_threads(1)
+
+from project_settings import project_path
+from toolbox.torchaudio.models.frcrn.configuration_frcrn import FRCRNConfig
+from toolbox.torchaudio.models.frcrn.modeling_frcrn import FRCRNPretrainedModel, MODEL_FILE
+
+logger = logging.getLogger("toolbox")
+
+
+class InferenceFRCRN(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)
+
+        logger.info(f"loading model; model_file: {self.pretrained_model_path_or_zip_file}")
+        config, model = self.load_models(self.pretrained_model_path_or_zip_file)
+        logger.info(f"model loading completed; model_file: {self.pretrained_model_path_or_zip_file}")
+
+        self.config = config
+        self.model = model
+        self.model.to(device)
+        self.model.eval()
+
+    def load_models(self, model_path: str):
+        model_path = Path(model_path)
+        if model_path.name.endswith(".zip"):
+            with zipfile.ZipFile(model_path.as_posix(), "r") as f_zip:
+                out_root = Path(tempfile.gettempdir()) / "nx_denoise"
+                out_root.mkdir(parents=True, exist_ok=True)
+                f_zip.extractall(path=out_root)
+            model_path = out_root / model_path.stem
+
+        config = FRCRNConfig.from_pretrained(
+            pretrained_model_name_or_path=model_path.as_posix(),
+        )
+        model = FRCRNPretrainedModel.from_pretrained(
+            pretrained_model_name_or_path=model_path.as_posix(),
+        )
+        model.to(self.device)
+        model.eval()
+
+        shutil.rmtree(model_path)
+        return config, model
+
+    def enhancement_by_ndarray(self, noisy_audio: np.ndarray) -> np.ndarray:
+        noisy_audio = torch.tensor(noisy_audio, dtype=torch.float32)
+        noisy_audio = noisy_audio.unsqueeze(dim=0)
+
+        # noisy_audio shape: [batch_size, n_samples]
+        enhanced_audio = self.enhancement_by_tensor(noisy_audio)
+        # noisy_audio shape: [n_samples,]
+        return enhanced_audio.cpu().numpy()
+
+    def enhancement_by_tensor(self, noisy_audio: torch.Tensor) -> torch.Tensor:
+        if torch.max(noisy_audio) > 1 or torch.min(noisy_audio) < -1:
+            raise AssertionError(f"The value range of audio samples should be between -1 and 1.")
+
+        # noisy_audio shape: [batch_size, num_samples]
+        noisy_audios = noisy_audio.to(self.device)
+
+        with torch.no_grad():
+            est_spec, est_wav, est_mask = self.model.forward(noisy_audios)
+
+        # shape: [batch_size, num_samples]
+        enhanced_audio = torch.unsqueeze(est_wav, dim=1)
+        # shape: [batch_size, 1, num_samples]
+
+        enhanced_audio = enhanced_audio[0]
+
+        # enhanced_audio shape: [channels, num_samples]
+        return enhanced_audio
+
+
+def main():
+    model_zip_file = project_path / "trained_models/frcrn-dns3.zip"
+    infer_model = InferenceFRCRN(model_zip_file)
+
+    sample_rate = 8000
+    noisy_audio_file = project_path / "data/examples/ai_agent/dfaaf264-b5e3-4ca2-b5cb-5b6d637d962d_section_5.wav"
+    noisy_audio, sample_rate = librosa.load(
+        noisy_audio_file.as_posix(),
+        sr=sample_rate,
+    )
+    duration = librosa.get_duration(y=noisy_audio, sr=sample_rate)
+    # noisy_audio = noisy_audio[int(7*sample_rate):int(9*sample_rate)]
+    noisy_audio = torch.tensor(noisy_audio, dtype=torch.float32)
+    noisy_audio = noisy_audio.unsqueeze(dim=0)
+
+    begin = time.time()
+    enhanced_audio = infer_model.enhancement_by_tensor(noisy_audio)
+    time_cost = time.time() - begin
+    print(f"enhanced_audio.shape: {enhanced_audio.shape}, time_cost: {time_cost:.4f}, audio_duration: {duration:.4f}, fpr: {time_cost / duration:.4f}")
+
+    filename = "enhanced_audio.wav"
+    torchaudio.save(filename, enhanced_audio.detach().cpu(), sample_rate)
+
+    return
+
+
+if __name__ == "__main__":
+    main()

toolbox/torchaudio/models/mpnet/yaml/config.yaml

@@ -25,3 +25,6 @@ dist_config:
   dist_backend: nccl
   dist_url: tcp://localhost:54321
   world_size: 1
+
+discriminator_dim: 32
+discriminator_in_channel: 2