update

examples/spectrum_unet_irm_aishell/step_2_train_model.py

@@ -167,12 +167,12 @@ class CollateFunction(object):
         # snr_db shape: [batch_size, 1, time_steps]
 
         # assert
-        if torch.any(torch.isnan(mix_spec_list)):
-            raise AssertionError("nan in mix_spec Tensor")
-        if torch.any(torch.isnan(speech_irm_list)):
-            raise AssertionError("nan in speech_irm Tensor")
-        if torch.any(torch.isnan(snr_db_list)):
-            raise AssertionError("nan in snr_db Tensor")
+        if torch.any(torch.isnan(mix_spec_list)) or torch.any(torch.isinf(mix_spec_list)):
+            raise AssertionError("nan in mix_spec_list")
+        if torch.any(torch.isnan(speech_irm_list)) or torch.any(torch.isinf(speech_irm_list)):
+            raise AssertionError("nan in speech_irm_list")
+        if torch.any(torch.isnan(snr_db_list)) or torch.any(torch.isinf(snr_db_list)):
+            raise AssertionError("nan in snr_db_list")
 
         return mix_spec_list, speech_irm_list, snr_db_list
 
@@ -290,15 +290,14 @@ def main():
             snr_db_target = snr_db.to(device)
 
             speech_irm_prediction, lsnr_prediction = model.forward(mix_spec)
-            if torch.any(torch.isnan(speech_irm_prediction)):
+            if torch.any(torch.isnan(speech_irm_prediction)) or torch.any(torch.isinf(speech_irm_prediction)):
                 raise AssertionError("nan in speech_irm_prediction")
-            if torch.any(torch.isnan(lsnr_prediction)):
-                raise AssertionError("nan 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")
             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)):
-                raise AssertionError("nan in snr_loss")
-            print(f"irm_loss: {irm_loss}, snr_loss: {snr_loss}")
+            if torch.any(torch.isnan(snr_loss)) or torch.any(torch.isinf(snr_loss)):
+                raise AssertionError("nan or inf in snr_loss")
             loss = irm_loss + 0 * snr_loss
             # loss = irm_loss
 

toolbox/torchaudio/models/spectrum_unet_irm/modeling_spectrum_unet_irm.py

@@ -392,9 +392,9 @@ class Encoder(nn.Module):
         emb = emb.flatten(2)
         # emb shape: [batch_size, time_steps, hidden_size * channels]
         emb, h = self.emb_gru.forward(emb, hidden_state)
-        print(f"emb: {torch.any(torch.isnan(emb))}")
+        print(f"emb: {torch.any(torch.isnan(emb)) or torch.any(torch.isinf(emb))}")
         lsnr = self.lsnr_fc(emb) * self.lsnr_scale + self.lsnr_offset
-        print(f"lsnr: {torch.any(torch.isnan(lsnr))}")
+        print(f"lsnr: {torch.any(torch.isnan(lsnr)) or torch.any(torch.isinf(lsnr))}")
         return e0, e1, e2, e3, emb, lsnr