update

examples/dfnet/run.sh

@@ -2,6 +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 dfnet-dns3 \
 --noise_dir "/data/tianxing/HuggingDatasets/nx_noise/data/noise/dns3-noise" \

examples/dfnet/step_1_prepare_data.py

@@ -104,7 +104,7 @@ def main():
     dataset = list()
 
     count = 0
-    process_bar = tqdm(desc="build dataset excel")
+    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:

examples/dfnet/step_2_train_model.py

@@ -25,6 +25,8 @@ 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.losses.irm import IRMLoss
+from toolbox.torchaudio.losses.snr import LocalSNRLoss
 from toolbox.torchaudio.metrics.pesq import run_pesq_score
 from toolbox.torchaudio.models.dfnet.configuration_dfnet import DfNetConfig
 from toolbox.torchaudio.models.dfnet.modeling_dfnet import DfNet, DfNetPretrainedModel
@@ -79,22 +81,20 @@ class CollateFunction(object):
             # 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"]
+            # snr_db: float = sample["snr_db"]
 
             clean_audios.append(clean_audio)
             noisy_audios.append(noisy_audio)
-            snr_db_list.append(snr_db)
 
         clean_audios = torch.stack(clean_audios)
         noisy_audios = torch.stack(noisy_audios)
-        snr_db_list = torch.tensor(snr_db_list)
 
         # 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, snr_db_list
+        return clean_audios, noisy_audios
 
 
 collate_fn = CollateFunction()
@@ -146,7 +146,7 @@ def main():
         num_workers=0 if platform.system() == "Windows" else os.cpu_count() // 2,
         collate_fn=collate_fn,
         pin_memory=False,
-        prefetch_factor=2,
+        prefetch_factor=None if platform.system() == "Windows" else 2,
     )
     valid_data_loader = DataLoader(
         dataset=valid_dataset,
@@ -157,7 +157,7 @@ def main():
         num_workers=0 if platform.system() == "Windows" else os.cpu_count() // 2,
         collate_fn=collate_fn,
         pin_memory=False,
-        prefetch_factor=2,
+        prefetch_factor=None if platform.system() == "Windows" else 2,
     )
 
     # models
@@ -222,7 +222,6 @@ def main():
         factor_mag=1.0,
         reduction="mean"
     ).to(device)
-    lsnr_loss_fn = nn.L1Loss(reduction="mean")
 
     # training loop
 
@@ -247,8 +246,10 @@ def main():
 
         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_batches = 0.
 
         progress_bar_train = tqdm(
@@ -256,20 +257,18 @@ def main():
             desc="Training; epoch-{}".format(epoch_idx),
         )
         for train_batch in train_data_loader:
-            clean_audios, noisy_audios, snr_db_list = train_batch
+            clean_audios, noisy_audios = train_batch
             clean_audios: torch.Tensor = clean_audios.to(device)
             noisy_audios: torch.Tensor = noisy_audios.to(device)
-            snr_db_list: torch.Tensor = snr_db_list.to(device)
 
             est_spec, est_wav, est_mask, lsnr = model.forward(noisy_audios)
 
-            print(f"est_mask.shape: {est_mask.shape}, est_mask.dtype: {est_mask.dtype}")
+            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)
-            # mr_stft_loss = mr_stft_loss_fn.forward(denoise_audios, clean_audios)
-            # neg_si_snr_loss = lsnr_loss_fn.forward(lsnr, snr_db_list)
+            lsnr_loss = model.lsnr_loss_fn(lsnr, clean_audios, noisy_audios)
 
-            loss = 1.0 * neg_si_snr_loss + 1.0 * mask_loss
+            loss = 1.0 * mr_stft_loss + 1.0 * neg_si_snr_loss + 1.0 * mask_loss + 1.0 * lsnr_loss
             if torch.any(torch.isnan(loss)) or torch.any(torch.isinf(loss)):
                 logger.info(f"find nan or inf in loss.")
                 continue
@@ -286,22 +285,28 @@ def main():
 
             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_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)
 
             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,
             })
 
             # evaluation
@@ -312,8 +317,10 @@ def main():
 
                     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_batches = 0.
 
                     progress_bar_train.close()
@@ -321,17 +328,18 @@ def main():
                         desc="Evaluation; steps-{}k".format(int(step_idx/1000)),
                     )
                     for eval_batch in valid_data_loader:
-                        clean_audios, noisy_audios, snr_db_list = eval_batch
+                        clean_audios, noisy_audios = eval_batch
                         clean_audios: torch.Tensor = clean_audios.to(device)
                         noisy_audios: torch.Tensor = noisy_audios.to(device)
-                        snr_db_list: torch.Tensor = snr_db_list.to(device)
 
                         est_spec, est_wav, est_mask, lsnr = model.forward(noisy_audios)
 
+                        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)
 
-                        loss = 1.0 * neg_si_snr_loss + 1.0 * mask_loss
+                        loss = 1.0 * mr_stft_loss + 1.0 * neg_si_snr_loss + 1.0 * mask_loss + 1.0 * lsnr_loss
                         if torch.any(torch.isnan(loss)) or torch.any(torch.isinf(loss)):
                             logger.info(f"find nan or inf in loss.")
                             continue
@@ -342,28 +350,36 @@ def main():
 
                         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_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)
 
                         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,
                         })
 
                     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_batches = 0.
 
                     progress_bar_eval.close()
@@ -393,7 +409,7 @@ def main():
                         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
@@ -407,8 +423,10 @@ def main():
                         "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,
                     }
                     metrics_filename = save_dir / "metrics_epoch.json"
                     with open(metrics_filename, "w", encoding="utf-8") as f:

examples/dfnet/yaml/config.yaml

@@ -31,10 +31,6 @@ encoder_emb_hidden_size: 256
 
 encoder_linear_groups: 32
 
-lsnr_max: 30
-lsnr_min: -15
-norm_tau: 1.
-
 decoder_emb_num_layers: 3
 decoder_emb_skip_op: "none"
 decoder_emb_linear_groups: 16
@@ -49,8 +45,15 @@ df_decoder_linear_groups: 16
 df_pathway_kernel_size_t: 5
 df_lookahead: 2
 
-# runtime
-use_post_filter: true
+# 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
@@ -63,9 +66,9 @@ max_epochs: 100
 clip_grad_norm: 10.0
 seed: 1234
 
-min_snr_db: -10
-max_snr_db: 20
-
 num_workers: 8
-batch_size: 32
+batch_size: 4
 eval_steps: 10000
+
+# runtime
+use_post_filter: true

examples/frcrn/step_1_prepare_data.py

@@ -104,7 +104,7 @@ def main():
     dataset = list()
 
     count = 0
-    process_bar = tqdm(desc="build dataset excel")
+    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):
             flag = random.random()

toolbox/torchaudio/losses/irm.py

@@ -93,6 +93,69 @@ class CIRMLoss(nn.Module):
         return loss
 
 
+class IRMLoss(nn.Module):
+    """
+    https://github.com/Rikorose/DeepFilterNet/blob/main/DeepFilterNet/df/loss.py#L25
+    """
+    def __init__(self,
+                 n_fft: int = 512,
+                 win_size: int = 512,
+                 hop_size: int = 256,
+                 center: bool = True,
+                 eps: float = 1e-8,
+                 reduction: str = "mean",
+                 ):
+        super(IRMLoss, self).__init__()
+        self.n_fft = n_fft
+        self.win_size = win_size
+        self.hop_size = hop_size
+        self.center = center
+        self.eps = eps
+        self.reduction = reduction
+
+        self.window = nn.Parameter(torch.hann_window(win_size), requires_grad=False)
+
+        if reduction not in ("sum", "mean"):
+            raise AssertionError(f"param reduction must be sum or mean.")
+
+    def forward(self, mask: torch.Tensor, clean: torch.Tensor, noisy: torch.Tensor):
+        if noisy.shape != clean.shape:
+            raise AssertionError("Input signals must have the same shape")
+        noise = noisy - clean
+
+        # clean_stft, noisy_stft shape: [b, f, t]
+        stft_clean = torch.stft(
+            clean,
+            n_fft=self.n_fft,
+            win_length=self.win_size,
+            hop_length=self.hop_size,
+            window=self.window,
+            center=self.center,
+            pad_mode="reflect",
+            normalized=False,
+            return_complex=True
+        )
+        stft_noise = torch.stft(
+            noise,
+            n_fft=self.n_fft,
+            win_length=self.win_size,
+            hop_length=self.hop_size,
+            window=self.window,
+            center=self.center,
+            pad_mode="reflect",
+            normalized=False,
+            return_complex=True
+        )
+
+        mag_clean = torch.abs(stft_clean)
+        mag_noise = torch.abs(stft_noise)
+
+        gth_irm_mask = (mag_clean / (mag_clean + mag_noise + self.eps)).clamp(0, 1)
+
+        loss = F.l1_loss(gth_irm_mask, mask, reduction=self.reduction)
+        return loss
+
+
 def main():
     batch_size = 2
     signal_length = 16000

toolbox/torchaudio/losses/snr.py

@@ -5,6 +5,9 @@ https://zhuanlan.zhihu.com/p/627039860
 """
 import torch
 import torch.nn as nn
+from torch.nn import functional as F
+
+from toolbox.torchaudio.modules.local_snr_target import LocalSnrTarget
 
 
 class NegativeSNRLoss(nn.Module):
@@ -83,6 +86,86 @@ class NegativeSISNRLoss(nn.Module):
         return -loss
 
 
+class LocalSNRLoss(nn.Module):
+    """
+    https://github.com/Rikorose/DeepFilterNet/blob/main/DeepFilterNet/df/modules.py#L816
+
+    """
+    def __init__(self,
+                 sample_rate: int = 8000,
+                 nfft: int = 512,
+                 win_size: int = 512,
+                 hop_size: int = 256,
+                 n_frame: int = 3,
+                 min_local_snr: int = -15,
+                 max_local_snr: int = 30,
+                 db: bool = True,
+                 factor: float = 1,
+                 reduction: str = "mean",
+                 eps: float = 1e-8,
+                 ):
+        super(LocalSNRLoss, self).__init__()
+        self.sample_rate = sample_rate
+        self.nfft = nfft
+        self.win_size = win_size
+        self.hop_size = hop_size
+
+        self.factor = factor
+        self.reduction = reduction
+        self.eps = eps
+
+        self.lsnr_fn = LocalSnrTarget(
+            sample_rate=sample_rate,
+            nfft=nfft,
+            win_size=win_size,
+            hop_size=hop_size,
+            n_frame=n_frame,
+            min_local_snr=min_local_snr,
+            max_local_snr=max_local_snr,
+            db=db,
+        )
+
+        self.window = nn.Parameter(torch.hann_window(win_size), requires_grad=False)
+
+    def forward(self, lsnr: torch.Tensor, clean: torch.Tensor, noisy: torch.Tensor):
+        if clean.shape != noisy.shape:
+            raise AssertionError("Input signals must have the same shape")
+        noise = noisy - clean
+
+        stft_clean = torch.stft(
+            clean,
+            n_fft=self.nfft,
+            win_length=self.win_size,
+            hop_length=self.hop_size,
+            window=self.window,
+            center=self.center,
+            pad_mode="reflect",
+            normalized=False,
+            return_complex=True
+        )
+        stft_noise = torch.stft(
+            noise,
+            n_fft=self.nfft,
+            win_length=self.win_size,
+            hop_length=self.hop_size,
+            window=self.window,
+            center=self.center,
+            pad_mode="reflect",
+            normalized=False,
+            return_complex=True
+        )
+
+        # lsnr shape: [b, 1, t]
+        lsnr = lsnr.squeeze(1)
+        # lsnr shape: [b, t]
+
+        lsnr_gth = self.lsnr_fn.forward(stft_clean, stft_noise)
+        # lsnr_gth shape: [b, t]
+
+        loss = F.mse_loss(lsnr, lsnr_gth) * self.factor
+        return loss
+
+
 def main():
     batch_size = 2
     signal_length = 16000

toolbox/torchaudio/models/dfnet/configuration_dfnet.py

@@ -31,10 +31,6 @@ class DfNetConfig(PretrainedConfig):
 
                  encoder_linear_groups: int = 32,
 
-                 lsnr_max: int = 30,
-                 lsnr_min: int = -15,
-                 norm_tau: float = 1.,
-
                  decoder_emb_num_layers: int = 3,
                  decoder_emb_skip_op: str = "none",
                  decoder_emb_linear_groups: int = 16,
@@ -49,7 +45,13 @@ class DfNetConfig(PretrainedConfig):
                  df_pathway_kernel_size_t: int = 5,
                  df_lookahead: int = 2,
 
-                 use_post_filter: bool = False,
+                 n_frame: int = 3,
+                 max_local_snr: int = 30,
+                 min_local_snr: int = -15,
+                 norm_tau: float = 1.,
+
+                 min_snr_db: float = -10,
+                 max_snr_db: float = 20,
 
                  lr: float = 0.001,
                  lr_scheduler: str = "CosineAnnealingLR",
@@ -59,13 +61,12 @@ class DfNetConfig(PretrainedConfig):
                  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,
 
+                 use_post_filter: bool = False,
+
                  **kwargs
                  ):
         super(DfNetConfig, self).__init__(**kwargs)
@@ -97,10 +98,6 @@ class DfNetConfig(PretrainedConfig):
         self.encoder_linear_groups = encoder_linear_groups
         self.encoder_combine_op = encoder_combine_op
 
-        self.lsnr_max = lsnr_max
-        self.lsnr_min = lsnr_min
-        self.norm_tau = norm_tau
-
         # decoder
         self.decoder_emb_num_layers = decoder_emb_num_layers
         self.decoder_emb_skip_op = decoder_emb_skip_op
@@ -117,10 +114,17 @@ class DfNetConfig(PretrainedConfig):
         self.df_pathway_kernel_size_t = df_pathway_kernel_size_t
         self.df_lookahead = df_lookahead
 
-        # runtime
-        self.use_post_filter = use_post_filter
+        # lsnr
+        self.n_frame = n_frame
+        self.max_local_snr = max_local_snr
+        self.min_local_snr = min_local_snr
+        self.norm_tau = norm_tau
+
+        # 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()
@@ -129,13 +133,13 @@ class DfNetConfig(PretrainedConfig):
         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
 
+        # runtime
+        self.use_post_filter = use_post_filter
+
 
 if __name__ == "__main__":
     pass

toolbox/torchaudio/models/dfnet/conv_stft.py

@@ -8,6 +8,7 @@ import torch
 import torch.nn as nn
 import torch.nn.functional as F
 from scipy.signal import get_window
+from sympy.physics.units import power
 
 
 def init_kernels(nfft: int, win_size: int, hop_size: int, win_type: str = None, inverse=False):
@@ -40,7 +41,7 @@ class ConvSTFT(nn.Module):
                  win_size: int,
                  hop_size: int,
                  win_type: str = "hamming",
-                 feature_type: str = "real",
+                 power: int = None,
                  requires_grad: bool = False):
         super(ConvSTFT, self).__init__()
 
@@ -57,23 +58,29 @@ class ConvSTFT(nn.Module):
 
         self.stride = hop_size
         self.dim = self.nfft
-        self.feature_type = feature_type
+        self.power = power
 
     def forward(self, inputs: torch.Tensor):
         if inputs.dim() == 2:
             inputs = torch.unsqueeze(inputs, 1)
 
-        outputs = F.conv1d(inputs, self.weight, stride=self.stride)
+        matrix = F.conv1d(inputs, self.weight, stride=self.stride)
+        dim = self.dim // 2 + 1
+        real = matrix[:, :dim, :]
+        imag = matrix[:, dim:, :]
+        spec = torch.complex(real, imag)
 
-        if self.feature_type == "complex":
-            return outputs
-        else:
-            dim = self.dim // 2 + 1
-            real = outputs[:, :dim, :]
-            imag = outputs[:, dim:, :]
+        if self.power is None:
+            return spec
+        elif self.power == 1:
             mags = torch.sqrt(real**2 + imag**2)
-            phase = torch.atan2(imag, real)
-            return mags, phase
+            # phase = torch.atan2(imag, real)
+            return mags
+        elif self.power == 2:
+            power = real**2 + imag**2
+            return power
+        else:
+            raise AssertionError
 
 
 class ConviSTFT(nn.Module):
@@ -83,7 +90,6 @@ class ConviSTFT(nn.Module):
                  hop_size: int,
                  nfft: int = None,
                  win_type: str = "hamming",
-                 feature_type: str = "real",
                  requires_grad: bool = False):
         super(ConviSTFT, self).__init__()
         if nfft is None:
@@ -100,45 +106,41 @@ class ConviSTFT(nn.Module):
 
         self.stride = hop_size
         self.dim = self.nfft
-        self.feature_type = feature_type
 
         self.register_buffer("window", window)
         self.register_buffer("enframe", torch.eye(win_size)[:, None, :])
 
     def forward(self,
-                inputs: torch.Tensor,
-                phase: torch.Tensor = None):
+                inputs: torch.Tensor):
         """
-        :param inputs: torch.Tensor, shape: [b, n+2, t] (complex spec) or [b, n//2+1, t] (mags)
-        :param phase: torch.Tensor, shape: [b, n//2+1, t]
+        :param inputs: torch.Tensor, shape: [b, f, t]
         :return:
         """
-        if phase is not None:
-            real = inputs * torch.cos(phase)
-            imag = inputs * torch.sin(phase)
-            inputs = torch.cat([real, imag], 1)
-        outputs = F.conv_transpose1d(inputs, self.weight, stride=self.stride)
+        inputs = torch.view_as_real(inputs)
+        matrix = torch.concat(tensors=[inputs[..., 0], inputs[..., 1]], dim=1)
+
+        waveform = F.conv_transpose1d(matrix, self.weight, stride=self.stride)
 
         # this is from torch-stft: https://github.com/pseeth/torch-stft
-        t = self.window.repeat(1, 1, inputs.size(-1))**2
+        t = self.window.repeat(1, 1, matrix.size(-1))**2
         coff = F.conv_transpose1d(t, self.enframe, stride=self.stride)
-        outputs = outputs / (coff + 1e-8)
-        return outputs
+        waveform = waveform / (coff + 1e-8)
+        return waveform
 
 
 def main():
-    stft = ConvSTFT(win_size=512, hop_size=200, feature_type="complex")
-    istft = ConviSTFT(win_size=512, hop_size=200, feature_type="complex")
+    stft = ConvSTFT(nfft=512, win_size=512, hop_size=200, power=None)
+    istft = ConviSTFT(nfft=512, win_size=512, hop_size=200)
 
     mixture = torch.rand(size=(1, 8000*40), dtype=torch.float32)
 
     spec = stft.forward(mixture)
     # shape: [batch_size, freq_bins, time_steps]
-    print(spec.shape)
+    print(f"spec.shape: {spec.shape}, spec.dtype: {spec.dtype}")
 
     waveform = istft.forward(spec)
     # shape: [batch_size, channels, num_samples]
-    print(waveform.shape)
+    print(f"waveform.shape: {waveform.shape}, waveform.dtype: {waveform.dtype}")
 
     return
 

toolbox/torchaudio/models/dfnet/modeling_dfnet.py

@@ -13,6 +13,7 @@ import torchaudio
 from toolbox.torchaudio.configuration_utils import CONFIG_FILE
 from toolbox.torchaudio.models.dfnet.configuration_dfnet import DfNetConfig
 from toolbox.torchaudio.models.dfnet.conv_stft import ConvSTFT, ConviSTFT
+from toolbox.torchaudio.modules.local_snr_target import LocalSnrTarget
 
 
 MODEL_FILE = "model.pt"
@@ -415,8 +416,8 @@ class Encoder(nn.Module):
             nn.Linear(self.embedding_output_size, 1),
             nn.Sigmoid()
         )
-        self.lsnr_scale = config.lsnr_max - config.lsnr_min
-        self.lsnr_offset = config.lsnr_min
+        self.lsnr_scale = config.max_local_snr - config.min_local_snr
+        self.lsnr_offset = config.min_local_snr
 
     def forward(self,
                 feat_power: torch.Tensor,
@@ -789,7 +790,7 @@ class DfNet(nn.Module):
     def __init__(self, config: DfNetConfig):
         super(DfNet, self).__init__()
         self.config = config
-        self.eps = 1e-8
+        self.eps = 1e-12
 
         self.freq_bins = self.config.nfft // 2 + 1
 
@@ -803,7 +804,7 @@ class DfNet(nn.Module):
             win_size=config.win_size,
             hop_size=config.hop_size,
             win_type=config.win_type,
-            feature_type="complex",
+            power=None,
             requires_grad=False
         )
         self.istft = ConviSTFT(
@@ -811,7 +812,6 @@ class DfNet(nn.Module):
             win_size=config.win_size,
             hop_size=config.hop_size,
             win_type=config.win_type,
-            feature_type="complex",
             requires_grad=False
         )
 
@@ -828,98 +828,121 @@ class DfNet(nn.Module):
 
         self.mask = Mask(use_post_filter=config.use_post_filter)
 
-    def forward(self,
-                noisy: torch.Tensor,
-                ):
-        if noisy.dim() == 2:
-            noisy = torch.unsqueeze(noisy, dim=1)
-        _, _, n_samples = noisy.shape
+        self.lsnr_fn = LocalSnrTarget(
+            sample_rate=config.sample_rate,
+            nfft=config.nfft,
+            win_size=config.win_size,
+            hop_size=config.hop_size,
+            n_frame=config.n_frame,
+            min_local_snr=config.min_local_snr,
+            max_local_snr=config.max_local_snr,
+            db=True,
+        )
+
+    def signal_prepare(self, signal: torch.Tensor) -> torch.Tensor:
+        if signal.dim() == 2:
+            signal = torch.unsqueeze(signal, dim=1)
+        _, _, n_samples = signal.shape
         remainder = (n_samples - self.win_size) % self.hop_size
         if remainder > 0:
             n_samples_pad = self.hop_size - remainder
-            noisy = F.pad(noisy, pad=(0, n_samples_pad), mode="constant", value=0)
+            signal = F.pad(signal, pad=(0, n_samples_pad), mode="constant", value=0)
+        return signal, n_samples
 
-        # [batch_size, freq_bins * 2, time_steps]
-        cmp_spec = self.stft.forward(noisy)
-        # [batch_size, 1, freq_bins * 2, time_steps]
-        cmp_spec = torch.unsqueeze(cmp_spec, 1)
+    def forward(self,
+                noisy: torch.Tensor,
+                ):
+        """
+        :param noisy:
+        :return:
+        est_spec: shape: [b, 257*2, t]
+        est_wav:  shape: [b, num_samples]
+        est_mask: shape: [b, 257, t]
+        lsnr:     shape: [b, 1, t]
+        """
+        noisy, n_samples = self.signal_prepare(noisy)
 
-        # [batch_size, 2, freq_bins, time_steps]
-        cmp_spec = torch.cat([
-            cmp_spec[:, :, :self.freq_bins, :],
-            cmp_spec[:, :, self.freq_bins:, :],
-        ], dim=1)
-        # n//2+1 -> n//2; 257 -> 256
+        # noisy shape: [b, num_samples_pad]
+        cmp_spec = self.stft.forward(noisy)
+        # cmp_spec shape: [b, f, t], torch.complex64
+        cmp_spec = torch.view_as_real(cmp_spec)
+        # cmp_spec shape: [b, f, t, 2]
+        cmp_spec = cmp_spec.permute(0, 3, 1, 2)
+        # cmp_spec shape: [b, 2, f, t]
         cmp_spec = cmp_spec[:, :, :-1, :]
+        # cmp_spec shape: [b, 2, spec_bins, t]
+        # n//2+1 -> n//2; 257 -> 256
 
         spec = torch.unsqueeze(cmp_spec, dim=4)
-        # [batch_size, 2, freq_bins, time_steps, 1]
+        # spec shape: [b, 2, spec_bins, t, 1]
         spec = spec.permute(0, 4, 3, 2, 1)
-        # spec shape: [batch_size, 1, time_steps, freq_bins, 2]
+        # spec shape: [b, 1, t, spec_bins, 2]
 
         feat_power = torch.sum(torch.square(spec), dim=-1)
-        # feat_power shape: [batch_size, 1, time_steps, spec_bins]
+        # feat_power shape: [b, 1, t, spec_bins]
 
         feat_spec = torch.transpose(cmp_spec, dim0=2, dim1=3)
-        # feat_spec shape: [batch_size, 2, time_steps, freq_bins]
+        # feat_spec shape: [b, 2, t, spec_bins]
         feat_spec = feat_spec[..., :self.df_decoder.df_bins]
-        # feat_spec shape: [batch_size, 2, time_steps, df_bins]
+        # feat_spec shape: [b, 2, t, df_bins]
 
         e0, e1, e2, e3, emb, c0, lsnr, h = self.encoder.forward(feat_power, feat_spec)
 
         mask = self.decoder.forward(emb, e3, e2, e1, e0)
-        # mask shape: [batch_size, 1, time_steps, spec_bins]
+        # mask shape: [b, 1, t, spec_bins]
         if torch.any(mask > 1) or torch.any(mask < 0):
             raise AssertionError
 
         spec_m = self.mask.forward(spec, mask)
 
-        # lsnr shape: [batch_size, time_steps, 1]
+        # lsnr shape: [b, t, 1]
         lsnr = torch.transpose(lsnr, dim0=2, dim1=1)
-        # lsnr shape: [batch_size, 1, time_steps]
+        # lsnr shape: [b, 1, t]
 
         df_coefs = self.df_decoder.forward(emb, c0)
         df_coefs = self.df_out_transform(df_coefs)
-        # df_coefs shape: [batch_size, df_order, time_steps, df_bins, 2]
+        # df_coefs shape: [b, df_order, t, df_bins, 2]
 
         spec_e = self.df_op.forward(spec.clone(), df_coefs)
-        # est_spec shape: [batch_size, 1, time_steps, spec_bins, 2]
+        # est_spec shape: [b, 1, t, spec_bins, 2]
 
         spec_e[..., self.df_decoder.df_bins:, :] = spec_m[..., self.df_decoder.df_bins:, :]
 
         spec_e = torch.squeeze(spec_e, dim=1)
         spec_e = spec_e.permute(0, 2, 1, 3)
-        # spec_e shape: [batch_size, spec_bins, time_steps, 2]
+        # spec_e shape: [b, spec_bins, t, 2]
 
         mask = torch.squeeze(mask, dim=1)
         mask = mask.permute(0, 2, 1)
-        # mask shape: [b, 256, t]
+        # mask shape: [b, spec_bins, t]
         est_mask = self.mask_transfer(mask)
-        # est_mask shape: [b, 257, t]
+        # est_mask shape: [b, f, t]
 
-        # spec_e shape: [b, 256, t, 2]
+        # spec_e shape: [b, spec_bins, t, 2]
         est_spec = self.spec_transfer(spec_e)
-        # est_spec shape: [b, 257*2, t]
+        # est_spec shape: [b, f, t], torch.complex64
+
         est_wav = self.istft.forward(est_spec)
         est_wav = torch.squeeze(est_wav, dim=1)
         est_wav = est_wav[:, :n_samples]
         # est_wav shape: [b, n_samples]
+
         return est_spec, est_wav, est_mask, lsnr
 
     def spec_transfer(self, spec_e: torch.Tensor) -> torch.Tensor:
-        # spec_e shape: [b, 256, t, 2]
+        # spec_e shape: [b, spec_bins, t, 2]
         b, _, t, _ = spec_e.shape
-        est_spec = torch.cat(tensors=[
-            torch.concat(tensors=[
+        est_spec = torch.complex(
+            real=torch.concat(tensors=[
                 spec_e[..., 0],
                 torch.zeros(size=(b, 1, t), dtype=spec_e.dtype).to(spec_e.device)
             ], dim=1),
-            torch.concat(tensors=[
+            imag=torch.concat(tensors=[
                 spec_e[..., 1],
                 torch.zeros(size=(b, 1, t), dtype=spec_e.dtype).to(spec_e.device)
             ], dim=1),
-        ], dim=1)
-        # est_spec shape: [b, 257*2, t]
+        )
+        # est_spec shape: [b, f, t]
         return est_spec
 
     def mask_transfer(self, mask: torch.Tensor) -> torch.Tensor:
@@ -934,29 +957,58 @@ class DfNet(nn.Module):
 
     def mask_loss_fn(self, est_mask: torch.Tensor, clean: torch.Tensor, noisy: torch.Tensor):
         """
-
-        :param est_mask: torch.Tensor, shape: [b, n+2, t]
+        :param est_mask: torch.Tensor, shape: [b, 257, t]
         :param clean:
         :param noisy:
         :return:
         """
-        clean_stft = self.stft(clean)
-        clean_re = clean_stft[:, :self.freq_bins, :]
-        clean_im = clean_stft[:, self.freq_bins:, :]
-        clean_power = clean_re ** 2 + clean_im ** 2
+        if noisy.shape != clean.shape:
+            raise AssertionError("Input signals must have the same shape")
+        noise = noisy - clean
+
+        clean, _ = self.signal_prepare(clean)
+        noise, _ = self.signal_prepare(noise)
 
-        noisy_stft = self.stft(noisy)
-        noisy_re = noisy_stft[:, :self.freq_bins, :]
-        noisy_im = noisy_stft[:, self.freq_bins:, :]
-        noisy_power = noisy_re ** 2 + noisy_im ** 2
+        stft_clean = self.stft.forward(clean)
+        mag_clean = torch.abs(stft_clean)
 
-        speech_irm = clean_power / (noisy_power + self.eps)
-        # speech_irm = torch.pow(speech_irm, self.irm_beta)
+        stft_noise = self.stft.forward(noise)
+        mag_noise = torch.abs(stft_noise)
 
-        loss = F.mse_loss(est_mask, speech_irm)
+        gth_irm_mask = (mag_clean / (mag_clean + mag_noise + self.eps)).clamp(0, 1)
+
+        loss = F.l1_loss(gth_irm_mask, est_mask, reduction="mean")
 
         return loss
 
+    def lsnr_loss_fn(self, lsnr: torch.Tensor, clean: torch.Tensor, noisy: torch.Tensor):
+        if noisy.shape != clean.shape:
+            raise AssertionError("Input signals must have the same shape")
+        noise = noisy - clean
+
+        clean, _ = self.signal_prepare(clean)
+        noise, _ = self.signal_prepare(noise)
+
+        stft_clean = self.stft.forward(clean)
+        stft_noise = self.stft.forward(noise)
+        # shape: [b, f, t]
+        stft_clean = torch.transpose(stft_clean, dim0=1, dim1=2)
+        stft_noise = torch.transpose(stft_noise, dim0=1, dim1=2)
+        # shape: [b, t, f]
+        stft_clean = torch.unsqueeze(stft_clean, dim=1)
+        stft_noise = torch.unsqueeze(stft_noise, dim=1)
+        # shape: [b, 1, t, f]
+
+        # lsnr shape: [b, 1, t]
+        lsnr = lsnr.squeeze(1)
+        # lsnr shape: [b, t]
+
+        lsnr_gth = self.lsnr_fn.forward(stft_clean, stft_noise)
+        # lsnr_gth shape: [b, t]
+
+        loss = F.mse_loss(lsnr, lsnr_gth)
+        return loss
+
 
 class DfNetPretrainedModel(DfNet):
     def __init__(self,
@@ -1011,8 +1063,12 @@ def main():
 
     noisy = torch.randn(size=(1, 16000), dtype=torch.float32)
 
-    output = model.forward(noisy)
-    print(output[1].shape)
+    est_spec, est_wav, est_mask, lsnr = model.forward(noisy)
+    print(f"est_spec.shape: {est_spec.shape}")
+    print(f"est_wav.shape: {est_wav.shape}")
+    print(f"est_mask.shape: {est_mask.shape}")
+    print(f"lsnr.shape: {lsnr.shape}")
+
     return
 
 

toolbox/torchaudio/modules/__init__.py

@@ -0,0 +1,6 @@
+#!/usr/bin/python3
+# -*- coding: utf-8 -*-
+
+
+if __name__ == "__main__":
+    pass

toolbox/torchaudio/modules/local_snr_target.py

@@ -0,0 +1,149 @@
+#!/usr/bin/python3
+# -*- coding: utf-8 -*-
+"""
+https://github.com/Rikorose/DeepFilterNet/blob/main/DeepFilterNet/df/modules.py#L816
+"""
+from typing import Tuple
+
+import torch
+import torch.nn as nn
+from torch.nn import functional as F
+import torchaudio
+
+
+def local_energy(spec: torch.Tensor, n_frame: int, device: torch.device) -> torch.Tensor:
+    if n_frame % 2 == 0:
+        n_frame += 1
+    n_frame_half = n_frame // 2
+
+    # spec shape: [b, c, t, f, 2]
+    spec = F.pad(spec.pow(2).sum(-1).sum(-1), (n_frame_half, n_frame_half, 0, 0))
+    # spec shape: [b, c, t-pad]
+
+    weight = torch.hann_window(n_frame, device=device, dtype=spec.dtype)
+    # w shape: [n_frame]
+
+    spec = spec.unfold(-1, size=n_frame, step=1) * weight
+    # x shape: [b, c, t, n_frame]
+
+    result = torch.sum(spec, dim=-1).div(n_frame)
+    # result shape: [b, c, t]
+    return result
+
+
+def local_snr(spec_clean: torch.Tensor,
+              spec_noise: torch.Tensor,
+              n_frame: int = 5,
+              db: bool = False,
+              eps: float = 1e-12,
+              ):
+    # [b, c, t, f]
+    spec_clean = torch.view_as_real(spec_clean)
+    spec_noise = torch.view_as_real(spec_noise)
+    # [b, c, t, f, 2]
+
+    energy_clean = local_energy(spec_clean, n_frame=n_frame, device=spec_clean.device)
+    energy_noise = local_energy(spec_noise, n_frame=n_frame, device=spec_noise.device)
+    # [b, c, t]
+
+    snr = energy_clean / energy_noise.clamp_min(eps)
+    # snr shape: [b, c, t]
+
+    if db:
+        snr = snr.clamp_min(eps).log10().mul(10)
+    return snr, energy_clean, energy_noise
+
+
+class LocalSnrTarget(nn.Module):
+    def __init__(self,
+                 sample_rate: int = 8000,
+                 nfft: int = 512,
+                 win_size: int = 512,
+                 hop_size: int = 256,
+
+                 n_frame: int = 3,
+
+                 min_local_snr: int = -15,
+                 max_local_snr: int = 30,
+
+                 db: bool = True,
+                 ):
+        super().__init__()
+        self.sample_rate = sample_rate
+        self.nfft = nfft
+        self.win_size = win_size
+        self.hop_size = hop_size
+
+        self.n_frame = n_frame
+
+        self.min_local_snr = min_local_snr
+        self.max_local_snr = max_local_snr
+
+        self.db = db
+
+    def forward(self,
+                spec_clean: torch.Tensor,
+                spec_noise: torch.Tensor,
+                ) -> torch.Tensor:
+        """
+
+        :param spec_clean: torch.complex, shape: [b, c, t, f]
+        :param spec_noise: torch.complex, shape: [b, c, t, f]
+        :return: lsnr, shape: [b, t]
+        """
+
+        lsnr, _, _ = local_snr(
+            spec_clean=spec_clean,
+            spec_noise=spec_noise,
+            n_frame=self.n_frame,
+            db=self.db,
+        )
+        # lsnr shape: [b, c, t]
+        lsnr = lsnr.clamp(self.min_local_snr, self.max_local_snr).squeeze(1)
+        # lsnr shape: [b, t]
+        return lsnr
+
+
+def main():
+    sample_rate = 8000
+    nfft = 512
+    win_size = 512
+    hop_size = 256
+    window_fn = "hamming"
+
+    transform = torchaudio.transforms.Spectrogram(
+        n_fft=nfft,
+        win_length=win_size,
+        hop_length=hop_size,
+        power=None,
+        window_fn=torch.hamming_window if window_fn == "hamming" else torch.hann_window,
+    )
+
+    noisy = torch.randn(size=(1, 16000), dtype=torch.float32)
+
+    spec = transform.forward(noisy)
+    spec = spec.permute(0, 2, 1)
+    spec = torch.unsqueeze(spec, dim=1)
+    print(f"spec.shape: {spec.shape}, spec.dtype: {spec.dtype}")
+
+    # [b, c, t, f]
+    # spec = torch.view_as_real(spec)
+    # [b, c, t, f, 2]
+
+    local = LocalSnrTarget(
+        sample_rate=sample_rate,
+        nfft=nfft,
+        win_size=win_size,
+        hop_size=hop_size,
+        n_frame=5,
+        min_local_snr=-15,
+        max_local_snr=30,
+        db=True,
+    )
+    lsnr_target = local.forward(spec, spec)
+    print(f"lsnr_target.shape: {lsnr_target.shape}")
+    return
+
+
+if __name__ == "__main__":
+    main()