update

examples/spectrum_dfnet_aishell/step_2_train_model.py

@@ -313,22 +313,21 @@ def main():
             snr_db_target = snr_db.to(device)
 
             speech_spec_prediction, speech_irm_prediction, lsnr_prediction = model.forward(mix_complex_spec)
-            # if torch.any(torch.isnan(speech_spec_prediction)) or torch.any(torch.isinf(speech_spec_prediction)):
-            #     raise AssertionError("nan or inf in speech_spec_prediction")
+            if torch.any(torch.isnan(speech_spec_prediction)) or torch.any(torch.isinf(speech_spec_prediction)):
+                raise AssertionError("nan or inf in speech_spec_prediction")
             if torch.any(torch.isnan(speech_irm_prediction)) or torch.any(torch.isinf(speech_irm_prediction)):
                 raise AssertionError("nan or inf in speech_irm_prediction")
-            # if torch.any(torch.isnan(lsnr_prediction)) or torch.any(torch.isinf(lsnr_prediction)):
-            #     raise AssertionError("nan or inf in lsnr_prediction")
+            if torch.any(torch.isnan(lsnr_prediction)) or torch.any(torch.isinf(lsnr_prediction)):
+                raise AssertionError("nan or inf in lsnr_prediction")
 
-            # speech_loss = speech_mse_loss.forward(speech_spec_prediction, torch.view_as_real(speech_complex_spec))
+            speech_loss = speech_mse_loss.forward(speech_spec_prediction, torch.view_as_real(speech_complex_spec))
             irm_loss = irm_mse_loss.forward(speech_irm_prediction, speech_irm_target)
-            # snr_loss = snr_mse_loss.forward(lsnr_prediction, snr_db_target)
+            snr_loss = snr_mse_loss.forward(lsnr_prediction, snr_db_target)
 
             # if torch.any(torch.isnan(snr_loss)) or torch.any(torch.isinf(snr_loss)):
             #     raise AssertionError("nan or inf in snr_loss")
 
-            # loss = speech_loss + irm_loss + snr_loss
-            loss = irm_loss
+            loss = speech_loss + irm_loss + snr_loss
 
             total_loss += loss.item()
             total_examples += mix_complex_spec.size(0)
@@ -361,19 +360,18 @@ def main():
 
             with torch.no_grad():
                 speech_spec_prediction, speech_irm_prediction, lsnr_prediction = model.forward(mix_complex_spec)
-                # if torch.any(torch.isnan(speech_spec_prediction)) or torch.any(torch.isinf(speech_spec_prediction)):
-                #     raise AssertionError("nan or inf in speech_spec_prediction")
+                if torch.any(torch.isnan(speech_spec_prediction)) or torch.any(torch.isinf(speech_spec_prediction)):
+                    raise AssertionError("nan or inf in speech_spec_prediction")
                 if torch.any(torch.isnan(speech_irm_prediction)) or torch.any(torch.isinf(speech_irm_prediction)):
                     raise AssertionError("nan or inf in speech_irm_prediction")
-                # if torch.any(torch.isnan(lsnr_prediction)) or torch.any(torch.isinf(lsnr_prediction)):
-                #     raise AssertionError("nan or inf in lsnr_prediction")
+                if torch.any(torch.isnan(lsnr_prediction)) or torch.any(torch.isinf(lsnr_prediction)):
+                    raise AssertionError("nan or inf in lsnr_prediction")
 
-                # speech_loss = speech_mse_loss.forward(speech_spec_prediction, torch.view_as_real(speech_complex_spec))
+                speech_loss = speech_mse_loss.forward(speech_spec_prediction, torch.view_as_real(speech_complex_spec))
                 irm_loss = irm_mse_loss.forward(speech_irm_prediction, speech_irm_target)
-                # snr_loss = snr_mse_loss.forward(lsnr_prediction, snr_db_target)
+                snr_loss = snr_mse_loss.forward(lsnr_prediction, snr_db_target)
 
-                # loss = speech_loss + irm_loss + snr_loss
-                loss = irm_loss
+                loss = speech_loss + irm_loss + snr_loss
 
                 total_loss += loss.item()
             total_examples += mix_complex_spec.size(0)