update

examples/spectrum_dfnet_aishell/run.sh

@@ -12,7 +12,7 @@ sh run.sh --stage 2 --stop_stage 2 --system_version centos --file_folder_name fi
 --noise_dir "/data/tianxing/HuggingDatasets/nx_noise/data/noise" \
 --speech_dir "/data/tianxing/HuggingDatasets/aishell/data_aishell/wav/train"
 
-sh run.sh --stage 3 --stop_stage 3 --system_version centos --file_folder_name file_dir \
+sh run.sh --stage 2 --stop_stage 3 --system_version centos --file_folder_name file_dir \
 --noise_dir "/data/tianxing/HuggingDatasets/nx_noise/data/noise" \
 --speech_dir "/data/tianxing/HuggingDatasets/aishell/data_aishell/wav/train"
 

examples/spectrum_dfnet_aishell/step_2_train_model.py

@@ -324,9 +324,6 @@ def main():
             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
 
             total_loss += loss.item()

examples/spectrum_dfnet_aishell/step_3_evaluation.py

@@ -19,7 +19,7 @@ import torch.nn as nn
 import torchaudio
 from tqdm import tqdm
 
-from toolbox.torchaudio.models.spectrum_unet_irm.modeling_spectrum_unet_irm import SpectrumUnetIRMPretrainedModel
+from toolbox.torchaudio.models.spectrum_dfnet.modeling_spectrum_dfnet import SpectrumDfNetPretrainedModel
 
 
 def get_args():
@@ -152,7 +152,7 @@ def main():
     logger.info("GPU available count: {}; device: {}".format(n_gpu, device))
 
     logger.info("prepare model")
-    model = SpectrumUnetIRMPretrainedModel.from_pretrained(
+    model = SpectrumDfNetPretrainedModel.from_pretrained(
         pretrained_model_name_or_path=args.model_dir,
     )
     model.to(device)