update

examples/spectrum_dfnet_aishell/step_2_train_model.py

@@ -313,22 +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 = 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)
@@ -361,19 +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 = 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

@@ -803,31 +803,30 @@ class SpectrumDfNet(nn.Module):
     def forward(self,
                 spec_complex: torch.Tensor,
                 ):
-        with torch.no_grad():
-            feat_power = torch.square(torch.abs(spec_complex))
-            feat_power = feat_power.unsqueeze(1).permute(0, 1, 3, 2)
-            # feat_power shape: [batch_size, spec_bins, time_steps]
-            # feat_power shape: [batch_size, 1, spec_bins, time_steps]
-            # feat_power shape: [batch_size, 1, time_steps, spec_bins]
-            feat_power = feat_power.detach()
-
-            # spec shape: [batch_size, spec_bins, time_steps]
-            feat_spec = torch.view_as_real(spec_complex)
-            # spec shape: [batch_size, spec_bins, time_steps, 2]
-            feat_spec = feat_spec.permute(0, 3, 2, 1)
-            # feat_spec shape: [batch_size, 2, time_steps, spec_bins]
-            feat_spec = feat_spec[..., :self.df_decoder.df_bins]
-            # feat_spec shape: [batch_size, 2, time_steps, df_bins]
-            feat_spec = feat_spec.detach()
-
-        # # spec shape: [batch_size, spec_bins, time_steps]
-        # spec = torch.unsqueeze(spec_complex, dim=1)
-        # # spec shape: [batch_size, 1, spec_bins, time_steps]
-        # spec = spec.permute(0, 1, 3, 2)
-        # # spec shape: [batch_size, 1, time_steps, spec_bins]
-        # spec = torch.view_as_real(spec)
-        # # spec shape: [batch_size, 1, time_steps, spec_bins, 2]
-        # spec = spec.detach()
+        feat_power = torch.square(torch.abs(spec_complex))
+        feat_power = feat_power.unsqueeze(1).permute(0, 1, 3, 2)
+        # feat_power shape: [batch_size, spec_bins, time_steps]
+        # feat_power shape: [batch_size, 1, spec_bins, time_steps]
+        # feat_power shape: [batch_size, 1, time_steps, spec_bins]
+        feat_power = feat_power.detach()
+
+        # spec shape: [batch_size, spec_bins, time_steps]
+        feat_spec = torch.view_as_real(spec_complex)
+        # spec shape: [batch_size, spec_bins, time_steps, 2]
+        feat_spec = feat_spec.permute(0, 3, 2, 1)
+        # feat_spec shape: [batch_size, 2, time_steps, spec_bins]
+        feat_spec = feat_spec[..., :self.df_decoder.df_bins]
+        # feat_spec shape: [batch_size, 2, time_steps, df_bins]
+        feat_spec = feat_spec.detach()
+
+        # spec shape: [batch_size, spec_bins, time_steps]
+        spec = torch.unsqueeze(spec_complex, dim=1)
+        # spec shape: [batch_size, 1, spec_bins, time_steps]
+        spec = spec.permute(0, 1, 3, 2)
+        # spec shape: [batch_size, 1, time_steps, spec_bins]
+        spec = torch.view_as_real(spec)
+        # spec shape: [batch_size, 1, time_steps, spec_bins, 2]
+        spec = spec.detach()
 
         e0, e1, e2, e3, emb, c0, lsnr, h = self.encoder.forward(feat_power, feat_spec)
 
@@ -836,31 +835,30 @@ class SpectrumDfNet(nn.Module):
         if torch.any(mask > 1) or torch.any(mask < 0):
             raise AssertionError
 
-        # spec_m = self.mask.forward(spec, mask)
+        spec_m = self.mask.forward(spec, mask)
 
         # lsnr shape: [batch_size, time_steps, 1]
         lsnr = torch.transpose(lsnr, dim0=2, dim1=1)
         # lsnr shape: [batch_size, 1, time_steps]
 
-        # df_coefs = self.df_decoder.forward(emb, c0)
-        # df_coefs = self.df_out_transform(df_coefs)
-        # # df_coefs shape: [batch_size, df_order, time_steps, df_bins, 2]
-        #
-        # spec_e = self.df_op.forward(spec.clone(), df_coefs)
-        # # spec_e shape: [batch_size, 1, time_steps, spec_bins, 2]
-        #
-        # spec_e[..., self.df_decoder.df_bins:, :] = spec_m[..., self.df_decoder.df_bins:, :]
-        #
-        # spec_e = torch.squeeze(spec_e, dim=1)
-        # spec_e = spec_e.permute(0, 2, 1, 3)
-        # # spec_e shape: [batch_size, spec_bins, time_steps, 2]
+        df_coefs = self.df_decoder.forward(emb, c0)
+        df_coefs = self.df_out_transform(df_coefs)
+        # df_coefs shape: [batch_size, df_order, time_steps, df_bins, 2]
+
+        spec_e = self.df_op.forward(spec.clone(), df_coefs)
+        # spec_e shape: [batch_size, 1, time_steps, spec_bins, 2]
+
+        spec_e[..., self.df_decoder.df_bins:, :] = spec_m[..., self.df_decoder.df_bins:, :]
+
+        spec_e = torch.squeeze(spec_e, dim=1)
+        spec_e = spec_e.permute(0, 2, 1, 3)
+        # spec_e shape: [batch_size, spec_bins, time_steps, 2]
 
         mask = torch.squeeze(mask, dim=1)
         mask = mask.permute(0, 2, 1)
         # mask shape: [batch_size, spec_bins, time_steps]
 
-        # return spec_e, mask, lsnr
-        return None, mask, lsnr
+        return spec_e, mask, lsnr
 
 
 class SpectrumDfNetPretrainedModel(SpectrumDfNet):