update

examples/conv_tasnet/step_2_train_model.py

@@ -256,8 +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
-            # 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 + 0.25 * mr_stft_loss
 
             denoise_audios_list_r = list(denoise_audios.detach().cpu().numpy())
             clean_audios_list_r = list(clean_audios.detach().cpu().numpy())
@@ -321,8 +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
-                # 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 + 0.25 * mr_stft_loss
 
                 denoise_audios_list_r = list(denoise_audios.detach().cpu().numpy())
                 clean_audios_list_r = list(clean_audios.detach().cpu().numpy())