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 3 --stop_stage 3 --system_version centos --file_folder_name file_dir \
+sh run.sh --stage 2 --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

@@ -140,8 +140,10 @@ class CollateFunction(object):
             # shape: [freq_dim, time_steps]
 
             snr_db: torch.Tensor = 10 * torch.log10(
-                speech_spec / (noise_spec + self.epsilon) + self.epsilon
+                speech_spec / (noise_spec + self.epsilon)
             )
+            snr_db = torch.clamp(snr_db, min=self.epsilon)
+
             snr_db_ = torch.unsqueeze(snr_db, dim=0)
             snr_db_ = torch.unsqueeze(snr_db_, dim=0)
             snr_db_ = self.make_unfold_snr_db(snr_db_, n_time_steps=3)
@@ -301,7 +303,7 @@ def main():
             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 + 1.0 * snr_loss
+            loss = irm_loss + 0.1 * snr_loss
             # loss = irm_loss
 
             total_loss += loss.item()
@@ -343,7 +345,7 @@ def main():
                 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 + 1.0 * snr_loss
+                loss = irm_loss + 0.1 * snr_loss
                 # loss = irm_loss
 
                 total_loss += loss.item()

requirements-python-3-9-9.txt

@@ -8,4 +8,6 @@ openpyxl==3.1.5
 torch==2.5.1
 torchaudio==2.5.1
 overrides==7.7.0
-torch-pesq==0.1.2
+torch-pesq
+torchmetrics
+torchmetrics[audio]

requirements.txt

@@ -8,4 +8,6 @@ openpyxl==3.1.5
 torch==2.5.1
 torchaudio==2.5.1
 overrides==7.7.0
-torch-pesq
+torch-pesq==0.1.2
+torchmetrics==1.6.1
+torchmetrics[audio]

toolbox/torch/training/metrics/stoi.py

@@ -0,0 +1,19 @@
+#!/usr/bin/python3
+# -*- coding: utf-8 -*-
+import torch
+from torchmetrics.functional.audio.stoi import short_time_objective_intelligibility
+
+
+# 假设 reference 和 degraded 是两个音频信号的张量
+reference = torch.randn(1, 16000)  # 参考信号
+degraded = torch.randn(1, 16000)   # 降质信号
+
+
+# 计算 STOI 分数
+stoi_score = short_time_objective_intelligibility(reference, degraded, fs=16000)
+
+print(f"STOI 分数: {stoi_score}")
+
+
+if __name__ == '__main__':
+    pass

toolbox/torchaudio/models/spectrum_unet_irm/modeling_spectrum_unet_irm.py

@@ -514,12 +514,14 @@ class Decoder(nn.Module):
 
 
 class SpectrumUnetIRM(nn.Module):
-    def __init__(self, config: SpectrumUnetIRMConfig):
+    def __init__(self, config: SpectrumUnetIRMConfig, eps: float = 1e-8):
         super(SpectrumUnetIRM, self).__init__()
         self.config = config
         self.encoder = Encoder(config)
         self.decoder = Decoder(config)
 
+        self.eps = eps
+
     def post_filter(self, mask: torch.Tensor, beta: float = 0.02) -> torch.Tensor:
         """
         总体上来说, 它会将 mask 中的值都调大一点. 可能是为了保留更多的声音以免损伤音质, 因为预测的 mask 肯定不是特别正确.