update

examples/spectrum_dfnet_aishell/step_2_train_model.py

@@ -328,7 +328,6 @@ def main():
             #     raise AssertionError("nan or inf in 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)
@@ -373,7 +372,6 @@ def main():
                 snr_loss = snr_mse_loss.forward(lsnr_prediction, snr_db_target)
 
                 loss = speech_loss + irm_loss + snr_loss
-                # loss = irm_loss + snr_loss
 
                 total_loss += loss.item()
             total_examples += mix_complex_spec.size(0)

examples/spectrum_dfnet_aishell/step_3_evaluation.py

@@ -94,8 +94,12 @@ istft = torchaudio.transforms.InverseSpectrogram(
 )
 
 
-def enhance(mix_spec_complex: torch.Tensor, speech_irm_prediction: torch.Tensor):
+def enhance(mix_spec_complex: torch.Tensor,
+            speech_spec_prediction: torch.Tensor,
+            speech_irm_prediction: torch.Tensor,
+            ):
     mix_spec_complex = mix_spec_complex.detach().cpu()
+    speech_spec_prediction = speech_spec_prediction.detach().cpu()
     speech_irm_prediction = speech_irm_prediction.detach().cpu()
 
     mask_speech = speech_irm_prediction
@@ -104,7 +108,8 @@ def enhance(mix_spec_complex: torch.Tensor, speech_irm_prediction: torch.Tensor)
     speech_spec = mix_spec_complex * mask_speech
     noise_spec = mix_spec_complex * mask_noise
 
-    speech_wave = istft.forward(speech_spec)
+    speech_wave = istft.forward(speech_spec_prediction)
+    # speech_wave = istft.forward(speech_spec)
     noise_wave = istft.forward(noise_spec)
 
     return speech_wave, noise_wave
@@ -212,6 +217,7 @@ def main():
         speech_spec = speech_spec[:, :-1, :]
         mix_spec = mix_spec[:, :-1, :]
 
+        speech_spec_complex: torch.Tensor = stft_complex.forward(speech_wave)
         mix_spec_complex: torch.Tensor = stft_complex.forward(mix_wave)
         # mix_spec_complex shape: [batch_size, freq_dim (257), time_steps, 2]
 
@@ -221,6 +227,7 @@ def main():
         snr_db: torch.Tensor = 10 * torch.log10(
             speech_spec / (noise_spec + 1e-8)
         )
+        snr_db = torch.clamp(snr_db, min=1e-8)
         snr_db = torch.mean(snr_db, dim=1, keepdim=True)
         # snr_db shape: [batch_size, 1, time_steps]
 
@@ -229,7 +236,7 @@ def main():
         snr_db_target = snr_db.to(device)
 
         with torch.no_grad():
-            speech_irm_prediction, lsnr_prediction = model.forward(mix_spec)
+            speech_spec_prediction, speech_irm_prediction, lsnr_prediction = model.forward(mix_spec_complex)
             irm_loss = irm_mse_loss.forward(speech_irm_prediction, speech_irm_target)
             # snr_loss = snr_mse_loss.forward(lsnr_prediction, snr_db_target)
             # loss = irm_loss + 0.1 * snr_loss
@@ -246,7 +253,7 @@ def main():
             dim=1,
         )
         # speech_irm_prediction shape: [batch_size, freq_dim (257), time_steps]
-        speech_wave_enhanced, noise_wave_enhanced = enhance(mix_spec_complex, speech_irm_prediction)
+        speech_wave_enhanced, noise_wave_enhanced = enhance(mix_spec_complex, speech_spec_prediction, speech_irm_prediction)
         save_audios(noise_wave, speech_wave, mix_wave, speech_wave_enhanced, noise_wave_enhanced, args.evaluation_audio_dir)
 
         total_loss += loss.item()