update

examples/frcrn/run.sh

@@ -4,7 +4,7 @@
 
 
 sh run.sh --stage 2 --stop_stage 2 --system_version centos --file_folder_name file_dir --final_model_name frcrn-20-512-nx-dns3 \
---config_file "yaml/config-20.yaml" \
+--config_file "yaml/config-10.yaml" \
 --noise_dir "/data/tianxing/HuggingDatasets/nx_noise/data/noise" \
 --speech_dir "/data/tianxing/HuggingDatasets/nx_noise/data/speech/dns3-speech"
 

examples/frcrn/step_2_train_model.py

@@ -34,6 +34,7 @@ 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.metrics.pesq import run_pesq_score
 from toolbox.torchaudio.models.frcrn.configuration_frcrn import FRCRNConfig
 from toolbox.torchaudio.models.frcrn.modeling_frcrn import FRCRN, FRCRNPretrainedModel
@@ -220,6 +221,14 @@ def main():
         raise AssertionError(f"invalid lr_scheduler: {config.lr_scheduler}")
 
     neg_si_snr_loss_fn = NegativeSISNRLoss(reduction="mean").to(device)
+    mr_stft_loss_fn = MultiResolutionSTFTLoss(
+        fft_size_list=[256, 512, 1024],
+        win_size_list=[256, 512, 1024],
+        hop_size_list=[128, 256, 512],
+        factor_sc=1.5,
+        factor_mag=1.0,
+        reduction="mean"
+    ).to(device)
 
     # training loop
 
@@ -248,6 +257,7 @@ 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_batches = 0.
@@ -264,10 +274,11 @@ def main():
             est_spec, est_wav, est_mask = model.forward(noisy_audios)
             denoise_audios = est_wav
 
+            mr_stft_loss = mr_stft_loss_fn.forward(denoise_audios, clean_audios)
             neg_si_snr_loss = neg_si_snr_loss_fn.forward(denoise_audios, clean_audios)
             mask_loss = model.mask_loss_fn(est_mask, 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
             if torch.any(torch.isnan(loss)) or torch.any(torch.isinf(loss)):
                 logger.info(f"find nan or inf in loss.")
                 continue
@@ -284,12 +295,14 @@ 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_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)
 
@@ -298,6 +311,7 @@ def main():
                 "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,
             })
@@ -311,6 +325,7 @@ 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_batches = 0.
@@ -327,10 +342,11 @@ def main():
                         est_spec, est_wav, est_mask = model.forward(noisy_audios)
                         denoise_audios = est_wav
 
+                        mr_stft_loss = mr_stft_loss_fn.forward(denoise_audios, clean_audios)
                         neg_si_snr_loss = neg_si_snr_loss_fn.forward(denoise_audios, clean_audios)
                         mask_loss = model.mask_loss_fn(est_mask, 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
                         if torch.any(torch.isnan(loss)) or torch.any(torch.isinf(loss)):
                             logger.info(f"find nan or inf in loss.")
                             continue
@@ -347,6 +363,7 @@ def main():
 
                         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)
 
@@ -355,12 +372,14 @@ def main():
                             "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,
                         })
 
                     total_pesq_score = 0.
                     total_loss = 0.
+                    total_mr_stft_loss = 0.
                     total_neg_si_snr_loss = 0.
                     total_mask_loss = 0.
                     total_batches = 0.

toolbox/torchaudio/models/frcrn/modeling_frcrn.py

@@ -300,25 +300,38 @@ class FRCRNPretrainedModel(FRCRN):
 def main():
     # model = FRCRN(
     #     use_complex_networks=True,
-    #     model_complexity=45,
+    #     model_complexity=-1,
+    #     model_depth=10,
+    #     padding_mode="zeros",
+    #     nfft=128,
+    #     win_size=128,
+    #     hop_size=64,
+    #     win_type="hann",
+    # )
+
+    # model = FRCRN(
+    #     use_complex_networks=True,
+    #     model_complexity=-1,
     #     model_depth=14,
     #     padding_mode="zeros",
-    #     nfft=512,
-    #     win_size=400,
-    #     hop_size=200,
+    #     nfft=640,
+    #     win_size=640,
+    #     hop_size=320,
     #     win_type="hann",
     # )
+
     model = FRCRN(
         use_complex_networks=True,
-        model_complexity=45,
-        model_depth=14,
+        model_complexity=20,
+        model_depth=20,
         padding_mode="zeros",
-        nfft=640,
-        win_size=640,
-        hop_size=320,
+        nfft=512,
+        win_size=512,
+        hop_size=256,
         win_type="hann",
     )
-    mixture = torch.rand(size=(1, 8000), dtype=torch.float32)
+
+    mixture = torch.rand(size=(1, 32000), dtype=torch.float32)
 
     est_spec, est_wav, est_mask = model.forward(mixture)
     print(est_spec.shape)

toolbox/torchaudio/models/frcrn/unet.py

@@ -339,19 +339,8 @@ class UNet(nn.Module):
         return cmp_spec
 
 
-def main():
+def main10():
     # [batch_size, 1, freq_bins, time_steps, 2]
-    # x = torch.rand(size=(1, 1, 257, 2000, 2))
-    # unet = UNet(
-    #     in_channels=1,
-    #     model_complexity=45,
-    #     model_depth=20,
-    #     use_complex_networks=True
-    # )
-    # print(unet)
-    # result = unet.forward(x)
-    # print(result.shape)
-
     # x = torch.rand(size=(1, 1, 65, 2000, 2))
     x = torch.rand(size=(1, 1, 65, 200, 2))
     unet = UNet(
@@ -366,5 +355,20 @@ def main():
     return
 
 
+def main20():
+    # [batch_size, 1, freq_bins, time_steps, 2]
+    x = torch.rand(size=(1, 1, 257, 2000, 2))
+    unet = UNet(
+        in_channels=1,
+        model_complexity=45,
+        model_depth=20,
+        use_complex_networks=True
+    )
+    print(unet)
+    result = unet.forward(x)
+    print(result.shape)
+    return
+
+
 if __name__ == "__main__":
-    main()
+    main20()