update

examples/conv_tasnet/step_2_train_model.py

@@ -256,7 +256,8 @@ def main():
             neg_stoi_loss = neg_stoi_loss_fn.forward(denoise_audios, clean_audios)
             mr_stft_loss = mr_stft_loss_fn.forward(denoise_audios, clean_audios)
 
-            loss = 0.25 * ae_loss + 0.25 * neg_si_snr_loss + 0.25 * neg_stoi_loss
+            loss = 0.25 * ae_loss + 0.25 * neg_si_snr_loss
+            # loss = 0.25 * ae_loss + 0.25 * neg_si_snr_loss + 0.25 * neg_stoi_loss
 
             denoise_audios_list_r = list(denoise_audios.detach().cpu().numpy())
             clean_audios_list_r = list(clean_audios.detach().cpu().numpy())
@@ -320,7 +321,8 @@ def main():
                 neg_stoi_loss = neg_stoi_loss_fn.forward(denoise_audios, clean_audios)
                 mr_stft_loss = mr_stft_loss_fn.forward(denoise_audios, clean_audios)
 
-                loss = 0.25 * ae_loss + 0.25 * neg_si_snr_loss + 0.25 * neg_stoi_loss
+                loss = 0.25 * ae_loss + 0.25 * neg_si_snr_loss
+                # loss = 0.25 * ae_loss + 0.25 * neg_si_snr_loss + 0.25 * neg_stoi_loss
 
                 denoise_audios_list_r = list(denoise_audios.detach().cpu().numpy())
                 clean_audios_list_r = list(clean_audios.detach().cpu().numpy())

toolbox/torchaudio/losses/spectral.py

@@ -191,9 +191,11 @@ class LogSTFTMagnitudeLoss(torch.nn.Module):
 
     def __init__(self,
                  reduction: str = "mean",
+                 eps: float = 1e-8,
                  ):
         super(LogSTFTMagnitudeLoss, self).__init__()
         self.reduction = reduction
+        self.eps = eps
 
         if reduction not in ("sum", "mean"):
             raise AssertionError(f"param reduction must be sum or mean.")
@@ -207,7 +209,7 @@ class LogSTFTMagnitudeLoss(torch.nn.Module):
         :param clean_magnitude: Tensor, shape: [batch_size, time_steps, freq_bins]
         :return:
         """
-        return F.l1_loss(torch.log(denoise_magnitude), torch.log(clean_magnitude))
+        return F.l1_loss(torch.log(denoise_magnitude + self.eps), torch.log(clean_magnitude + self.eps))
 
 
 class STFTLoss(torch.nn.Module):