update

examples/dfnet/yaml/config-512.yaml

@@ -0,0 +1,74 @@
+model_name: "dfnet"
+
+# spec
+sample_rate: 8000
+n_fft: 512
+win_length: 200
+hop_length: 80
+
+spec_bins: 256
+
+# model
+conv_channels: 64
+conv_kernel_size_input:
+  - 3
+  - 3
+conv_kernel_size_inner:
+  - 1
+  - 3
+conv_lookahead: 0
+
+convt_kernel_size_inner:
+  - 1
+  - 3
+
+embedding_hidden_size: 256
+encoder_combine_op: "concat"
+
+encoder_emb_skip_op: "none"
+encoder_emb_linear_groups: 16
+encoder_emb_hidden_size: 256
+
+encoder_linear_groups: 32
+
+decoder_emb_num_layers: 3
+decoder_emb_skip_op: "none"
+decoder_emb_linear_groups: 16
+decoder_emb_hidden_size: 256
+
+df_decoder_hidden_size: 256
+df_num_layers: 2
+df_order: 5
+df_bins: 96
+df_gru_skip: "grouped_linear"
+df_decoder_linear_groups: 16
+df_pathway_kernel_size_t: 5
+df_lookahead: 2
+
+# lsnr
+n_frame: 3
+lsnr_max: 30
+lsnr_min: -15
+norm_tau: 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: 8
+batch_size: 32
+eval_steps: 10000
+
+# runtime
+use_post_filter: true

examples/dfnet/yaml/config.yaml

@@ -2,14 +2,14 @@ model_name: "dfnet"
 
 # spec
 sample_rate: 8000
-n_fft: 512
-win_length: 200
+n_fft: 160
+win_length: 160
 hop_length: 80
 
-spec_bins: 256
+spec_bins: 80
 
 # model
-conv_channels: 64
+conv_channels: 32
 conv_kernel_size_input:
   - 3
   - 3
@@ -22,26 +22,26 @@ convt_kernel_size_inner:
   - 1
   - 3
 
-embedding_hidden_size: 256
+embedding_hidden_size: 80
 encoder_combine_op: "concat"
 
 encoder_emb_skip_op: "none"
-encoder_emb_linear_groups: 16
-encoder_emb_hidden_size: 256
+encoder_emb_linear_groups: 5
+encoder_emb_hidden_size: 80
 
-encoder_linear_groups: 32
+encoder_linear_groups: 10
 
 decoder_emb_num_layers: 3
 decoder_emb_skip_op: "none"
-decoder_emb_linear_groups: 16
-decoder_emb_hidden_size: 256
+decoder_emb_linear_groups: 5
+decoder_emb_hidden_size: 80
 
-df_decoder_hidden_size: 256
+df_decoder_hidden_size: 80
 df_num_layers: 2
 df_order: 5
-df_bins: 96
+df_bins: 30
 df_gru_skip: "grouped_linear"
-df_decoder_linear_groups: 16
+df_decoder_linear_groups: 5
 df_pathway_kernel_size_t: 5
 df_lookahead: 2
 

main.py

@@ -61,49 +61,22 @@ def shell(cmd: str):
 
 
 denoise_engines = {
-    "mpnet-nx-speech-1-epoch": {
+    "mpnet-nx-speech": {
         "infer_cls": InferenceMPNet,
         "kwargs": {
-            "pretrained_model_path_or_zip_file": (
-                        project_path / "trained_models/mpnet-nx-speech-1-epoch.zip").as_posix()
-        }
-    },
-    "mpnet-nx-speech-20-epoch": {
-        "infer_cls": InferenceMPNet,
-        "kwargs": {
-            "pretrained_model_path_or_zip_file": (
-                    project_path / "trained_models/mpnet-nx-speech-20-epoch.zip").as_posix()
-        }
-    },
-    "mpnet-nx-speech-33-epoch-best": {
-        "infer_cls": InferenceMPNet,
-        "kwargs": {
-            "pretrained_model_path_or_zip_file": (
-                    project_path / "trained_models/mpnet-nx-speech-33-epoch-best.zip").as_posix()
-        }
-    },
-    "mpnet-aishell-1-epoch": {
-        "infer_cls": InferenceMPNet,
-        "kwargs": {
-            "pretrained_model_path_or_zip_file": (project_path / "trained_models/mpnet-aishell-1-epoch.zip").as_posix()
-        }
-    },
-    "mpnet-aishell-11-epoch": {
-        "infer_cls": InferenceMPNet,
-        "kwargs": {
-            "pretrained_model_path_or_zip_file": (project_path / "trained_models/mpnet-aishell-11-epoch.zip").as_posix()
+            "pretrained_model_path_or_zip_file": (project_path / "trained_models/mpnet-nx-speech.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()
+            "pretrained_model_path_or_zip_file": (project_path / "trained_models/frcrn-dns3.zip").as_posix()
         }
     },
 }
 
 
-@lru_cache(maxsize=3)
+@lru_cache(maxsize=1)
 def load_denoise_model(infer_cls, **kwargs):
     infer_engine = infer_cls(**kwargs)
 

toolbox/torchaudio/models/dfnet/modeling_dfnet.py

@@ -215,7 +215,10 @@ class GroupedLinear(nn.Module):
 
     def forward(self, x: torch.Tensor) -> torch.Tensor:
         # x: [..., I]
-        b, t, _ = x.shape
+        b, t, f = x.shape
+        if f != self.input_size:
+            raise AssertionError
+
         # new_shape = list(x.shape)[:-1] + [self.groups, self.ws]
         new_shape = (b, t, self.groups, self.ws)
         x = x.view(new_shape)
@@ -633,8 +636,9 @@ class DfDecoder(nn.Module):
             GroupedLinear(
                 input_size=self.df_decoder_hidden_size,
                 hidden_size=out_dim,
-                groups=config.df_decoder_linear_groups
-            ),
+                groups=config.df_decoder_linear_groups,
+                # groups = self.df_bins // 5,
+        ),
             nn.Tanh()
         )
         self.df_fc_a = nn.Sequential(

toolbox/torchaudio/modules/erb_bands.py

@@ -0,0 +1,124 @@
+#!/usr/bin/python3
+# -*- coding: utf-8 -*-
+import math
+
+import numpy as np
+
+
+def freq2erb(freq_hz: float) -> float:
+    """
+    https://www.cnblogs.com/LXP-Never/p/16011229.html
+    1 / (24.7 * 9.265) = 0.00436976
+    """
+    return 9.265 * math.log(freq_hz / (24.7 * 9.265) + 1)
+
+
+def erb2freq(n_erb: float) -> float:
+    return 24.7 * 9.265 * (math.exp(n_erb / 9.265) - 1)
+
+
+def get_erb_widths(sample_rate: int, fft_size: int, erb_bins: int, min_freq_bins_for_erb: int) -> np.ndarray:
+    """
+    https://github.com/Rikorose/DeepFilterNet/blob/main/libDF/src/lib.rs
+    :param sample_rate:
+    :param fft_size:
+    :param erb_bins: erb (Equivalent Rectangular Bandwidth) 等效矩形带宽的通道数.
+    :param min_freq_bins_for_erb: Minimum number of frequency bands per erb band
+    :return:
+    """
+    nyq_freq = sample_rate / 2.
+    freq_width: float = sample_rate / fft_size
+
+    min_erb: float = freq2erb(0.)
+    max_erb: float = freq2erb(nyq_freq)
+
+    erb = [0] * erb_bins
+    step = (max_erb - min_erb) / erb_bins
+
+    prev_freq_bin = 0
+    freq_over = 0
+    for i in range(1, erb_bins + 1):
+        f = erb2freq(min_erb + i * step)
+        freq_bin = int(round(f / freq_width))
+        freq_bins = freq_bin - prev_freq_bin - freq_over
+
+        if freq_bins < min_freq_bins_for_erb:
+            freq_over = min_freq_bins_for_erb - freq_bins
+            freq_bins = min_freq_bins_for_erb
+        else:
+            freq_over = 0
+        erb[i - 1] = freq_bins
+        prev_freq_bin = freq_bin
+
+    erb[erb_bins - 1] += 1
+    too_large = sum(erb) - (fft_size / 2 + 1)
+    if too_large > 0:
+        erb[erb_bins - 1] -= too_large
+    return np.array(erb, dtype=np.uint64)
+
+
+def get_erb_filter_bank(erb_widths: np.ndarray,
+                        sample_rate: int,
+                        normalized: bool = True,
+                        inverse: bool = False,
+                        ):
+    num_freq_bins = int(np.sum(erb_widths))
+    num_erb_bins = len(erb_widths)
+
+    fb: np.ndarray = np.zeros(shape=(num_freq_bins, num_erb_bins))
+
+    points = np.cumsum([0] + erb_widths.tolist()).astype(int)[:-1]
+    for i, (b, w) in enumerate(zip(points.tolist(), erb_widths.tolist())):
+        fb[b: b + w, i] = 1
+
+    if inverse:
+        fb = fb.T
+        if not normalized:
+            fb /= np.sum(fb, axis=1, keepdims=True)
+    else:
+        if normalized:
+            fb /= np.sum(fb, axis=0)
+    return fb
+
+
+def spec2erb(spec: np.ndarray, erb_fb: np.ndarray, db: bool = True):
+    """
+    ERB filterbank and transform to decibel scale.
+
+    :param spec: Spectrum of shape [B, C, T, F].
+    :param erb_fb: ERB filterbank array of shape [B] containing the ERB widths,
+            where B are the number of ERB bins.
+    :param db: Whether to transform the output into decibel scale. Defaults to `True`.
+    :return:
+    """
+    # complex spec to power spec. (real * real + image * image)
+    spec_ = np.abs(spec) ** 2
+
+    # spec to erb feature.
+    erb_feat = np.matmul(spec_, erb_fb)
+
+    if db:
+        erb_feat = 10 * np.log10(erb_feat + 1e-10)
+
+    erb_feat = np.array(erb_feat, dtype=np.float32)
+    return erb_feat
+
+
+def main():
+    erb_widths = get_erb_widths(
+        sample_rate=8000,
+        fft_size=512,
+        erb_bins=32,
+        min_freq_bins_for_erb=2,
+    )
+    erb_fb = get_erb_filter_bank(
+        erb_widths=erb_widths,
+        sample_rate=8000,
+    )
+    print(erb_fb.shape)
+
+    return
+
+
+if __name__ == "__main__":
+    main()