update

examples/spectrum_unet_irm_aishell/run.sh

@@ -12,7 +12,7 @@ sh run.sh --stage 1 --stop_stage 3 --system_version centos --file_folder_name fi
 --noise_dir "/data/tianxing/HuggingDatasets/nx_noise/data/noise" \
 --speech_dir "/data/tianxing/HuggingDatasets/aishell/data_aishell/wav/train"
 
-sh run.sh --stage 2 --stop_stage 2 --system_version centos --file_folder_name file_dir \
+sh run.sh --stage 3 --stop_stage 3 --system_version centos --file_folder_name file_dir \
 --noise_dir "/data/tianxing/HuggingDatasets/nx_noise/data/noise" \
 --speech_dir "/data/tianxing/HuggingDatasets/aishell/data_aishell/wav/train"
 

examples/spectrum_unet_irm_aishell/step_2_train_model.py

@@ -295,10 +295,13 @@ def main():
             if torch.any(torch.isnan(lsnr_prediction)) or torch.any(torch.isinf(lsnr_prediction)):
                 raise AssertionError("nan or inf in lsnr_prediction")
             irm_loss = irm_mse_loss.forward(speech_irm_prediction, speech_irm_target)
+            lsnr_prediction = (lsnr_prediction - config.lsnr_min) / (config.lsnr_max - config.lsnr_min)
+            if torch.max(lsnr_prediction) > 1 or torch.min(lsnr_prediction) < 0:
+                raise AssertionError(f"expected lsnr_prediction between 0 and 1.")
             snr_loss = snr_mse_loss.forward(lsnr_prediction, snr_db_target)
             if torch.any(torch.isnan(snr_loss)) or torch.any(torch.isinf(snr_loss)):
                 raise AssertionError("nan or inf in snr_loss")
-            loss = irm_loss + 0.1 * snr_loss
+            loss = irm_loss + 1.0 * snr_loss
             # loss = irm_loss
 
             total_loss += loss.item()
@@ -336,8 +339,11 @@ def main():
                 if torch.any(torch.isnan(lsnr_prediction)) or torch.any(torch.isinf(lsnr_prediction)):
                     raise AssertionError("nan or inf in lsnr_prediction")
                 irm_loss = irm_mse_loss.forward(speech_irm_prediction, speech_irm_target)
+                lsnr_prediction = (lsnr_prediction - config.lsnr_min) / (config.lsnr_max - config.lsnr_min)
+                if torch.max(lsnr_prediction) > 1 or torch.min(lsnr_prediction) < 0:
+                    raise AssertionError(f"expected lsnr_prediction between 0 and 1.")
                 snr_loss = snr_mse_loss.forward(lsnr_prediction, snr_db_target)
-                loss = irm_loss + 0.1 * snr_loss
+                loss = irm_loss + 1.0 * snr_loss
                 # loss = irm_loss
 
                 total_loss += loss.item()

examples/spectrum_unet_irm_aishell/yaml/config.yaml

@@ -33,3 +33,6 @@ decoder_emb_num_layers: 3
 decoder_emb_skip_op: "none"
 decoder_emb_linear_groups: 16
 decoder_emb_hidden_size: 256
+
+# runtime
+use_post_filter: true

examples/test.py

@@ -1,18 +0,0 @@
-#!/usr/bin/python3
-# -*- coding: utf-8 -*-
-import torch
-
-speech_spec = torch.tensor([0], dtype=torch.float32)
-noise_spec = torch.tensor([0], dtype=torch.float32)
-epsilon = 1e-8
-
-
-result = torch.log10(
-    speech_spec / (noise_spec + epsilon) + epsilon
-)
-
-print(result)
-
-
-if __name__ == '__main__':
-    pass

requirements-python-3-9-9.txt

@@ -8,3 +8,4 @@ openpyxl==3.1.5
 torch==2.5.1
 torchaudio==2.5.1
 overrides==7.7.0
+torch-pesq==0.1.2

requirements.txt

@@ -8,3 +8,4 @@ openpyxl==3.1.5
 torch==2.5.1
 torchaudio==2.5.1
 overrides==7.7.0
+torch-pesq

toolbox/torch/training/__init__.py

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

toolbox/torch/training/metrics/__init__.py

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

toolbox/torch/training/metrics/pesq.py

@@ -0,0 +1,108 @@
+#!/usr/bin/python3
+# -*- coding: utf-8 -*-
+from typing import Optional
+
+import torch
+from torch_pesq import PesqLoss
+
+
+class Pesq(object):
+    def __init__(self):
+        pass
+
+
+class CategoricalAccuracy(object):
+    def __init__(self, top_k: int = 1, tie_break: bool = False) -> None:
+        if top_k > 1 and tie_break:
+            raise AssertionError("Tie break in Categorical Accuracy "
+                                 "can be done only for maximum (top_k = 1)")
+        if top_k <= 0:
+            raise AssertionError("top_k passed to Categorical Accuracy must be > 0")
+        self._top_k = top_k
+        self._tie_break = tie_break
+        self.correct_count = 0.
+        self.total_count = 0.
+
+    def __call__(self,
+                 predictions: torch.Tensor,
+                 gold_labels: torch.Tensor,
+                 mask: Optional[torch.Tensor] = None):
+
+        # predictions, gold_labels, mask = self.unwrap_to_tensors(predictions, gold_labels, mask)
+
+        # Some sanity checks.
+        num_classes = predictions.size(-1)
+        if gold_labels.dim() != predictions.dim() - 1:
+            raise AssertionError("gold_labels must have dimension == predictions.size() - 1 but "
+                                 "found tensor of shape: {}".format(predictions.size()))
+        if (gold_labels >= num_classes).any():
+            raise AssertionError("A gold label passed to Categorical Accuracy contains an id >= {}, "
+                                 "the number of classes.".format(num_classes))
+
+        predictions = predictions.view((-1, num_classes))
+        gold_labels = gold_labels.view(-1).long()
+        if not self._tie_break:
+            # Top K indexes of the predictions (or fewer, if there aren't K of them).
+            # Special case topk == 1, because it's common and .max() is much faster than .topk().
+            if self._top_k == 1:
+                top_k = predictions.max(-1)[1].unsqueeze(-1)
+            else:
+                top_k = predictions.topk(min(self._top_k, predictions.shape[-1]), -1)[1]
+
+            # This is of shape (batch_size, ..., top_k).
+            correct = top_k.eq(gold_labels.unsqueeze(-1)).float()
+        else:
+            # prediction is correct if gold label falls on any of the max scores. distribute score by tie_counts
+            max_predictions = predictions.max(-1)[0]
+            max_predictions_mask = predictions.eq(max_predictions.unsqueeze(-1))
+            # max_predictions_mask is (rows X num_classes) and gold_labels is (batch_size)
+            # ith entry in gold_labels points to index (0-num_classes) for ith row in max_predictions
+            # For each row check if index pointed by gold_label is was 1 or not (among max scored classes)
+            correct = max_predictions_mask[torch.arange(gold_labels.numel()).long(), gold_labels].float()
+            tie_counts = max_predictions_mask.sum(-1)
+            correct /= tie_counts.float()
+            correct.unsqueeze_(-1)
+
+        if mask is not None:
+            correct *= mask.view(-1, 1).float()
+            self.total_count += mask.sum()
+        else:
+            self.total_count += gold_labels.numel()
+        self.correct_count += correct.sum()
+
+    def get_metric(self, reset: bool = False):
+        """
+        Returns
+        -------
+        The accumulated accuracy.
+        """
+        if self.total_count > 1e-12:
+            accuracy = float(self.correct_count) / float(self.total_count)
+        else:
+            accuracy = 0.0
+        if reset:
+            self.reset()
+        return {'accuracy': accuracy}
+
+    def reset(self):
+        self.correct_count = 0.0
+        self.total_count = 0.0
+
+
+def main():
+    pesq = PesqLoss(0.5,
+                    sample_rate=8000,
+                    )
+
+    reference = torch.randn(1, 44100)
+    degraded = torch.randn(1, 44100)
+
+    mos = pesq.mos(reference, degraded)
+    loss = pesq(reference, degraded)
+
+    print(mos, loss)
+    return
+
+
+if __name__ == '__main__':
+    main()

toolbox/torchaudio/models/spectrum_unet_irm/configuration_specturm_unet_irm.py

@@ -33,6 +33,7 @@ class SpectrumUnetIRMConfig(PretrainedConfig):
                  decoder_emb_linear_groups: int = 16,
                  decoder_emb_hidden_size: int = 256,
 
+                 use_post_filter: bool = False,
                  **kwargs
                  ):
         super(SpectrumUnetIRMConfig, self).__init__(**kwargs)
@@ -67,6 +68,9 @@ class SpectrumUnetIRMConfig(PretrainedConfig):
         self.decoder_emb_linear_groups = decoder_emb_linear_groups
         self.decoder_emb_hidden_size = decoder_emb_hidden_size
 
+        # runtime
+        self.use_post_filter = use_post_filter
+
 
 if __name__ == "__main__":
     pass

toolbox/torchaudio/models/spectrum_unet_irm/modeling_spectrum_unet_irm.py

@@ -570,6 +570,9 @@ class SpectrumUnetIRM(nn.Module):
         mask = torch.transpose(mask, dim0=2, dim1=1)
         lsnr = torch.transpose(lsnr, dim0=2, dim1=1)
 
+        if not self.training and self.config.use_post_filter:
+            mask = self.post_filter(mask)
+
         # mask shape: [batch_size, freq_dim, time_steps]
         # lsnr shape: [batch_size, 1, time_steps]
         return mask, lsnr

toolbox/torchaudio/models/spectrum_unet_irm/yaml/config.yaml

@@ -33,3 +33,6 @@ decoder_emb_num_layers: 3
 decoder_emb_skip_op: "none"
 decoder_emb_linear_groups: 16
 decoder_emb_hidden_size: 256
+
+# runtime
+use_post_filter: true