update

examples/mpnet_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 1 --stop_stage 2 --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/mpnet_aishell/step_2_train_model.py

@@ -49,7 +49,6 @@ def get_args():
     parser.add_argument("--num_serialized_models_to_keep", default=10, type=int)
     parser.add_argument("--patience", default=5, type=int)
     parser.add_argument("--serialization_dir", default="serialization_dir", type=str)
-    parser.add_argument("--seed", default=0, type=int)
 
     parser.add_argument("--config_file", default="config.yaml", type=str)
 
@@ -79,110 +78,31 @@ def logging_config(file_dir: str):
 
 
 class CollateFunction(object):
-    def __init__(self,
-                 n_fft: int = 512,
-                 win_length: int = 200,
-                 hop_length: int = 80,
-                 window_fn: str = "hamming",
-                 irm_beta: float = 1.0,
-                 epsilon: float = 1e-8,
-                 ):
-        self.n_fft = n_fft
-        self.win_length = win_length
-        self.hop_length = hop_length
-        self.window_fn = window_fn
-        self.irm_beta = irm_beta
-        self.epsilon = epsilon
-
-        self.transform = torchaudio.transforms.Spectrogram(
-            n_fft=self.n_fft,
-            win_length=self.win_length,
-            hop_length=self.hop_length,
-            power=2.0,
-            window_fn=torch.hamming_window if window_fn == "hamming" else torch.hann_window,
-        )
-
-    @staticmethod
-    def make_unfold_snr_db(x: torch.Tensor, n_time_steps: int = 3):
-        batch_size, channels, freq_dim, time_steps = x.shape
-
-        # kernel: [freq_dim, n_time_step]
-        kernel_size = (freq_dim, n_time_steps)
-
-        # pad
-        pad = n_time_steps // 2
-        x = torch.concat(tensors=[
-            x[:, :, :, :pad],
-            x,
-            x[:, :, :, -pad:],
-        ], dim=-1)
-
-        x = F.unfold(
-            input=x,
-            kernel_size=kernel_size,
-        )
-        # x shape: [batch_size, fold, time_steps]
-        return x
+    def __init__(self):
+        pass
 
     def __call__(self, batch: List[dict]):
-        mix_spec_list = list()
-        speech_irm_list = list()
-        snr_db_list = list()
+        clean_audios = list()
+        noisy_audios = list()
+
         for sample in batch:
-            noise_wave: torch.Tensor = sample["noise_wave"]
-            speech_wave: torch.Tensor = sample["speech_wave"]
-            mix_wave: torch.Tensor = sample["mix_wave"]
+            # noise_wave: torch.Tensor = sample["noise_wave"]
+            clean_audio: torch.Tensor = sample["speech_wave"]
+            noisy_audio: torch.Tensor = sample["mix_wave"]
             # snr_db: float = sample["snr_db"]
 
-            noise_spec = self.transform.forward(noise_wave)
-            speech_spec = self.transform.forward(speech_wave)
-            mix_spec = self.transform.forward(mix_wave)
-
-            # noise_irm = noise_spec / (noise_spec + speech_spec)
-            speech_irm = speech_spec / (noise_spec + speech_spec + self.epsilon)
-            speech_irm = torch.pow(speech_irm, self.irm_beta)
-
-            # noise_spec, speech_spec, mix_spec, speech_irm
-            # shape: [freq_dim, time_steps]
-
-            snr_db: torch.Tensor = 10 * torch.log10(
-                speech_spec / (noise_spec + self.epsilon)
-            )
-            snr_db = torch.clamp(snr_db, min=self.epsilon)
-
-            snr_db_ = torch.unsqueeze(snr_db, dim=0)
-            snr_db_ = torch.unsqueeze(snr_db_, dim=0)
-            snr_db_ = self.make_unfold_snr_db(snr_db_, n_time_steps=3)
-            snr_db_ = torch.squeeze(snr_db_, dim=0)
-            # snr_db_ shape: [fold, time_steps]
-
-            snr_db = torch.mean(snr_db_, dim=0, keepdim=True)
-            # snr_db shape: [1, time_steps]
-
-            mix_spec_list.append(mix_spec)
-            speech_irm_list.append(speech_irm)
-            snr_db_list.append(snr_db)
-
-        mix_spec_list = torch.stack(mix_spec_list)
-        speech_irm_list = torch.stack(speech_irm_list)
-        snr_db_list = torch.stack(snr_db_list)  # shape: (batch_size, time_steps, 1)
-
-        mix_spec_list = mix_spec_list[:, :-1, :]
-        speech_irm_list = speech_irm_list[:, :-1, :]
+            clean_audios.append(clean_audio)
+            noisy_audios.append(noisy_audio)
 
-        # mix_spec_list shape: [batch_size, freq_dim, time_steps]
-        # speech_irm_list shape: [batch_size, freq_dim, time_steps]
-        # snr_db shape: [batch_size, 1, time_steps]
+        clean_audios = torch.stack(clean_audios)
+        noisy_audios = torch.stack(noisy_audios)
 
         # assert
-        if torch.any(torch.isnan(mix_spec_list)) or torch.any(torch.isinf(mix_spec_list)):
-            raise AssertionError("nan or inf in mix_spec_list")
-        if torch.any(torch.isnan(speech_irm_list)) or torch.any(torch.isinf(speech_irm_list)):
-            raise AssertionError("nan or inf in speech_irm_list")
-        if torch.any(torch.isnan(snr_db_list)) or torch.any(torch.isinf(snr_db_list)):
-            raise AssertionError("nan or inf in snr_db_list")
-
-        return mix_spec_list, speech_irm_list, snr_db_list
+        if torch.any(torch.isnan(clean_audios)) or torch.any(torch.isinf(clean_audios)):
+            raise AssertionError("nan or inf in clean_audios")
+        if torch.any(torch.isnan(noisy_audios)) or torch.any(torch.isinf(noisy_audios)):
+            raise AssertionError("nan or inf in noisy_audios")
+        return clean_audios, noisy_audios
 
 
 collate_fn = CollateFunction()