update

README.md

@@ -61,7 +61,7 @@ docker run -itd \
 --ipc=host \
 -v /data/tianxing/HuggingDatasets/nx_noise/data:/data/tianxing/HuggingDatasets/nx_noise/data \
 -v /data/tianxing/PycharmProjects/cc_vad:/data/tianxing/PycharmProjects/cc_vad \
-python:3.12
+python:3.12 /bin/bash
 
 
 查看GPU

examples/silero_vad_by_webrtcvad/run.sh

@@ -2,19 +2,9 @@
 
 : <<'END'
 
-sh run.sh --stage 2 --stop_stage 2 --system_version windows --file_folder_name file_dir --final_model_name dfnet-nx-speech \
---noise_dir "E:/Users/tianx/HuggingDatasets/nx_noise/data/noise" \
---speech_dir "E:/Users/tianx/HuggingDatasets/nx_noise/data/speech"
-
-sh run.sh --stage 2 --stop_stage 2 --system_version centos --file_folder_name file_dir --final_model_name dfnet2-nx-dns3 \
---noise_dir "/data/tianxing/HuggingDatasets/nx_noise/data/noise" \
---speech_dir "/data/tianxing/HuggingDatasets/nx_noise/data/speech/dns3-speech"
-
-sh run.sh --stage 2 --stop_stage 2 --system_version centos --file_folder_name file_dir --final_model_name dfnet2-nx2 \
---noise_dir "/data/tianxing/HuggingDatasets/nx_noise/data/noise/nx-noise" \
---speech_dir "/data/tianxing/HuggingDatasets/nx_noise/data/speech/nx-speech2"
-
-sh run.sh --stage 2 --stop_stage 2 --system_version centos --file_folder_name dfnet2-nx2-dns3 --final_model_name dfnet2-nx2-dns3 \
+bash run.sh --stage 1 --stop_stage 1 --system_version centos \
+--file_folder_name silero-vad-by-webrtcvad-nx2-dns3 \
+--final_model_name silero-vad-by-webrtcvad-nx2-dns3 \
 --noise_dir "/data/tianxing/HuggingDatasets/nx_noise/data/noise/" \
 --speech_dir "/data/tianxing/HuggingDatasets/nx_noise/data/speech/"
 

examples/silero_vad_by_webrtcvad/step_2_train_model.py

@@ -1,8 +1,5 @@
 #!/usr/bin/python3
 # -*- coding: utf-8 -*-
-"""
-https://github.com/Rikorose/DeepFilterNet
-"""
 import argparse
 import json
 import logging
@@ -13,9 +10,7 @@ from pathlib import Path
 import random
 import sys
 import shutil
-from typing import List
-
-from fontTools.varLib.plot import stops
+from typing import List, Tuple
 
 pwd = os.path.abspath(os.path.dirname(__file__))
 sys.path.append(os.path.join(pwd, "../../"))
@@ -27,12 +22,13 @@ from torch.nn import functional as F
 from torch.utils.data.dataloader import DataLoader
 from tqdm import tqdm
 
-from toolbox.torch.utils.data.dataset.denoise_jsonl_dataset import DenoiseJsonlDataset
-from toolbox.torchaudio.losses.snr import NegativeSISNRLoss
-from toolbox.torchaudio.losses.spectral import LSDLoss, MultiResolutionSTFTLoss
-from toolbox.torchaudio.metrics.pesq import run_pesq_score
-from toolbox.torchaudio.models.dfnet2.configuration_dfnet2 import DfNet2Config
-from toolbox.torchaudio.models.dfnet2.modeling_dfnet2 import DfNet2, DfNet2PretrainedModel
+from toolbox.torch.utils.data.dataset.vad_jsonl_dataset import VadJsonlDataset
+from toolbox.torchaudio.models.vad.silero_vad.configuration_silero_vad import SileroVadConfig
+from toolbox.torchaudio.models.vad.silero_vad.modeling_silero_vad import SileroVadModel, SileroVadPretrainedModel
+from toolbox.torchaudio.losses.vad_loss.base_vad_loss import BaseVadLoss
+from toolbox.torchaudio.losses.bce_loss import BCELoss
+from toolbox.torchaudio.losses.dice_loss import DiceLoss
+from toolbox.torchaudio.metrics.vad_metrics.vad_accuracy import VadAccuracy
 
 
 def get_args():
@@ -76,28 +72,23 @@ class CollateFunction(object):
         pass
 
     def __call__(self, batch: List[dict]):
-        clean_audios = list()
         noisy_audios = list()
-        snr_db_list = list()
+        batch_vad_segments = list()
 
         for sample in batch:
-            # noise_wave: torch.Tensor = sample["noise_wave"]
-            clean_audio: torch.Tensor = sample["speech_wave"]
-            noisy_audio: torch.Tensor = sample["mix_wave"]
-            # snr_db: float = sample["snr_db"]
+            noisy_wave: torch.Tensor = sample["noisy_wave"]
+            vad_segments: List[Tuple[float, float]] = sample["vad_segments"]
 
-            clean_audios.append(clean_audio)
-            noisy_audios.append(noisy_audio)
+            noisy_audios.append(noisy_wave)
+            batch_vad_segments.append(vad_segments)
 
-        clean_audios = torch.stack(clean_audios)
         noisy_audios = torch.stack(noisy_audios)
 
         # assert
-        if torch.any(torch.isnan(clean_audios)) or torch.any(torch.isinf(clean_audios)):
-            raise AssertionError("nan or inf in clean_audios")
         if torch.any(torch.isnan(noisy_audios)) or torch.any(torch.isinf(noisy_audios)):
             raise AssertionError("nan or inf in noisy_audios")
-        return clean_audios, noisy_audios
+
+        return noisy_audios, batch_vad_segments
 
 
 collate_fn = CollateFunction()
@@ -106,7 +97,7 @@ collate_fn = CollateFunction()
 def main():
     args = get_args()
 
-    config = DfNet2Config.from_pretrained(
+    config = SileroVadConfig.from_pretrained(
         pretrained_model_name_or_path=args.config_file,
     )
 
@@ -125,7 +116,7 @@ def main():
     logger.info(f"GPU available count: {n_gpu}; device: {device}")
 
     # datasets
-    train_dataset = DenoiseJsonlDataset(
+    train_dataset = VadJsonlDataset(
         jsonl_file=args.train_dataset,
         expected_sample_rate=config.sample_rate,
         max_wave_value=32768.0,
@@ -133,7 +124,7 @@ def main():
         max_snr_db=config.max_snr_db,
         # skip=225000,
     )
-    valid_dataset = DenoiseJsonlDataset(
+    valid_dataset = VadJsonlDataset(
         jsonl_file=args.valid_dataset,
         expected_sample_rate=config.sample_rate,
         max_wave_value=32768.0,
@@ -165,7 +156,7 @@ def main():
 
     # models
     logger.info(f"prepare models. config_file: {args.config_file}")
-    model = DfNet2PretrainedModel(config).to(device)
+    model = SileroVadPretrainedModel(config).to(device)
     model.to(device)
     model.train()
 
@@ -210,25 +201,17 @@ def main():
     else:
         raise AssertionError(f"invalid lr_scheduler: {config.lr_scheduler}")
 
-    neg_si_snr_loss_fn = NegativeSISNRLoss(reduction="mean").to(device)
-    mr_stft_loss_fn = MultiResolutionSTFTLoss(
-        fft_size_list=[256, 512, 1024],
-        win_size_list=[256, 512, 1024],
-        hop_size_list=[128, 256, 512],
-        factor_sc=1.5,
-        factor_mag=1.0,
-        reduction="mean"
-    ).to(device)
+    bce_loss_fn = BCELoss(reduction="mean").to(device)
+    dice_loss_fn = DiceLoss(reduction="mean").to(device)
+
+    vad_accuracy_metrics_fn = VadAccuracy(threshold=0.5)
 
     # training loop
 
     # state
-    average_pesq_score = 1000000000
     average_loss = 1000000000
-    average_mr_stft_loss = 1000000000
-    average_neg_si_snr_loss = 1000000000
-    average_mask_loss = 1000000000
-    average_lsnr_loss = 1000000000
+    average_bce_loss = 1000000000
+    average_dice_loss = 1000000000
 
     model_list = list()
     best_epoch_idx = None
@@ -246,13 +229,11 @@ def main():
 
         # train
         model.train()
+        vad_accuracy_metrics_fn.reset()
 
-        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_bce_loss = 0.
+        total_dice_loss = 0.
         total_batches = 0.
 
         progress_bar_train = tqdm(
@@ -260,28 +241,24 @@ def main():
             desc="Training; epoch-{}".format(epoch_idx),
         )
         for train_batch in train_data_loader:
-            clean_audios, noisy_audios = train_batch
-            clean_audios: torch.Tensor = clean_audios.to(device)
+            noisy_audios, batch_vad_segments = train_batch
             noisy_audios: torch.Tensor = noisy_audios.to(device)
+            # noisy_audios shape: [b, num_samples]
+            num_samples = noisy_audios.shape[-1]
+
+            predictions = model.forward(noisy_audios)
 
-            est_spec, est_wav, est_mask, lsnr = model.forward(noisy_audios)
-            # est_wav shape: [b, 1, n_samples]
-            est_wav = torch.squeeze(est_wav, dim=1)
-            # est_wav shape: [b, n_samples]
+            targets = BaseVadLoss.get_targets(predictions, batch_vad_segments, duration=num_samples / config.sample_rate)
 
-            mr_stft_loss = mr_stft_loss_fn.forward(est_wav, clean_audios)
-            neg_si_snr_loss = neg_si_snr_loss_fn.forward(est_wav, clean_audios)
-            mask_loss = model.mask_loss_fn(est_mask, clean_audios, noisy_audios)
-            lsnr_loss = model.lsnr_loss_fn(lsnr, clean_audios, noisy_audios)
+            bce_loss = bce_loss_fn.forward(predictions, targets)
+            dice_loss = dice_loss_fn.forward(predictions, targets)
 
-            loss = 1.0 * mr_stft_loss + 1.0 * neg_si_snr_loss + 1.0 * mask_loss + 0.01 * lsnr_loss
+            loss = 1.0 * bce_loss + 1.0 * dice_loss
             if torch.any(torch.isnan(loss)) or torch.any(torch.isinf(loss)):
                 logger.info(f"find nan or inf in loss. continue.")
                 continue
 
-            denoise_audios_list_r = list(est_wav.detach().cpu().numpy())
-            clean_audios_list_r = list(clean_audios.detach().cpu().numpy())
-            pesq_score = run_pesq_score(clean_audios_list_r, denoise_audios_list_r, sample_rate=config.sample_rate, mode="nb")
+            vad_accuracy_metrics_fn.__call__(predictions, targets)
 
             optimizer.zero_grad()
             loss.backward()
@@ -289,30 +266,25 @@ def main():
             optimizer.step()
             lr_scheduler.step()
 
-            total_pesq_score += pesq_score
             total_loss += loss.item()
-            total_mr_stft_loss += mr_stft_loss.item()
-            total_neg_si_snr_loss += neg_si_snr_loss.item()
-            total_mask_loss += mask_loss.item()
-            total_lsnr_loss += lsnr_loss.item()
+            total_bce_loss += bce_loss.item()
+            total_dice_loss += dice_loss.item()
             total_batches += 1
 
-            average_pesq_score = round(total_pesq_score / total_batches, 4)
             average_loss = round(total_loss / total_batches, 4)
-            average_mr_stft_loss = round(total_mr_stft_loss / total_batches, 4)
-            average_neg_si_snr_loss = round(total_neg_si_snr_loss / total_batches, 4)
-            average_mask_loss = round(total_mask_loss / total_batches, 4)
-            average_lsnr_loss = round(total_lsnr_loss / total_batches, 4)
+            average_bce_loss = round(total_bce_loss / total_batches, 4)
+            average_dice_loss = round(total_dice_loss / total_batches, 4)
+
+            metrics = vad_accuracy_metrics_fn.get_metric()
+            accuracy = metrics["accuracy"]
 
             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,
-                "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,
+                "average_bce_loss": average_bce_loss,
+                "average_dice_loss": average_dice_loss,
+                "accuracy": accuracy,
             })
 
             # evaluation
@@ -322,13 +294,11 @@ def main():
                     torch.cuda.empty_cache()
 
                     model.eval()
+                    vad_accuracy_metrics_fn.reset()
 
-                    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_bce_loss = 0.
+                    total_dice_loss = 0.
                     total_batches = 0.
 
                     progress_bar_train.close()
@@ -336,63 +306,52 @@ def main():
                         desc="Evaluation; steps-{}k".format(int(step_idx/1000)),
                     )
                     for eval_batch in valid_data_loader:
-                        clean_audios, noisy_audios = eval_batch
-                        clean_audios: torch.Tensor = clean_audios.to(device)
+                        noisy_audios, batch_vad_segments = train_batch
                         noisy_audios: torch.Tensor = noisy_audios.to(device)
+                        # noisy_audios shape: [b, num_samples]
+                        num_samples = noisy_audios.shape[-1]
+
+                        predictions = model.forward(noisy_audios)
 
-                        est_spec, est_wav, est_mask, lsnr = model.forward(noisy_audios)
-                        # est_wav shape: [b, 1, n_samples]
-                        est_wav = torch.squeeze(est_wav, dim=1)
-                        # est_wav shape: [b, n_samples]
+                        targets = BaseVadLoss.get_targets(predictions, batch_vad_segments, duration=num_samples / config.sample_rate)
 
-                        mr_stft_loss = mr_stft_loss_fn.forward(est_wav, clean_audios)
-                        neg_si_snr_loss = neg_si_snr_loss_fn.forward(est_wav, clean_audios)
-                        mask_loss = model.mask_loss_fn(est_mask, clean_audios, noisy_audios)
-                        lsnr_loss = model.lsnr_loss_fn(lsnr, clean_audios, noisy_audios)
+                        bce_loss = bce_loss_fn.forward(predictions, targets)
+                        dice_loss = dice_loss_fn.forward(predictions, targets)
 
-                        loss = 1.0 * mr_stft_loss + 1.0 * neg_si_snr_loss + 1.0 * mask_loss + 0.01 * lsnr_loss
+                        loss = 1.0 * bce_loss + 1.0 * dice_loss
                         if torch.any(torch.isnan(loss)) or torch.any(torch.isinf(loss)):
                             logger.info(f"find nan or inf in loss. continue.")
                             continue
 
-                        denoise_audios_list_r = list(est_wav.detach().cpu().numpy())
-                        clean_audios_list_r = list(clean_audios.detach().cpu().numpy())
-                        pesq_score = run_pesq_score(clean_audios_list_r, denoise_audios_list_r, sample_rate=config.sample_rate, mode="nb")
+                        vad_accuracy_metrics_fn.__call__(predictions, targets)
 
-                        total_pesq_score += pesq_score
                         total_loss += loss.item()
-                        total_mr_stft_loss += mr_stft_loss.item()
-                        total_neg_si_snr_loss += neg_si_snr_loss.item()
-                        total_mask_loss += mask_loss.item()
-                        total_lsnr_loss += lsnr_loss.item()
+                        total_bce_loss += bce_loss.item()
+                        total_dice_loss += dice_loss.item()
                         total_batches += 1
 
-                        average_pesq_score = round(total_pesq_score / total_batches, 4)
                         average_loss = round(total_loss / total_batches, 4)
-                        average_mr_stft_loss = round(total_mr_stft_loss / total_batches, 4)
-                        average_neg_si_snr_loss = round(total_neg_si_snr_loss / total_batches, 4)
-                        average_mask_loss = round(total_mask_loss / total_batches, 4)
-                        average_lsnr_loss = round(total_lsnr_loss / total_batches, 4)
+                        average_bce_loss = round(total_bce_loss / total_batches, 4)
+                        average_dice_loss = round(total_dice_loss / total_batches, 4)
+
+                        metrics = vad_accuracy_metrics_fn.get_metric()
+                        accuracy = metrics["accuracy"]
 
                         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,
-                            "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,
+                            "average_bce_loss": average_bce_loss,
+                            "average_dice_loss": average_dice_loss,
+                            "accuracy": accuracy,
                         })
 
                 model.train()
+                vad_accuracy_metrics_fn.reset()
 
-                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_bce_loss = 0.
+                total_dice_loss = 0.
                 total_batches = 0.
 
                 progress_bar_eval.close()
@@ -418,12 +377,12 @@ def main():
                 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:
+                    best_metric = accuracy
+                elif accuracy >= best_metric:
                     # great is better.
                     best_epoch_idx = epoch_idx
                     best_step_idx = step_idx
-                    best_metric = average_pesq_score
+                    best_metric = accuracy
                 else:
                     pass
 
@@ -431,12 +390,11 @@ def main():
                     "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,
+                    "bce_loss": average_bce_loss,
+                    "dice_loss": average_dice_loss,
+
+                    "accuracy": accuracy,
                 }
                 metrics_filename = save_dir / "metrics_epoch.json"
                 with open(metrics_filename, "w", encoding="utf-8") as f:

toolbox/torchaudio/losses/bce_loss.py

@@ -0,0 +1,44 @@
+#!/usr/bin/python3
+# -*- coding: utf-8 -*-
+from typing import List, Tuple
+
+import torch
+import torch.nn as nn
+
+from toolbox.torchaudio.losses.vad_loss.base_vad_loss import BaseVadLoss
+
+
+class BCELoss(BaseVadLoss):
+    """
+    Binary Cross-Entropy Loss, BCE Loss
+    """
+    def __init__(self,
+                 reduction: str = "mean",
+                 ):
+        super(BCELoss, self).__init__()
+        self.reduction = reduction
+
+        self.bce_loss_fn = nn.BCELoss(reduction=reduction)
+
+    def forward(self, inputs: torch.Tensor, targets: torch.Tensor):
+        """
+        :param inputs: torch.Tensor, shape: [b, t, 1]. vad prob, after sigmoid activation.
+        :param targets: shape as `inputs`.
+        :return:
+        """
+        loss = self.bce_loss_fn.forward(inputs, targets)
+        return loss
+
+
+def main():
+    inputs = torch.zeros(size=(1, 198, 1), dtype=torch.float32)
+
+    loss_fn = BCELoss()
+
+    loss = loss_fn.forward(inputs, inputs)
+    print(loss)
+    return
+
+
+if __name__ == "__main__":
+    main()

toolbox/torchaudio/losses/dice_loss.py

@@ -0,0 +1,61 @@
+#!/usr/bin/python3
+# -*- coding: utf-8 -*-
+from typing import List, Tuple
+
+import torch
+import torch.nn as nn
+
+
+class DiceLoss(nn.Module):
+    def __init__(self,
+                 reduction: str = "mean",
+                 eps: float = 1e-6,
+                 ):
+        super(DiceLoss, self).__init__()
+        self.reduction = reduction
+        self.eps = eps
+
+        if reduction not in ("sum", "mean"):
+            raise AssertionError(f"param reduction must be sum or mean.")
+
+    def forward(self, inputs: torch.Tensor, targets: torch.Tensor):
+        """
+        :param inputs: torch.Tensor, shape: [b, t, 1]. vad prob, after sigmoid activation.
+        :param targets: shape as `inputs`.
+        :return:
+        """
+        inputs_ = torch.squeeze(inputs, dim=-1)
+        targets_ = torch.squeeze(targets, dim=-1)
+        # shape: [b, t]
+
+        intersection = (inputs_ * targets_).sum(dim=-1)
+        union = (inputs_ + targets_).sum(dim=-1)
+        # shape: [b,]
+
+        dice = (2. * intersection + self.eps) / (union + self.eps)
+        # shape: [b,]
+
+        loss = 1. - dice
+        # shape: [b,]
+
+        if self.reduction == "mean":
+            loss = torch.mean(loss)
+        elif self.reduction == "sum":
+            loss = torch.sum(loss)
+        else:
+            raise AssertionError
+        return loss
+
+
+def main():
+    inputs = torch.zeros(size=(1, 198, 1), dtype=torch.float32)
+
+    loss_fn = DiceLoss()
+
+    loss = loss_fn.forward(inputs, inputs)
+    print(loss)
+    return
+
+
+if __name__ == "__main__":
+    main()

toolbox/torchaudio/models/vad/silero_vad/modeling_silero_vad.py

@@ -8,9 +8,13 @@ https://github.com/snakers4/silero-vad
 
 https://github.com/snakers4/silero-vad/blob/master/src/silero_vad/data/silero_vad.jit
 """
+import os
+from typing import Optional, Union
+
 import torch
 import torch.nn as nn
 
+from toolbox.torchaudio.configuration_utils import CONFIG_FILE
 from toolbox.torchaudio.models.vad.silero_vad.configuration_silero_vad import SileroVadConfig
 from toolbox.torchaudio.modules.conv_stft import ConvSTFT
 
@@ -134,6 +138,52 @@ class SileroVadModel(nn.Module):
         return x
 
 
+class SileroVadPretrainedModel(SileroVadModel):
+    def __init__(self,
+                 config: SileroVadConfig,
+                 ):
+        super(SileroVadPretrainedModel, self).__init__(
+            config=config,
+        )
+
+    @classmethod
+    def from_pretrained(cls, pretrained_model_name_or_path, **kwargs):
+        config = SileroVadConfig.from_pretrained(pretrained_model_name_or_path, **kwargs)
+
+        model = cls(config)
+
+        if os.path.isdir(pretrained_model_name_or_path):
+            ckpt_file = os.path.join(pretrained_model_name_or_path, MODEL_FILE)
+        else:
+            ckpt_file = pretrained_model_name_or_path
+
+        with open(ckpt_file, "rb") as f:
+            state_dict = torch.load(f, map_location="cpu", weights_only=True)
+        model.load_state_dict(state_dict, strict=True)
+        return model
+
+    def save_pretrained(self,
+                        save_directory: Union[str, os.PathLike],
+                        state_dict: Optional[dict] = None,
+                        ):
+
+        model = self
+
+        if state_dict is None:
+            state_dict = model.state_dict()
+
+        os.makedirs(save_directory, exist_ok=True)
+
+        # save state dict
+        model_file = os.path.join(save_directory, MODEL_FILE)
+        torch.save(state_dict, model_file)
+
+        # save config
+        config_file = os.path.join(save_directory, CONFIG_FILE)
+        self.config.to_yaml_file(config_file)
+        return save_directory
+
+
 def main():
     config = SileroVadConfig()
     model = SileroVadModel(config=config)