update

examples/conv_tasnet/run.sh

@@ -3,25 +3,10 @@
 : <<'END'
 
 
-sh run.sh --stage 2 --stop_stage 2 --system_version windows --file_folder_name file_dir --final_model_name mpnet-aishell-20250224 \
---noise_dir "E:/Users/tianx/HuggingDatasets/nx_noise/data/noise" \
---speech_dir "E:/programmer/asr_datasets/aishell/data_aishell/wav/train"
-
-
-sh run.sh --stage 3 --stop_stage 3 --system_version centos --file_folder_name file_dir --final_model_name mpnet-aishell-20250224 \
---noise_dir "/data/tianxing/HuggingDatasets/nx_noise/data/noise" \
---speech_dir "/data/tianxing/HuggingDatasets/aishell/data_aishell/wav/train"
-
-sh run.sh --stage 2 --stop_stage 2 --system_version centos --file_folder_name file_dir --final_model_name nx-clean-unet-aishell-20250228 \
---noise_dir "/data/tianxing/HuggingDatasets/nx_noise/data/noise" \
---speech_dir "/data/tianxing/HuggingDatasets/aishell/data_aishell/wav/train" \
---max_epochs 100
-
-
-sh run.sh --stage 1 --stop_stage 2 --system_version centos --file_folder_name file_dir --final_model_name mpnet-nx-speech-20250224 \
---noise_dir "/data/tianxing/HuggingDatasets/nx_noise/data/noise" \
---speech_dir "/data/tianxing/HuggingDatasets/nx_noise/data/speech" \
---max_epochs 100 --max_count 10000
+sh run.sh --stage 1 --stop_stage 1 --system_version centos --file_folder_name file_dir --final_model_name conv-tasnet-dns3-20250319 \
+--noise_dir "/data/tianxing/HuggingDatasets/nx_noise/data/noise/dns3-noise" \
+--speech_dir "/data/tianxing/HuggingDatasets/nx_noise/data/speech/dns3-speech" \
+--max_epochs 200
 
 
 END

examples/conv_tasnet/step_1_prepare_data.py

@@ -54,28 +54,30 @@ def filename_generator(data_dir: str):
         yield filename.as_posix()
 
 
-def target_second_signal_generator(data_dir: str, duration: int = 2, sample_rate: int = 8000):
+def target_second_signal_generator(data_dir: str, duration: int = 2, sample_rate: int = 8000, max_epoch: int = 20000):
     data_dir = Path(data_dir)
-    for filename in data_dir.glob("**/*.wav"):
-        signal, _ = librosa.load(filename.as_posix(), sr=sample_rate)
-        raw_duration = librosa.get_duration(y=signal, sr=sample_rate)
-
-        if raw_duration < duration:
-            # print(f"duration less than {duration} s. skip filename: {filename.as_posix()}")
-            continue
-        if signal.ndim != 1:
-            raise AssertionError(f"expected ndim 1, instead of {signal.ndim}")
-
-        signal_length = len(signal)
-        win_size = int(duration * sample_rate)
-        for begin in range(0, signal_length - win_size, win_size):
-            row = {
-                "filename": filename.as_posix(),
-                "raw_duration": round(raw_duration, 4),
-                "offset": round(begin / sample_rate, 4),
-                "duration": round(duration, 4),
-            }
-            yield row
+    for epoch_idx in range(max_epoch):
+        for filename in data_dir.glob("**/*.wav"):
+            signal, _ = librosa.load(filename.as_posix(), sr=sample_rate)
+            raw_duration = librosa.get_duration(y=signal, sr=sample_rate)
+
+            if raw_duration < duration:
+                # print(f"duration less than {duration} s. skip filename: {filename.as_posix()}")
+                continue
+            if signal.ndim != 1:
+                raise AssertionError(f"expected ndim 1, instead of {signal.ndim}")
+
+            signal_length = len(signal)
+            win_size = int(duration * sample_rate)
+            for begin in range(0, signal_length - win_size, win_size):
+                row = {
+                    "epoch_idx": epoch_idx,
+                    "filename": filename.as_posix(),
+                    "raw_duration": round(raw_duration, 4),
+                    "offset": round(begin / sample_rate, 4),
+                    "duration": round(duration, 4),
+                }
+                yield row
 
 
 def get_dataset(args):
@@ -88,12 +90,14 @@ def get_dataset(args):
     noise_generator = target_second_signal_generator(
         noise_dir.as_posix(),
         duration=args.duration,
-        sample_rate=args.target_sample_rate
+        sample_rate=args.target_sample_rate,
+        max_epoch=100000,
     )
     speech_generator = target_second_signal_generator(
         speech_dir.as_posix(),
         duration=args.duration,
-        sample_rate=args.target_sample_rate
+        sample_rate=args.target_sample_rate,
+        max_epoch=1,
     )
 
     dataset = list()
@@ -155,7 +159,6 @@ def get_dataset(args):
     return
 
 
-
 def split_dataset(args):
     """分割训练集, 测试集"""
     file_dir = Path(args.file_dir)

examples/conv_tasnet/step_2_train_model.py

@@ -29,7 +29,7 @@ from toolbox.torch.utils.data.dataset.denoise_excel_dataset import DenoiseExcelD
 from toolbox.torchaudio.models.conv_tasnet.configuration_conv_tasnet import ConvTasNetConfig
 from toolbox.torchaudio.models.conv_tasnet.modeling_conv_tasnet import ConvTasNet, ConvTasNetPretrainedModel
 from toolbox.torchaudio.losses.snr import NegativeSISNRLoss
-from toolbox.torchaudio.losses.spectral import LSDLoss
+from toolbox.torchaudio.losses.spectral import LSDLoss, MultiResolutionSTFTLoss
 from toolbox.torchaudio.losses.perceptual import NegSTOILoss
 from toolbox.torchaudio.metrics.pesq import run_pesq_score
 
@@ -39,7 +39,7 @@ def get_args():
     parser.add_argument("--train_dataset", default="train.xlsx", type=str)
     parser.add_argument("--valid_dataset", default="valid.xlsx", type=str)
 
-    parser.add_argument("--max_epochs", default=100, type=int)
+    parser.add_argument("--max_epochs", default=200, type=int)
 
     parser.add_argument("--num_serialized_models_to_keep", default=10, type=int)
     parser.add_argument("--patience", default=5, type=int)
@@ -201,6 +201,12 @@ def main():
     neg_si_snr_loss_fn = NegativeSISNRLoss(reduction="mean").to(device)
     neg_stoi_loss_fn = NegSTOILoss(sample_rate=8000, reduction="mean").to(device)
     lds_loss_fn = LSDLoss(reduction="mean").to(device)
+    mr_stft_loss_fn = MultiResolutionSTFTLoss(
+        fft_size_list=[256, 512, 1024],
+        win_size_list=[120, 240, 480],
+        hop_size_list=[25, 50, 100],
+        reduction="mean"
+    ).to(device)
 
     # training loop
 
@@ -245,8 +251,9 @@ def main():
             neg_si_snr_loss = neg_si_snr_loss_fn.forward(denoise_audios, clean_audios)
             neg_stoi_loss = neg_stoi_loss_fn.forward(denoise_audios, clean_audios)
             lds_loss = lds_loss_fn.forward(denoise_audios, clean_audios)
+            mr_stft_loss = mr_stft_loss_fn.forward(denoise_audios, clean_audios)
 
-            loss = 0.25 * ae_loss + 0.25 * neg_si_snr_loss + 0.25 * neg_stoi_loss + 0.25 * lds_loss
+            loss = 0.25 * ae_loss + 0.25 * neg_si_snr_loss + 0.25 * neg_stoi_loss + 0.25 * lds_loss + 0.25 * mr_stft_loss
 
             denoise_audios_list_r = list(denoise_audios.detach().cpu().numpy())
             clean_audios_list_r = list(clean_audios.detach().cpu().numpy())

examples/conv_tasnet/yaml/config.yaml

@@ -1,42 +1,17 @@
-model_name: "nx_clean_unet"
+model_name: "conv_tasnet"
 
 sample_rate: 8000
-segment_size: 16000
-n_fft: 512
-win_size: 200
-hop_size: 80
+segment_size: 4
 
-down_sampling_num_layers: 6
-down_sampling_in_channels: 1
-down_sampling_hidden_channels: 64
-down_sampling_kernel_size: 4
-down_sampling_stride: 2
+win_size: 20
+freq_bins: 256
+bottleneck_channels: 256
+num_speakers: 1
+num_blocks: 4
+num_sub_blocks: 8
+sub_blocks_channels: 512
+sub_blocks_kernel_size: 3
 
-causal_in_channels: 1
-causal_out_channels: 1
-causal_kernel_size: 3
-causal_bias: false
-causal_separable: true
-causal_f_stride: 1
-causal_num_layers: 5
-
-tsfm_hidden_size: 256
-tsfm_attention_heads: 8
-tsfm_num_blocks: 6
-tsfm_dropout_rate: 0.1
-tsfm_max_length: 512
-tsfm_chunk_size: 1
-tsfm_num_left_chunks: 128
-tsfm_num_right_chunks: 4
-
-discriminator_dim: 32
-discriminator_in_channel: 2
-
-compress_factor: 0.3
-
-batch_size: 64
-learning_rate: 0.0005
-adam_b1: 0.8
-adam_b2: 0.99
-lr_decay: 0.99
-seed: 1234
+norm_type: "gLN"
+causal: false
+mask_nonlinear: "relu"

examples/data_preprocess/dns_challenge_to_8k/process_dns_challenge_clean_read_speech.py

@@ -57,9 +57,11 @@ def get_args():
         # default=r"E:\programmer\asr_datasets\dns-challenge\datasets.clean.read_speech\datasets\clean",
         # default=r"E:\programmer\asr_datasets\dns-challenge\datasets.clean.mandarin_speech\datasets\clean\mandarin_speech",
         # default=r"E:\programmer\asr_datasets\dns-challenge\datasets.clean.singing_voice\datasets\clean\singing_voice",
-        # default=r"E:\programmer\asr_datasets\dns-challenge\datasets.clean.german_speech\datasets\clean\german_speech",
+        # default=r"E:\programmer\asr_datasets\dns-challenge\datasets.clean.french_data\datasets\clean\french_data",
+        default=r"E:\programmer\asr_datasets\dns-challenge\datasets.clean.german_speech\datasets\clean\german_speech",
         # default=r"E:\programmer\asr_datasets\dns-challenge\datasets.clean.italian_speech\datasets\clean\italian_speech",
-        default=r"E:\programmer\asr_datasets\dns-challenge\datasets.clean.russian_speech\datasets\clean\russian_speech",
+        # default=r"E:\programmer\asr_datasets\dns-challenge\datasets.clean.russian_speech\datasets\clean\russian_speech",
+        # default=r"E:\programmer\asr_datasets\dns-challenge\datasets.clean.spanish_speech\datasets\clean\spanish_speech",
         type=str
     )
     parser.add_argument(
@@ -67,9 +69,11 @@ def get_args():
         # default=r"E:\programmer\asr_datasets\denoise\dns-clean-read-speech-8k",
         # default=r"E:\programmer\asr_datasets\denoise\dns-clean-mandarin-speech-8k",
         # default=r"E:\programmer\asr_datasets\denoise\dns-clean-singing-voice-8k",
-        # default=r"E:\programmer\asr_datasets\denoise\dns-clean-german-speech-8k",
+        # default=r"E:\programmer\asr_datasets\denoise\dns-clean-french-speech-8k",
+        default=r"E:\programmer\asr_datasets\denoise\dns-clean-german-speech-8k",
         # default=r"E:\programmer\asr_datasets\denoise\dns-clean-italian-speech-8k",
-        default=r"E:\programmer\asr_datasets\denoise\dns-clean-russian-speech-8k",
+        # default=r"E:\programmer\asr_datasets\denoise\dns-clean-russian-speech-8k",
+        # default=r"E:\programmer\asr_datasets\denoise\dns-clean-spanish-speech-8k",
         type=str
     )
     parser.add_argument("--sample_rate", default=8000, type=int)
@@ -87,24 +91,26 @@ def main():
     # finished_set
     finished_set = set()
     for filename in tqdm(output_dir.glob("**/*.wav")):
-        name = filename.stem
-        finished_set.add(name)
+        filename = Path(filename)
+        relative_name = filename.relative_to(output_dir)
+        relative_name_ = relative_name.as_posix()
+        finished_set.add(relative_name_)
     print(f"finished_set count: {len(finished_set)}")
 
     for filename in tqdm(data_dir.glob("**/*.wav")):
-        label = filename.parts[-2]
-        name = filename.stem
         relative_name = filename.relative_to(data_dir)
-        # print(f"filename: {filename.as_posix()}")
-        if name in finished_set:
+        relative_name_ = relative_name.as_posix()
+        if relative_name_ in finished_set:
             continue
-        finished_set.add(name)
+        finished_set.add(relative_name_)
 
         try:
-            signal, _ = librosa.load(filename.as_posix(), sr=args.sample_rate)
+            signal, _ = librosa.load(filename.as_posix(), mono=False, sr=args.sample_rate)
         except Exception:
             print(f"skip file: {filename.as_posix()}")
             continue
+        if signal.ndim != 1:
+            raise AssertionError
 
         signal = signal * (1 << 15)
         signal = np.array(signal, dtype=np.int16)

toolbox/torchaudio/losses/spectral.py

@@ -283,6 +283,7 @@ class MultiResolutionSTFTLoss(torch.nn.Module):
                  hop_size_list: List[int] = None,
                  factor_sc=0.1,
                  factor_mag=0.1,
+                 reduction: str = "mean",
                  ):
         super(MultiResolutionSTFTLoss, self).__init__()
         fft_size_list = fft_size_list or [1024, 2048, 512]
@@ -299,6 +300,7 @@ class MultiResolutionSTFTLoss(torch.nn.Module):
                     n_fft=n_fft,
                     win_size=win_size,
                     hop_size=hop_size,
+                    reduction=reduction,
                 )
             )
 

toolbox/torchaudio/models/conv_tasnet/modeling_conv_tasnet.py

@@ -388,6 +388,7 @@ class ConvTasNet(nn.Module):
         est_mask = self.separator.forward(mixture_w)
         # est_mask shape: [batch_size, num_speakers, freq_bins, time_steps]
         est_source = self.decoder.forward(mixture_w, est_mask)
+        # est_source shape: [batch_size, num_speakers, num_samples]
 
         num_samples1 = mixture.size(-1)
         num_samples2 = est_source.size(-1)