update

examples/dfnet2/step_2_train_model.py

@@ -41,7 +41,7 @@ def get_args():
     parser.add_argument("--valid_dataset", default="valid.jsonl", type=str)
 
     parser.add_argument("--num_serialized_models_to_keep", default=15, type=int)
-    parser.add_argument("--patience", default=10, type=int)
+    parser.add_argument("--patience", default=30, type=int)
     parser.add_argument("--serialization_dir", default="serialization_dir", type=str)
 
     parser.add_argument("--config_file", default="config.yaml", type=str)
@@ -274,7 +274,7 @@ def main():
             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 * mr_stft_loss + 1.0 * neg_si_snr_loss + 1.0 * mask_loss + 0.3 * lsnr_loss
+            loss = 1.0 * mr_stft_loss + 1.0 * neg_si_snr_loss + 1.0 * mask_loss + 0.01 * lsnr_loss
             if torch.any(torch.isnan(loss)) or torch.any(torch.isinf(loss)):
                 logger.info(f"find nan or inf in loss.")
                 continue
@@ -350,7 +350,7 @@ def main():
                         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 * mr_stft_loss + 1.0 * neg_si_snr_loss + 1.0 * mask_loss + 0.3 * lsnr_loss
+                        loss = 1.0 * mr_stft_loss + 1.0 * neg_si_snr_loss + 1.0 * mask_loss + 0.01 * lsnr_loss
                         if torch.any(torch.isnan(loss)) or torch.any(torch.isinf(loss)):
                             logger.info(f"find nan or inf in loss.")
                             continue

examples/dtln/step_2_train_model.py

@@ -40,7 +40,7 @@ def get_args():
     parser.add_argument("--valid_dataset", default="valid.jsonl", type=str)
 
     parser.add_argument("--num_serialized_models_to_keep", default=15, type=int)
-    parser.add_argument("--patience", default=10, type=int)
+    parser.add_argument("--patience", default=30, type=int)
     parser.add_argument("--serialization_dir", default="serialization_dir", type=str)
 
     parser.add_argument("--config_file", default="config.yaml", type=str)

toolbox/torchaudio/losses/spectral.py

@@ -218,7 +218,7 @@ class LogSTFTMagnitudeLoss(torch.nn.Module):
         loss = F.l1_loss(torch.log(denoise_magnitude + self.eps), torch.log(clean_magnitude + self.eps))
 
         if torch.any(torch.isnan(loss)) or torch.any(torch.isinf(loss)):
-            raise AssertionError("SpectralConvergenceLoss, nan or inf in loss")
+            raise AssertionError("LogSTFTMagnitudeLoss, nan or inf in loss")
 
         return loss
 

toolbox/torchaudio/modules/local_snr_target.py

@@ -17,7 +17,9 @@ def local_energy(spec: torch.Tensor, n_frame: int, device: torch.device) -> torc
     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 = 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)