update

examples/conv_tasnet/step_2_train_model.py

@@ -235,6 +235,7 @@ def main():
         total_ae_loss = 0.
         total_neg_si_snr_loss = 0.
         total_neg_stoi_loss = 0.
+        total_mr_stft_loss = 0.
         total_batches = 0.
         progress_bar = tqdm(
             desc="Training; epoch-{}".format(idx_epoch),
@@ -255,8 +256,7 @@ 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 = mr_stft_loss
-            # loss = 0.25 * ae_loss + 0.25 * neg_si_snr_loss + 0.25 * neg_stoi_loss + 0.25 * mr_stft_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())
@@ -272,6 +272,7 @@ def main():
             total_ae_loss += ae_loss.item()
             total_neg_si_snr_loss += neg_si_snr_loss.item()
             total_neg_stoi_loss += neg_stoi_loss.item()
+            total_mr_stft_loss += mr_stft_loss.item()
             total_batches += 1
 
             average_pesq_score = round(total_pesq_score / total_batches, 4)
@@ -279,6 +280,7 @@ def main():
             average_ae_loss = round(total_ae_loss / total_batches, 4)
             average_neg_si_snr_loss = round(total_neg_si_snr_loss / total_batches, 4)
             average_neg_stoi_loss = round(total_neg_stoi_loss / total_batches, 4)
+            average_mr_stft_loss = round(total_mr_stft_loss / total_batches, 4)
 
             progress_bar.update(1)
             progress_bar.set_postfix({
@@ -287,6 +289,7 @@ def main():
                 "ae_loss": average_ae_loss,
                 "neg_si_snr_loss": average_neg_si_snr_loss,
                 "neg_stoi_loss": average_neg_stoi_loss,
+                "mr_stft_loss": average_mr_stft_loss,
             })
 
         # evaluation
@@ -317,7 +320,7 @@ 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 + 0.25 * mr_stft_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())
@@ -328,6 +331,7 @@ def main():
                 total_ae_loss += ae_loss.item()
                 total_neg_si_snr_loss += neg_si_snr_loss.item()
                 total_neg_stoi_loss += neg_stoi_loss.item()
+                total_mr_stft_loss += mr_stft_loss.item()
                 total_batches += 1
 
                 average_pesq_score = round(total_pesq_score / total_batches, 4)
@@ -335,6 +339,7 @@ def main():
                 average_ae_loss = round(total_ae_loss / total_batches, 4)
                 average_neg_si_snr_loss = round(total_neg_si_snr_loss / total_batches, 4)
                 average_neg_stoi_loss = round(total_neg_stoi_loss / total_batches, 4)
+                average_mr_stft_loss = round(total_mr_stft_loss / total_batches, 4)
 
                 progress_bar.update(1)
                 progress_bar.set_postfix({
@@ -343,6 +348,8 @@ def main():
                     "ae_loss": average_ae_loss,
                     "neg_si_snr_loss": average_neg_si_snr_loss,
                     "neg_stoi_loss": average_neg_stoi_loss,
+                    "mr_stft_loss": average_mr_stft_loss,
+
                 })
 
         # scheduler