add dfnet2

examples/dfnet2/run.sh

@@ -6,7 +6,7 @@ sh run.sh --stage 2 --stop_stage 2 --system_version windows --file_folder_name f
 --noise_dir "E:/Users/tianx/HuggingDatasets/nx_noise/data/noise" \
 --speech_dir "E:/Users/tianx/HuggingDatasets/nx_noise/data/speech"
 
-sh run.sh --stage 2 --stop_stage 2 --system_version centos --file_folder_name file_dir --final_model_name dfnet-nx-dns3 \
+sh run.sh --stage 2 --stop_stage 2 --system_version centos --file_folder_name file_dir --final_model_name dfnet2-nx-dns3 \
 --noise_dir "/data/tianxing/HuggingDatasets/nx_noise/data/noise" \
 --speech_dir "/data/tianxing/HuggingDatasets/nx_noise/data/speech/dns3-speech"
 

examples/dfnet2/step_2_train_model.py

@@ -265,6 +265,9 @@ def main():
             noisy_audios: torch.Tensor = noisy_audios.to(device)
 
             est_spec, est_wav, est_mask, lsnr = model.forward(noisy_audios)
+            # est_wav shape: [b, 1, n_samples]
+            est_wav = torch.squeeze(est_wav, dim=1)
+            # est_wav shape: [b, n_samples]
 
             mr_stft_loss = mr_stft_loss_fn.forward(est_wav, clean_audios)
             neg_si_snr_loss = neg_si_snr_loss_fn.forward(est_wav, clean_audios)
@@ -336,6 +339,9 @@ def main():
                         noisy_audios: torch.Tensor = noisy_audios.to(device)
 
                         est_spec, est_wav, est_mask, lsnr = model.forward(noisy_audios)
+                        # est_wav shape: [b, 1, n_samples]
+                        est_wav = torch.squeeze(est_wav, dim=1)
+                        # est_wav shape: [b, n_samples]
 
                         mr_stft_loss = mr_stft_loss_fn.forward(est_wav, clean_audios)
                         neg_si_snr_loss = neg_si_snr_loss_fn.forward(est_wav, clean_audios)