update

examples/conv_tasnet/run.sh

@@ -3,7 +3,7 @@
 : <<'END'
 
 
-sh run.sh --stage 2 --stop_stage 2 --system_version centos --file_folder_name file_dir --final_model_name conv-tasnet-dns3-20250319 \
+sh run.sh --stage 1 --stop_stage 2 --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 400

examples/conv_tasnet/step_2_train_model.py

@@ -139,11 +139,17 @@ def main():
         jsonl_file=args.train_dataset,
         expected_sample_rate=config.sample_rate,
         max_wave_value=32768.0,
+        min_snr_db=config.min_snr_db,
+        max_snr_db=config.max_snr_db,
+        # skip=625000,
     )
     valid_dataset = DenoiseJsonlDataset(
         jsonl_file=args.valid_dataset,
         expected_sample_rate=config.sample_rate,
         max_wave_value=32768.0,
+        min_snr_db=config.min_snr_db,
+        max_snr_db=config.max_snr_db,
+        # skip=625000,
     )
     train_data_loader = DataLoader(
         dataset=train_dataset,

examples/conv_tasnet/yaml/config.yaml

@@ -16,6 +16,9 @@ norm_type: "gLN"
 causal: false
 mask_nonlinear: "relu"
 
+min_snr_db: -10
+max_snr_db: 20
+
 lr: 0.001
 lr_scheduler: "CosineAnnealingLR"
 lr_scheduler_kwargs:

toolbox/torch/utils/data/dataset/denoise_jsonl_dataset.py

@@ -1,18 +1,13 @@
 #!/usr/bin/python3
 # -*- coding: utf-8 -*-
 import json
-import os
 import random
 from typing import List
 
 import librosa
 import numpy as np
-import pandas as pd
-from scipy.io import wavfile
 import torch
-import torchaudio
 from torch.utils.data import Dataset, IterableDataset
-from tqdm import tqdm
 
 
 class DenoiseJsonlDataset(IterableDataset):
@@ -22,13 +17,19 @@ class DenoiseJsonlDataset(IterableDataset):
                  resample: bool = False,
                  max_wave_value: float = 1.0,
                  buffer_size: int = 1000,
+                 min_snr_db: float = None,
+                 max_snr_db: float = None,
                  eps: float = 1e-8,
+                 skip: int = 0,
                  ):
         self.jsonl_file = jsonl_file
         self.expected_sample_rate = expected_sample_rate
         self.resample = resample
         self.max_wave_value = max_wave_value
+        self.min_snr_db = min_snr_db
+        self.max_snr_db = max_snr_db
         self.eps = eps
+        self.skip = skip
 
         self.buffer_size = buffer_size
         self.buffer_samples: List[dict] = list()
@@ -36,6 +37,12 @@ class DenoiseJsonlDataset(IterableDataset):
     def __iter__(self):
         iterable_source = self.iterable_source()
 
+        try:
+            for _ in range(self.skip):
+                next(iterable_source)
+        except StopIteration:
+            pass
+
         # 初始填充缓冲区
         try:
             for _ in range(self.buffer_size):
@@ -74,7 +81,10 @@ class DenoiseJsonlDataset(IterableDataset):
                 speech_offset = row["speech_offset"]
                 speech_duration = row["speech_duration"]
 
-                snr_db = row["snr_db"]
+                if self.min_snr_db is None or self.max_snr_db is None:
+                    snr_db = row["snr_db"]
+                else:
+                    snr_db = random.uniform(self.min_snr_db, self.max_snr_db)
 
                 sample = {
                     "noise_filename": noise_filename,

toolbox/torchaudio/models/conv_tasnet/configuration_conv_tasnet.py

@@ -27,12 +27,15 @@ class ConvTasNetConfig(PretrainedConfig):
                  causal: bool = False,
                  mask_nonlinear: str = "relu",
 
-                 lr: float = 1e-3,
-                 eval_steps: int = 25000,
+                 min_snr_db: float = -10,
+                 max_snr_db: float = 20,
 
+                 lr: float = 1e-3,
                  lr_scheduler: str = "CosineAnnealingLR",
                  lr_scheduler_kwargs: dict = None,
 
+                 eval_steps: int = 25000,
+
                  **kwargs
                  ):
         super(ConvTasNetConfig, self).__init__(**kwargs)
@@ -53,12 +56,15 @@ class ConvTasNetConfig(PretrainedConfig):
         self.causal = causal
         self.mask_nonlinear = mask_nonlinear
 
-        self.lr = lr
-        self.eval_steps = eval_steps
+        self.min_snr_db = min_snr_db
+        self.max_snr_db = max_snr_db
 
+        self.lr = lr
         self.lr_scheduler = lr_scheduler
         self.lr_scheduler_kwargs = lr_scheduler_kwargs or dict()
 
+        self.eval_steps = eval_steps
+
 
 if __name__ == "__main__":
     pass