update

examples/cnn_vad_by_webrtcvad/step_1_prepare_data.py

@@ -210,7 +210,7 @@ def main():
                 "random1": random1,
             }
             row = json.dumps(row, ensure_ascii=False)
-            if random2 < (1 / 300 / 1):
+            if random2 < (2 / 300):
                 fvalid.write(f"{row}\n")
             else:
                 ftrain.write(f"{row}\n")

examples/cnn_vad_by_webrtcvad/step_4_train_model.py

@@ -38,7 +38,7 @@ def get_args():
     parser.add_argument("--valid_dataset", default="valid-vad.jsonl", type=str)
 
     parser.add_argument("--num_serialized_models_to_keep", default=15, type=int)
-    parser.add_argument("--patience", default=30, type=int)
+    parser.add_argument("--patience", default=10, type=int)
     parser.add_argument("--serialization_dir", default="serialization_dir", type=str)
 
     parser.add_argument("--config_file", default="yaml/config.yaml", type=str)
@@ -74,22 +74,28 @@ class CollateFunction(object):
 
     def __call__(self, batch: List[dict]):
         noisy_audios = list()
+        clean_audios = list()
         batch_vad_segments = list()
 
         for sample in batch:
             noisy_wave: torch.Tensor = sample["noisy_wave"]
+            speech_wave: torch.Tensor = sample["speech_wave"]
             vad_segments: List[Tuple[float, float]] = sample["vad_segments"]
 
             noisy_audios.append(noisy_wave)
+            clean_audios.append(speech_wave)
             batch_vad_segments.append(vad_segments)
 
         noisy_audios = torch.stack(noisy_audios)
+        clean_audios = torch.stack(clean_audios)
 
         # assert
         if torch.any(torch.isnan(noisy_audios)) or torch.any(torch.isinf(noisy_audios)):
             raise AssertionError("nan or inf in noisy_audios")
+        if torch.any(torch.isnan(clean_audios)) or torch.any(torch.isinf(clean_audios)):
+            raise AssertionError("nan or inf in clean_audios")
 
-        return noisy_audios, batch_vad_segments
+        return noisy_audios, clean_audios, batch_vad_segments
 
 
 collate_fn = CollateFunction()
@@ -214,6 +220,7 @@ def main():
     average_loss = 1000000000
     average_bce_loss = 1000000000
     average_dice_loss = 1000000000
+    average_lsnr_loss = 1000000000
 
     accuracy = -1
     f1 = -1
@@ -242,6 +249,7 @@ def main():
         total_loss = 0.
         total_bce_loss = 0.
         total_dice_loss = 0.
+        total_lsnr_loss = 0.
         total_batches = 0.
 
         progress_bar_train = tqdm(
@@ -249,19 +257,21 @@ def main():
             desc="Training; epoch-{}".format(epoch_idx),
         )
         for train_batch in train_data_loader:
-            noisy_audios, batch_vad_segments = train_batch
+            noisy_audios, clean_audios, batch_vad_segments = train_batch
             noisy_audios: torch.Tensor = noisy_audios.to(device)
+            clean_audios: torch.Tensor = clean_audios.to(device)
             # noisy_audios shape: [b, num_samples]
             num_samples = noisy_audios.shape[-1]
 
-            logits, probs = model.forward(noisy_audios)
+            logits, probs, lsnr = model.forward(noisy_audios)
 
             targets = BaseVadLoss.get_targets(probs, batch_vad_segments, duration=num_samples / config.sample_rate)
 
             bce_loss = bce_loss_fn.forward(probs, targets)
             dice_loss = dice_loss_fn.forward(probs, targets)
+            lsnr_loss = model.lsnr_loss_fn(lsnr, clean_audios, noisy_audios)
 
-            loss = 1.0 * bce_loss + 1.0 * dice_loss
+            loss = 1.0 * bce_loss + 1.0 * dice_loss + 0.3 * lsnr_loss
             if torch.any(torch.isnan(loss)) or torch.any(torch.isinf(loss)):
                 logger.info(f"find nan or inf in loss. continue.")
                 continue
@@ -278,11 +288,13 @@ def main():
             total_loss += loss.item()
             total_bce_loss += bce_loss.item()
             total_dice_loss += dice_loss.item()
+            total_lsnr_loss += lsnr_loss.item()
             total_batches += 1
 
             average_loss = round(total_loss / total_batches, 4)
             average_bce_loss = round(total_bce_loss / total_batches, 4)
             average_dice_loss = round(total_dice_loss / total_batches, 4)
+            average_lsnr_loss = round(total_lsnr_loss / total_batches, 4)
 
             metrics = vad_accuracy_metrics_fn.get_metric()
             accuracy = metrics["accuracy"]
@@ -297,6 +309,7 @@ def main():
                 "loss": average_loss,
                 "bce_loss": average_bce_loss,
                 "dice_loss": average_dice_loss,
+                "lsnr_loss": average_lsnr_loss,
                 "accuracy": accuracy,
                 "f1": f1,
                 "precision": precision,
@@ -316,6 +329,7 @@ def main():
                     total_loss = 0.
                     total_bce_loss = 0.
                     total_dice_loss = 0.
+                    total_lsnr_loss = 0.
                     total_batches = 0.
 
                     progress_bar_train.close()
@@ -323,19 +337,21 @@ def main():
                         desc="Evaluation; steps-{}k".format(int(step_idx/1000)),
                     )
                     for eval_batch in valid_data_loader:
-                        noisy_audios, batch_vad_segments = train_batch
+                        noisy_audios, clean_audios, batch_vad_segments = eval_batch
                         noisy_audios: torch.Tensor = noisy_audios.to(device)
+                        clean_audios: torch.Tensor = clean_audios.to(device)
                         # noisy_audios shape: [b, num_samples]
                         num_samples = noisy_audios.shape[-1]
 
-                        logits, probs = model.forward(noisy_audios)
+                        logits, probs, lsnr = model.forward(noisy_audios)
 
                         targets = BaseVadLoss.get_targets(probs, batch_vad_segments, duration=num_samples / config.sample_rate)
 
                         bce_loss = bce_loss_fn.forward(probs, targets)
                         dice_loss = dice_loss_fn.forward(probs, targets)
+                        lsnr_loss = model.lsnr_loss_fn(lsnr, clean_audios, noisy_audios)
 
-                        loss = 1.0 * bce_loss + 1.0 * dice_loss
+                        loss = 1.0 * bce_loss + 1.0 * dice_loss + 0.3 * lsnr_loss
                         if torch.any(torch.isnan(loss)) or torch.any(torch.isinf(loss)):
                             logger.info(f"find nan or inf in loss. continue.")
                             continue
@@ -346,11 +362,13 @@ def main():
                         total_loss += loss.item()
                         total_bce_loss += bce_loss.item()
                         total_dice_loss += dice_loss.item()
+                        total_lsnr_loss += lsnr_loss.item()
                         total_batches += 1
 
                         average_loss = round(total_loss / total_batches, 4)
                         average_bce_loss = round(total_bce_loss / total_batches, 4)
                         average_dice_loss = round(total_dice_loss / total_batches, 4)
+                        average_lsnr_loss = round(total_lsnr_loss / total_batches, 4)
 
                         metrics = vad_accuracy_metrics_fn.get_metric()
                         accuracy = metrics["accuracy"]
@@ -365,6 +383,7 @@ def main():
                             "loss": average_loss,
                             "bce_loss": average_bce_loss,
                             "dice_loss": average_dice_loss,
+                            "lsnr_loss": average_lsnr_loss,
                             "accuracy": accuracy,
                             "f1": f1,
                             "precision": precision,
@@ -378,6 +397,7 @@ def main():
                 total_loss = 0.
                 total_bce_loss = 0.
                 total_dice_loss = 0.
+                total_lsnr_loss = 0.
                 total_batches = 0.
 
                 progress_bar_eval.close()
@@ -419,8 +439,12 @@ def main():
                     "loss": average_loss,
                     "bce_loss": average_bce_loss,
                     "dice_loss": average_dice_loss,
+                    "lsnr_loss": average_lsnr_loss,
 
                     "accuracy": accuracy,
+                    "f1": f1,
+                    "precision": precision,
+                    "recall": recall,
                 }
                 metrics_filename = save_dir / "metrics_epoch.json"
                 with open(metrics_filename, "w", encoding="utf-8") as f:

examples/cnn_vad_by_webrtcvad/yaml/config.yaml

@@ -1,4 +1,4 @@
-model_name: "fsmn_vad"
+model_name: "cnn_vad"
 
 # spec
 sample_rate: 8000
@@ -8,19 +8,36 @@ hop_size: 80
 win_type: hann
 
 # model
-fsmn_input_size: 257
-fsmn_input_affine_size: 140
-fsmn_hidden_size: 250
-fsmn_basic_block_layers: 4
-fsmn_basic_block_hidden_size: 128
-fsmn_basic_block_lorder: 20
-fsmn_basic_block_rorder: 0
-fsmn_basic_block_lstride: 1
-fsmn_basic_block_rstride: 0
-fsmn_output_affine_size: 140
-fsmn_output_size: 1
+conv2d_block_param_list:
+  - batch_norm: true
+    in_channels: 1
+    out_channels: 4
+    kernel_size: 3
+    padding: "same"
+    dilation: 3
+    activation: relu
+    dropout: 0.1
+  - in_channels: 4
+    out_channels: 4
+    kernel_size: 5
+    padding: "same"
+    dilation: 3
+    activation: relu
+    dropout: 0.1
+  - in_channels: 4
+    out_channels: 4
+    kernel_size: 3
+    padding: "same"
+    dilation: 2
+    activation: relu
+    dropout: 0.1
+encoder_output_size: 1028
 
-use_softmax: false
+# lsnr
+n_frame: 3
+min_local_snr_db: -15
+max_local_snr_db: 30
+norm_tau: 1.
 
 # data
 min_snr_db: -10

examples/fsmn_vad_by_webrtcvad/step_4_train_model.py

@@ -323,7 +323,7 @@ def main():
                         desc="Evaluation; steps-{}k".format(int(step_idx/1000)),
                     )
                     for eval_batch in valid_data_loader:
-                        noisy_audios, batch_vad_segments = train_batch
+                        noisy_audios, batch_vad_segments = eval_batch
                         noisy_audios: torch.Tensor = noisy_audios.to(device)
                         # noisy_audios shape: [b, num_samples]
                         num_samples = noisy_audios.shape[-1]

examples/silero_vad_by_webrtcvad/step_4_train_model.py

@@ -323,7 +323,7 @@ def main():
                         desc="Evaluation; steps-{}k".format(int(step_idx/1000)),
                     )
                     for eval_batch in valid_data_loader:
-                        noisy_audios, batch_vad_segments = train_batch
+                        noisy_audios, batch_vad_segments = eval_batch
                         noisy_audios: torch.Tensor = noisy_audios.to(device)
                         # noisy_audios shape: [b, num_samples]
                         num_samples = noisy_audios.shape[-1]

toolbox/torch/utils/data/dataset/vad_padding_jsonl_dataset.py

@@ -162,6 +162,7 @@ class VadPaddingJsonlDataset(IterableDataset):
 
         result = {
             "noisy_wave": noisy_wave,
+            "speech_wave": speech_wave,
             "vad_segments": vad_segments,
         }
         return result

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

@@ -46,7 +46,12 @@ class CNNVadConfig(PretrainedConfig):
                  win_type: str = "hann",
 
                  conv2d_block_param_list: list = None,
-                 classifier_input_size: int = 1028,
+                 encoder_output_size: int = 1028,
+
+                 n_frame: int = 3,
+                 min_local_snr_db: float = -15,
+                 max_local_snr_db: float = 30,
+                 norm_tau: float = 1.,
 
                  min_snr_db: float = -10,
                  max_snr_db: float = 20,
@@ -75,7 +80,13 @@ class CNNVadConfig(PretrainedConfig):
 
         # encoder
         self.conv2d_block_param_list = conv2d_block_param_list or DEFAULT_CONV2D_BLOCK_PARAM_LIST
-        self.classifier_input_size = classifier_input_size
+        self.encoder_output_size = encoder_output_size
+
+        # lsnr
+        self.n_frame = n_frame
+        self.min_local_snr_db = min_local_snr_db
+        self.max_local_snr_db = max_local_snr_db
+        self.norm_tau = norm_tau
 
         # data snr
         self.min_snr_db = min_snr_db

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

@@ -5,10 +5,12 @@ from typing import Dict, List, Optional, Tuple, Union
 
 import torch
 import torch.nn as nn
+from torch.nn import functional as F
 
 from toolbox.torchaudio.configuration_utils import CONFIG_FILE
 from toolbox.torchaudio.models.vad.cnn_vad.configuration_cnn_vad import CNNVadConfig
 from toolbox.torchaudio.modules.conv_stft import ConvSTFT
+from toolbox.torchaudio.modules.local_snr_target import LocalSnrTarget
 
 
 MODEL_FILE = "model.pt"
@@ -77,20 +79,26 @@ class Conv2dBlock(nn.Module):
 
 class CNNVadModel(nn.Module):
     def __init__(self,
+                 sample_rate: int,
                  nfft: int,
                  win_size: int,
                  hop_size: int,
                  win_type: str,
                  conv2d_block_param_list: List[dict],
-                 classifier_input_size: int,
+                 encoder_output_size: int,
+                 min_local_snr: float = -15,
+                 max_local_snr: float = 30,
                  ):
         super(CNNVadModel, self).__init__()
+        self.sample_rate = sample_rate
         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_input_size = classifier_input_size
+        self.encoder_output_size = encoder_output_size
+        self.min_local_snr = min_local_snr
+        self.max_local_snr = max_local_snr
 
         self.eps = 1e-12
 
@@ -102,6 +110,14 @@ class CNNVadModel(nn.Module):
             power=1,
             requires_grad=False
         )
+        self.complex_stft = ConvSTFT(
+            nfft=nfft,
+            win_size=win_size,
+            hop_size=hop_size,
+            win_type=win_type,
+            power=None,
+            requires_grad=False
+        )
 
         self.cnn_encoder_list = nn.ModuleList(modules=[
             Conv2dBlock(
@@ -117,14 +133,34 @@ class CNNVadModel(nn.Module):
             for param in conv2d_block_param_list
         ])
 
-        self.classifier = nn.Sequential(
-            nn.Linear(classifier_input_size, 32),
+        # vad
+        self.vad_fc = nn.Sequential(
+            nn.Linear(encoder_output_size, 32),
             nn.ReLU(),
             nn.Linear(32, 1),
         )
-
         self.sigmoid = nn.Sigmoid()
 
+        # lsnr
+        self.lsnr_fc = nn.Sequential(
+            nn.Linear(encoder_output_size, 1),
+            nn.Sigmoid()
+        )
+        self.lsnr_scale = self.max_local_snr - self.min_local_snr
+        self.lsnr_offset = self.min_local_snr
+
+        # lsnr
+        self.lsnr_fn = LocalSnrTarget(
+            sample_rate=self.sample_rate,
+            nfft=self.nfft,
+            win_size=self.win_size,
+            hop_size=self.hop_size,
+            n_frame=self.n_frame,
+            min_local_snr=self.min_local_snr,
+            max_local_snr=self.max_local_snr,
+            db=True,
+        )
+
     def forward(self, signal: torch.Tensor):
         if signal.dim() == 2:
             signal = torch.unsqueeze(signal, dim=1)
@@ -148,11 +184,46 @@ class CNNVadModel(nn.Module):
         x = torch.reshape(x, shape=(b, t, c*d))
         # x: [b, t, c*d]
 
-        logits = self.classifier.forward(x)
+        logits = self.vad_fc.forward(x)
         # logits shape: [b, t, 1]
         probs = self.sigmoid.forward(logits)
         # probs shape: [b, t, 1]
-        return logits, probs
+
+        lsnr = self.lsnr_fc.forward(x) * self.lsnr_scale + self.lsnr_offset
+        # lsnr shape: [b, t, 1]
+
+        return logits, probs, lsnr
+
+
+    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
+
+        if clean.dim() == 2:
+            clean = torch.unsqueeze(clean, dim=1)
+        if noise.dim() == 2:
+            noise = torch.unsqueeze(noise, dim=1)
+
+        stft_clean = self.complex_stft.forward(clean)
+        stft_noise = self.complex_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 CNNVadPretrainedModel(CNNVadModel):
@@ -165,7 +236,9 @@ class CNNVadPretrainedModel(CNNVadModel):
             hop_size=config.hop_size,
             win_type=config.win_type,
             conv2d_block_param_list=config.conv2d_block_param_list,
-            classifier_input_size=config.classifier_input_size,
+            encoder_output_size=config.encoder_output_size,
+            min_local_snr=config.min_local_snr,
+            max_local_snr=config.max_local_snr,
         )
         self.config = config
 
@@ -214,9 +287,10 @@ def main():
 
     noisy = torch.randn(size=(1, 16000), dtype=torch.float32)
 
-    logits, probs = model.forward(noisy)
-    print(f"probs: {probs}")
+    logits, probs, lsnr = model.forward(noisy)
+    print(f"logits.shape: {logits.shape}")
     print(f"probs.shape: {probs.shape}")
+    print(f"lsnr.shape: {lsnr.shape}")
 
     return
 

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

@@ -31,7 +31,13 @@ conv2d_block_param_list:
     dilation: 2
     activation: relu
     dropout: 0.1
-classifier_input_size: 1028
+encoder_output_size: 1028
+
+# lsnr
+n_frame: 3
+min_local_snr_db: -15
+max_local_snr_db: 30
+norm_tau: 1.
 
 # data
 min_snr_db: -10

toolbox/torchaudio/models/vad/fsmn_vad/inference_fsmn_vad.py

@@ -92,8 +92,15 @@ def get_args():
         # default=(project_path / "data/examples/hado/b556437e-c68b-4f6d-9eed-2977c29db887.wav").as_posix(),
         # default=(project_path / "data/examples/hado/eae93a33-8ee0-4d86-8f85-cac5116ae6ef.wav").as_posix(),
         # default=(project_path / "data/examples/speech/active_media_r_0ba69730-66a4-4ecd-8929-ef58f18f4612_2.wav").as_posix(),
-        default=(project_path / "data/examples/speech/active_media_r_2a2f472b-a0b8-4fd5-b1c4-1aedc5d2ce57_0.wav").as_posix(),
-        # default=r"D:\Users\tianx\HuggingDatasets\nx_noise\data\speech\nx-speech\en-SG\2025-06-17\active_media_r_0af6bd3a-9aef-4bef-935b-63abfb4d46d8_5.wav",
+        # default=(project_path / "data/examples/speech/active_media_r_2a2f472b-a0b8-4fd5-b1c4-1aedc5d2ce57_0.wav").as_posix(),
+        # default=r"D:\Users\tianx\HuggingDatasets\nx_noise\data\speech\nx-speech\en-SG\2025-05-29\active_media_r_1d4edd08-c6db-41a1-a349-7a22ac36f684_6.wav",
+        # default=r"D:\Users\tianx\HuggingDatasets\nx_noise\data\speech\nx-speech\en-SG\2025-05-29\active_media_r_04f6d842-488e-4e34-967b-2980fdd877c7_5.wav",
+        # default=r"D:\Users\tianx\HuggingDatasets\nx_noise\data\speech\nx-speech\en-SG\2025-05-29\active_media_r_7f6670aa-5600-44c0-9bce-77c1d2b739c7_8.wav",
+        # default=r"D:\Users\tianx\HuggingDatasets\nx_noise\data\speech\nx-speech\en-SG\2025-05-29\active_media_r_1187ff81-3a38-4b0b-846f-b81ad6540ce9_5.wav",
+        # default=r"D:\Users\tianx\HuggingDatasets\nx_noise\data\speech\nx-speech\en-SG\2025-05-29\active_media_r_e44bbfaa-f332-4c02-90a3-cc98505d9a1b_3.wav",
+        # default=r"D:\Users\tianx\HuggingDatasets\nx_noise\data\speech\nx-speech\en-SG\2025-05-29\active_media_r_f89cf1af-f556-42fd-9a42-6c9431002a12_11.wav",
+        # default=r"D:\Users\tianx\HuggingDatasets\nx_noise\data\speech\nx-speech\en-SG\2025-05-29\active_media_r_f89cf1af-f556-42fd-9a42-6c9431002a12_15.wav",
+        default=r"D:\Users\tianx\HuggingDatasets\nx_noise\data\speech\nx-speech\en-SG\2025-05-29\active_media_w_8b6e28e2-a238-4c8c-b2e3-426b1fca149b_6.wav",
         type=str,
     )
     args = parser.parse_args()

toolbox/torchaudio/models/vad/ten_vad/__init__.py

@@ -1,12 +0,0 @@
-#!/usr/bin/python3
-# -*- coding: utf-8 -*-
-"""
-https://huggingface.co/TEN-framework/ten-vad
-https://zhuanlan.zhihu.com/p/1906832842756976909
-https://github.com/TEN-framework/ten-vad
-
-"""
-
-
-if __name__ == "__main__":
-    pass

toolbox/torchaudio/modules/local_snr_target.py

@@ -0,0 +1,151 @@
+#!/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 = spec.pow(2).sum(-1).sum(-1)
+    # spec shape: [b, c, t]
+    spec = F.pad(spec, (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()