update

examples/conv_tasnet/step_2_train_model.py

@@ -1,7 +1,7 @@
 #!/usr/bin/python3
 # -*- coding: utf-8 -*-
 """
-https://github.com/yxlu-0102/MP-SENet/blob/main/train.py
+https://github.com/kaituoxu/Conv-TasNet/tree/master/src
 """
 import argparse
 import json
@@ -42,14 +42,11 @@ def get_args():
     parser.add_argument("--max_epochs", default=200, type=int)
 
     parser.add_argument("--batch_size", default=8, type=int)
-    parser.add_argument("--learning_rate", default=1e-3, type=float)
     parser.add_argument("--num_serialized_models_to_keep", default=10, type=int)
     parser.add_argument("--patience", default=5, type=int)
     parser.add_argument("--serialization_dir", default="serialization_dir", type=str)
     parser.add_argument("--seed", default=1234, 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()
@@ -171,7 +168,7 @@ def main():
 
     # optimizer
     logger.info("prepare optimizer, lr_scheduler, loss_fn, categorical_accuracy")
-    optimizer = torch.optim.AdamW(model.parameters(), args.learning_rate)
+    optimizer = torch.optim.AdamW(model.parameters(), config.lr)
 
     # resume training
     last_epoch = -1
@@ -197,10 +194,21 @@ def main():
             state_dict = torch.load(f, map_location="cpu", weights_only=True)
         optimizer.load_state_dict(state_dict)
 
-    lr_scheduler = torch.optim.lr_scheduler.MultiStepLR(
-        optimizer,
-        milestones=[10000, 20000, 30000, 40000, 50000], gamma=0.5
-    )
+    if config.lr_scheduler == "CosineAnnealingLR":
+        lr_scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(
+            optimizer,
+            last_epoch=last_epoch,
+            # T_max=10 * config.eval_steps,
+            # eta_min=0.01 * config.lr,
+            **config.lr_scheduler_kwargs,
+        )
+    elif config.lr_scheduler == "MultiStepLR":
+        lr_scheduler = torch.optim.lr_scheduler.MultiStepLR(
+            optimizer,
+            milestones=[10000, 20000, 30000, 40000, 50000], gamma=0.5
+        )
+    else:
+        raise AssertionError(f"invalid lr_scheduler: {config.lr_scheduler}")
 
     ae_loss_fn = nn.L1Loss(reduction="mean").to(device)
     neg_si_snr_loss_fn = NegativeSISNRLoss(reduction="mean").to(device)
@@ -209,6 +217,8 @@ def main():
         fft_size_list=[256, 512, 1024],
         win_size_list=[120, 240, 480],
         hop_size_list=[25, 50, 100],
+        factor_sc=1.5,
+        factor_mag=1.0,
         reduction="mean"
     ).to(device)
 
@@ -222,7 +232,7 @@ def main():
     average_neg_stoi_loss = 1000000000
 
     model_list = list()
-    best_idx_epoch = None
+    best_steps = None
     best_metric = None
     patience_count = 0
 
@@ -260,7 +270,10 @@ def main():
             mr_stft_loss = mr_stft_loss_fn.forward(denoise_audios, clean_audios)
 
             # loss = 0.25 * ae_loss + 0.25 * neg_si_snr_loss
-            loss = 0.25 * ae_loss + 0.25 * neg_si_snr_loss + 0.25 * neg_stoi_loss + 0.25 * mr_stft_loss
+            # loss = 0.25 * ae_loss + 0.25 * neg_si_snr_loss + 0.25 * neg_stoi_loss + 0.25 * mr_stft_loss
+            # loss = 1.0 * ae_loss + 0.8 * neg_si_snr_loss + 0.5 * mr_stft_loss + 0.3 * neg_stoi_loss
+            # loss = 1.0 * ae_loss + 0.8 * neg_si_snr_loss + 0.7 * mr_stft_loss + 0.5 * neg_stoi_loss
+            loss = 2.0 * mr_stft_loss + 0.8 * ae_loss + 0.7 * neg_si_snr_loss + 0.5 * neg_stoi_loss
 
             denoise_audios_list_r = list(denoise_audios.detach().cpu().numpy())
             clean_audios_list_r = list(clean_audios.detach().cpu().numpy())
@@ -288,6 +301,7 @@ def main():
 
             progress_bar_train.update(1)
             progress_bar_train.set_postfix({
+                "lr": lr_scheduler.get_last_lr()[0],
                 "pesq_score": average_pesq_score,
                 "loss": average_loss,
                 "ae_loss": average_ae_loss,
@@ -298,7 +312,7 @@ def main():
 
             # evaluation
             total_steps += 1
-            if total_steps % args.eval_steps == 0:
+            if total_steps % config.eval_steps == 0:
                 with torch.no_grad():
                     torch.cuda.empty_cache()
 
@@ -311,7 +325,7 @@ def main():
 
                     progress_bar_train.close()
                     progress_bar_eval = tqdm(
-                        desc="Evaluation; step-{}".format(total_steps),
+                        desc="Evaluation; step-{}k".format(int(total_steps/1000)),
                     )
                     for eval_batch in valid_data_loader:
                         clean_audios, noisy_audios = eval_batch
@@ -327,7 +341,10 @@ def main():
                         mr_stft_loss = mr_stft_loss_fn.forward(denoise_audios, clean_audios)
 
                         # loss = 0.25 * ae_loss + 0.25 * neg_si_snr_loss
-                        loss = 0.25 * ae_loss + 0.25 * neg_si_snr_loss + 0.25 * neg_stoi_loss + 0.25 * mr_stft_loss
+                        # loss = 0.25 * ae_loss + 0.25 * neg_si_snr_loss + 0.25 * neg_stoi_loss + 0.25 * mr_stft_loss
+                        # loss = 1.0 * ae_loss + 0.8 * neg_si_snr_loss + 0.5 * mr_stft_loss + 0.3 * neg_stoi_loss
+                        # loss = 1.0 * ae_loss + 0.8 * neg_si_snr_loss + 0.7 * mr_stft_loss + 0.5 * neg_stoi_loss
+                        loss = 2.0 * mr_stft_loss + 0.8 * ae_loss + 0.7 * neg_si_snr_loss + 0.5 * neg_stoi_loss
 
                         denoise_audios_list_r = list(denoise_audios.detach().cpu().numpy())
                         clean_audios_list_r = list(clean_audios.detach().cpu().numpy())
@@ -350,6 +367,7 @@ def main():
 
                         progress_bar_eval.update(1)
                         progress_bar_eval.set_postfix({
+                            "lr": lr_scheduler.get_last_lr()[0],
                             "pesq_score": average_pesq_score,
                             "loss": average_loss,
                             "ae_loss": average_ae_loss,
@@ -373,7 +391,7 @@ def main():
                     )
 
                     # save path
-                    save_dir = serialization_dir / "steps-{}".format(total_steps)
+                    save_dir = serialization_dir / "steps-{}k".format(int(total_steps/1000))
                     save_dir.mkdir(parents=True, exist_ok=False)
 
                     # save models
@@ -389,18 +407,18 @@ def main():
 
                     # save metric
                     if best_metric is None:
-                        best_idx_epoch = idx_epoch
+                        best_steps = total_steps
                         best_metric = average_pesq_score
                     elif average_pesq_score > best_metric:
                         # great is better.
-                        best_idx_epoch = idx_epoch
+                        best_steps = total_steps
                         best_metric = average_pesq_score
                     else:
                         pass
 
                     metrics = {
                         "idx_epoch": idx_epoch,
-                        "best_idx_epoch": best_idx_epoch,
+                        "best_steps": best_steps,
                         "pesq_score": average_pesq_score,
                         "loss": average_loss,
                         "ae_loss": average_ae_loss,

examples/conv_tasnet/yaml/config.yaml

@@ -15,3 +15,11 @@ sub_blocks_kernel_size: 3
 norm_type: "gLN"
 causal: false
 mask_nonlinear: "relu"
+
+lr: 0.001
+lr_scheduler: "CosineAnnealingLR"
+lr_scheduler_kwargs:
+  T_max: 250000
+  eta_min: 0.00001
+
+eval_steps: 25000

toolbox/torchaudio/models/clean_unet/inference_clean_unet.py

@@ -79,6 +79,7 @@ class InferenceCleanUNet(object):
         # enhanced_audio shape: [channels, num_samples]
         return enhanced_audio
 
+
 def main():
     model_zip_file = project_path / "trained_models/clean-unet-aishell-18-epoch.zip"
     infer_mpnet = InferenceCleanUNet(model_zip_file)
@@ -100,5 +101,5 @@ def main():
     return
 
 
-if __name__ == '__main__':
+if __name__ == "__main__":
     main()

toolbox/torchaudio/models/conv_tasnet/configuration_conv_tasnet.py

@@ -27,6 +27,12 @@ class ConvTasNetConfig(PretrainedConfig):
                  causal: bool = False,
                  mask_nonlinear: str = "relu",
 
+                 lr: float = 1e-3,
+                 eval_steps: int = 25000,
+
+                 lr_scheduler: str = "CosineAnnealingLR",
+                 lr_scheduler_kwargs: dict = None,
+
                  **kwargs
                  ):
         super(ConvTasNetConfig, self).__init__(**kwargs)
@@ -47,6 +53,12 @@ class ConvTasNetConfig(PretrainedConfig):
         self.causal = causal
         self.mask_nonlinear = mask_nonlinear
 
+        self.lr = lr
+        self.eval_steps = eval_steps
+
+        self.lr_scheduler = lr_scheduler
+        self.lr_scheduler_kwargs = lr_scheduler_kwargs or dict()
+
 
 if __name__ == "__main__":
     pass

toolbox/torchaudio/models/conv_tasnet/inference_conv_tasnet.py

@@ -0,0 +1,112 @@
+#!/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.conv_tasnet.configuration_conv_tasnet import ConvTasNetConfig
+from toolbox.torchaudio.models.conv_tasnet.modeling_conv_tasnet import ConvTasNetPretrainedModel, MODEL_FILE
+
+logger = logging.getLogger("toolbox")
+
+
+class InferenceConvTasNet(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 = ConvTasNetConfig.from_pretrained(
+            pretrained_model_name_or_path=model_path.as_posix(),
+        )
+        model = ConvTasNetPretrainedModel.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():
+            enhanced_audios = self.model.forward(noisy_audios)
+            # enhanced_audio shape: [batch_size, channels, num_samples]
+            # enhanced_audios = torch.squeeze(enhanced_audios, dim=1)
+
+        enhanced_audio = enhanced_audios[0]
+
+        # enhanced_audio shape: [channels, num_samples]
+        return enhanced_audio
+
+
+def main():
+    model_zip_file = project_path / "trained_models/conv-tasnet-dns3-575k-steps.zip"
+    infer_conv_tasnet = InferenceConvTasNet(model_zip_file)
+
+    sample_rate = 8000
+    noisy_audio_file = project_path / "data/examples/ai_agent/dfaaf264-b5e3-4ca2-b5cb-5b6d637d962d_section_1.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_conv_tasnet.enhancement_by_tensor(noisy_audio)
+    time_cost = time.time() - begin
+    print(f"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()