update

Dockerfile

@@ -12,6 +12,8 @@ RUN pip install --no-cache-dir --upgrade -r /code/requirements.txt
 
 RUN bash install.sh --stage 1 --stop_stage 2 --system_version centos
 
+RUN mkdir -p logs
+
 USER user
 
 ENV HOME=/home/user \

examples/fsmn_vad/step_1_prepare_data.py

@@ -1,156 +0,0 @@
-#!/usr/bin/python3
-# -*- coding: utf-8 -*-
-import argparse
-import json
-import os
-from pathlib import Path
-import random
-import sys
-
-pwd = os.path.abspath(os.path.dirname(__file__))
-sys.path.append(os.path.join(pwd, "../../"))
-
-import librosa
-import numpy as np
-from tqdm import tqdm
-
-
-def get_args():
-    parser = argparse.ArgumentParser()
-    parser.add_argument(
-        "--noise_dir",
-        default=r"E:\Users\tianx\HuggingDatasets\nx_noise\data\noise",
-        type=str
-    )
-    parser.add_argument(
-        "--speech_dir",
-        default=r"E:\programmer\asr_datasets\aishell\data_aishell\wav\train",
-        type=str
-    )
-
-    parser.add_argument("--train_dataset", default="train.jsonl", type=str)
-    parser.add_argument("--valid_dataset", default="valid.jsonl", type=str)
-
-    parser.add_argument("--duration", default=6.0, type=float)
-    parser.add_argument("--min_snr_db", default=-10, type=float)
-    parser.add_argument("--max_snr_db", default=20, type=float)
-
-    parser.add_argument("--target_sample_rate", default=8000, type=int)
-
-    parser.add_argument("--max_count", default=-1, type=int)
-
-    args = parser.parse_args()
-    return args
-
-
-def target_second_signal_generator(data_dir: str, duration: int = 6, sample_rate: int = 8000, max_epoch: int = 20000):
-    data_dir = Path(data_dir)
-    for epoch_idx in range(max_epoch):
-        for filename in data_dir.glob("**/*.wav"):
-            signal, _ = librosa.load(filename.as_posix(), sr=sample_rate)
-            raw_duration = librosa.get_duration(y=signal, sr=sample_rate)
-
-            if raw_duration < duration:
-                # print(f"duration less than {duration} s. skip filename: {filename.as_posix()}")
-                continue
-            if signal.ndim != 1:
-                raise AssertionError(f"expected ndim 1, instead of {signal.ndim}")
-
-            signal_length = len(signal)
-            win_size = int(duration * sample_rate)
-            for begin in range(0, signal_length - win_size, win_size):
-                if np.sum(signal[begin: begin+win_size]) == 0:
-                    continue
-                row = {
-                    "epoch_idx": epoch_idx,
-                    "filename": filename.as_posix(),
-                    "raw_duration": round(raw_duration, 4),
-                    "offset": round(begin / sample_rate, 4),
-                    "duration": round(duration, 4),
-                }
-                yield row
-
-
-def main():
-    args = get_args()
-
-    noise_dir = Path(args.noise_dir)
-    speech_dir = Path(args.speech_dir)
-
-    train_dataset = Path(args.train_dataset)
-    valid_dataset = Path(args.valid_dataset)
-    train_dataset.parent.mkdir(parents=True, exist_ok=True)
-    valid_dataset.parent.mkdir(parents=True, exist_ok=True)
-
-    noise_generator = target_second_signal_generator(
-        noise_dir.as_posix(),
-        duration=args.duration,
-        sample_rate=args.target_sample_rate,
-        max_epoch=100000,
-    )
-    speech_generator = target_second_signal_generator(
-        speech_dir.as_posix(),
-        duration=args.duration,
-        sample_rate=args.target_sample_rate,
-        max_epoch=1,
-    )
-
-    count = 0
-    process_bar = tqdm(desc="build dataset jsonl")
-    with open(args.train_dataset, "w", encoding="utf-8") as ftrain, open(args.valid_dataset, "w", encoding="utf-8") as fvalid:
-        for noise, speech in zip(noise_generator, speech_generator):
-            if count >= args.max_count > 0:
-                break
-
-            # row
-            noise_filename = noise["filename"]
-            noise_raw_duration = noise["raw_duration"]
-            noise_offset = noise["offset"]
-            noise_duration = noise["duration"]
-
-            speech_filename = speech["filename"]
-            speech_raw_duration = speech["raw_duration"]
-            speech_offset = speech["offset"]
-            speech_duration = speech["duration"]
-
-            # row
-            random1 = random.random()
-            random2 = random.random()
-
-            row = {
-                "count": count,
-
-                "noise_filename": noise_filename,
-                "noise_raw_duration": noise_raw_duration,
-                "noise_offset": noise_offset,
-                "noise_duration": noise_duration,
-
-                "speech_filename": speech_filename,
-                "speech_raw_duration": speech_raw_duration,
-                "speech_offset": speech_offset,
-                "speech_duration": speech_duration,
-
-                "snr_db": random.uniform(args.min_snr_db, args.max_snr_db),
-
-                "random1": random1,
-            }
-            row = json.dumps(row, ensure_ascii=False)
-            if random2 < (1 / 300 / 1):
-                fvalid.write(f"{row}\n")
-            else:
-                ftrain.write(f"{row}\n")
-
-            count += 1
-            duration_seconds = count * args.duration
-            duration_hours = duration_seconds / 3600
-
-            process_bar.update(n=1)
-            process_bar.set_postfix({
-                "duration_hours": round(duration_hours, 4),
-            })
-
-    return
-
-
-if __name__ == "__main__":
-    main()

examples/fsmn_vad_by_webrtcvad/run.sh

@@ -2,7 +2,7 @@
 
 : <<'END'
 
-bash run.sh --stage 1 --stop_stage 1 --system_version centos \
+bash run.sh --stage 2 --stop_stage 3 --system_version centos \
 --file_folder_name fsmn-vad-by-webrtcvad-nx2-dns3 \
 --final_model_name fsmn-vad-by-webrtcvad-nx2-dns3 \
 --noise_dir "/data/tianxing/HuggingDatasets/nx_noise/data/noise/" \

toolbox/torchaudio/models/vad/cnn_vad/configuration_cnn_vad.py

@@ -0,0 +1,73 @@
+#!/usr/bin/python3
+# -*- coding: utf-8 -*-
+from typing import Tuple
+
+from toolbox.torchaudio.configuration_utils import PretrainedConfig
+
+
+class CNNVadConfig(PretrainedConfig):
+    def __init__(self,
+                 sample_rate: int = 8000,
+                 nfft: int = 512,
+                 win_size: int = 240,
+                 hop_size: int = 80,
+                 win_type: str = "hann",
+
+                 conv2d_block_param_list: list = None,
+                 classifier_hidden_size: int = 128,
+
+                 min_snr_db: float = -10,
+                 max_snr_db: float = 20,
+
+                 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,
+
+                 num_workers: int = 4,
+                 batch_size: int = 4,
+                 eval_steps: int = 25000,
+
+                 **kwargs
+                 ):
+        super(CNNVadConfig, self).__init__(**kwargs)
+        # transform
+        self.sample_rate = sample_rate
+        self.nfft = nfft
+        self.win_size = win_size
+        self.hop_size = hop_size
+        self.win_type = win_type
+
+        # encoder
+        self.conv2d_block_param_list = conv2d_block_param_list
+        self.classifier_hidden_size = classifier_hidden_size
+
+        # data snr
+        self.min_snr_db = min_snr_db
+        self.max_snr_db = max_snr_db
+
+        # train
+        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.num_workers = num_workers
+        self.batch_size = batch_size
+        self.eval_steps = eval_steps
+
+
+def main():
+    config = SileroVadConfig()
+    config.to_yaml_file("config.yaml")
+    return
+
+
+if __name__ == "__main__":
+    main()

toolbox/torchaudio/models/vad/cnn_vad/modeling_cnn_vad.py

@@ -1,6 +1,169 @@
 #!/usr/bin/python3
 # -*- coding: utf-8 -*-
+import os
+from typing import Dict, List, Optional, Tuple, 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
+
+
+MODEL_FILE = "model.pt"
+
+
+name2activation = {
+    "relu": nn.ReLU,
+}
+
+
+class Conv2dBlock(nn.Module):
+    def __init__(self,
+                 in_channels: int,
+                 out_channels: int,
+                 kernel_size: Union[int, Tuple[int, int]],
+                 stride: Tuple[int, int],
+                 padding: str = 0,
+                 dilation: int = 1,
+                 batch_norm: bool = False,
+                 activation: str = None,
+                 dropout: float = None,
+                 ):
+        super().__init__()
+        self.in_channels = in_channels
+        self.out_channels = out_channels
+        self.kernel_size: Tuple[int, int] = kernel_size if isinstance(kernel_size, tuple) else (kernel_size, kernel_size)
+
+        if batch_norm:
+            self.batch_norm = nn.BatchNorm2d(in_channels)
+        else:
+            self.batch_norm = None
+
+        self.conv = nn.Conv2d(
+            in_channels,
+            out_channels,
+            kernel_size=kernel_size,
+            stride=stride,
+            padding=(padding,),
+            dilation=(dilation,),
+        )
+
+        if activation is None:
+            self.activation = None
+        else:
+            self.activation = name2activation[activation]()
+
+        if dropout is not None:
+            self.dropout = nn.Dropout(p=dropout)
+        else:
+            self.dropout = None
+
+    def forward(self, x: torch.Tensor):
+
+        if self.batch_norm is not None:
+            x = self.batch_norm(x)
+
+        x = self.conv(x)
+
+        if self.activation is not None:
+            x = self.activation(x)
+
+        if self.dropout is not None:
+            x = self.dropout(x)
+
+        return x
+
+
+class CNNVadModel(nn.Module):
+    def __init__(self,
+                 nfft: int,
+                 win_size: int,
+                 hop_size: int,
+                 win_type: str,
+                 conv2d_block_param_list: List[dict],
+                 classifier_hidden_size: int,
+                 ):
+        super(CNNVadModel, self).__init__()
+        self.nfft = nfft
+        self.win_size = win_size
+        self.hop_size = hop_size
+        self.win_type = win_type
+        self.conv2d_block_param_list = conv2d_block_param_list
+        self.classifier_hidden_size = classifier_hidden_size
+
+        self.eps = 1e-12
+
+        self.stft = ConvSTFT(
+            nfft=nfft,
+            win_size=win_size,
+            hop_size=hop_size,
+            win_type=win_type,
+            power=1,
+            requires_grad=False
+        )
+
+        self.cnn_encoder_list = nn.ModuleList(modules=[
+            Conv2dBlock(
+                batch_norm=param["batch_norm"],
+                in_channels=param["in_channels"],
+                out_channels=param["out_channels"],
+                kernel_size=param["kernel_size"],
+                stride=param["stride"],
+                dilation=param["dilation"],
+                activation=param["activation"],
+                dropout=param["dropout"],
+            )
+            for param in conv2d_block_param_list
+        ])
+
+        self.classifier = nn.Sequential(
+            nn.Linear(classifier_hidden_size, 32),
+            nn.ReLU(),
+            nn.Linear(32, 1),
+        )
+
+        self.sigmoid = nn.Sigmoid()
+
+    def forward(self, signal: torch.Tensor):
+        if signal.dim() == 2:
+            signal = torch.unsqueeze(signal, dim=1)
+        _, _, num_samples = signal.shape
+        # signal shape [b, 1, num_samples]
+
+        mags = self.stft.forward(signal)
+        # mags shape: [b, f, t]
+
+        x = torch.transpose(mags, dim0=1, dim1=2)
+        # x shape: [b, t, f]
+
+        x = self.linear.forward(x)
+        # x shape: [b, t, f']
+
+        x = self.encoder.forward(x)
+        # x shape: [b, t, f]
+
+        x, _ = self.lstm.forward(x)
+        logits = self.classifier.forward(x)
+        # logits shape: [b, t, 1]
+        probs = self.sigmoid.forward(logits)
+        # probs shape: [b, t, 1]
+        return logits, probs
+
+
+def main():
+    config = SileroVadConfig()
+    model = SileroVadModel(config=config)
+
+    noisy = torch.randn(size=(1, 16000), dtype=torch.float32)
+
+    logits, probs = model.forward(noisy)
+    print(f"logits: {probs}")
+    print(f"logits.shape: {logits.shape}")
+
+    return
 
 
 if __name__ == "__main__":
-    pass
+    main()

toolbox/torchaudio/models/vad/cnn_vad/yaml/config.yaml

@@ -0,0 +1,52 @@
+model_name: "cnn_vad"
+
+# spec
+sample_rate: 8000
+nfft: 512
+win_size: 240
+hop_size: 80
+win_type: hann
+
+# model
+conv2d_block_param_list:
+  - batch_norm: true
+    in_channels: 1
+    out_channels: 4
+    kernel_size: 3
+    stride: 1
+    dilation: 3
+    activation: relu
+    dropout: 0.1
+  - in_channels: 4
+    out_channels: 4
+    kernel_size: 5
+    stride: 2
+    dilation: 3
+    activation: relu
+    dropout: 0.1
+  - in_channels: 4
+    out_channels: 4
+    kernel_size: 3
+    stride: 1
+    dilation: 2
+    activation: relu
+    dropout: 0.1
+
+# data
+min_snr_db: -10
+max_snr_db: 20
+
+# 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
+
+num_workers: 4
+batch_size: 128
+eval_steps: 25000