update

examples/spectrum_dfnet_aishell/step_2_train_model.py

@@ -313,21 +313,22 @@ 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")
 
-            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)
 
             # 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 = speech_loss + irm_loss + snr_loss
+            loss = irm_loss + snr_loss
 
             total_loss += loss.item()
             total_examples += mix_complex_spec.size(0)
@@ -360,18 +361,19 @@ 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")
 
-                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)
 
-                loss = speech_loss + irm_loss + snr_loss
+                # loss = speech_loss + irm_loss + snr_loss
+                loss = irm_loss + snr_loss
 
                 total_loss += loss.item()
             total_examples += mix_complex_spec.size(0)

toolbox/torchaudio/models/spectrum_dfnet/modeling_spectrum_dfnet.py

@@ -858,7 +858,8 @@ class SpectrumDfNet(nn.Module):
         mask = mask.permute(0, 2, 1)
         # mask shape: [batch_size, spec_bins, time_steps]
 
-        return spec_e, mask, lsnr
+        # return spec_e, mask, lsnr
+        return None, mask, lsnr
 
 
 class SpectrumDfNetPretrainedModel(SpectrumDfNet):