update

examples/dfnet2/run.sh

@@ -29,7 +29,7 @@ limit=10
 noise_dir=/data/tianxing/HuggingDatasets/nx_noise/data/noise
 speech_dir=/data/tianxing/HuggingDatasets/aishell/data_aishell/wav/train
 
-max_count=10000000
+max_count=-1
 
 nohup_name=nohup.out
 

examples/dfnet2/step_1_prepare_data.py

@@ -33,13 +33,13 @@ def get_args():
     parser.add_argument("--train_dataset", default="train.jsonl", type=str)
     parser.add_argument("--valid_dataset", default="valid.jsonl", type=str)
 
-    parser.add_argument("--duration", default=4.0, type=float)
+    parser.add_argument("--duration", default=2.0, type=float)
     parser.add_argument("--min_snr_db", default=-10, type=float)
     parser.add_argument("--max_snr_db", default=20, type=float)
 
     parser.add_argument("--target_sample_rate", default=8000, type=int)
 
-    parser.add_argument("--max_count", default=10000, type=int)
+    parser.add_argument("--max_count", default=-1, type=int)
 
     args = parser.parse_args()
     return args

examples/dfnet2/yaml/config.yaml

@@ -68,7 +68,7 @@ clip_grad_norm: 10.0
 seed: 1234
 
 num_workers: 8
-batch_size: 64
+batch_size: 96
 eval_steps: 10000
 
 # runtime

examples/dtln/run.sh

@@ -31,7 +31,7 @@ limit=10
 noise_dir=/data/tianxing/HuggingDatasets/nx_noise/data/noise
 speech_dir=/data/tianxing/HuggingDatasets/aishell/data_aishell/wav/train
 
-max_count=10000000
+max_count=-1
 
 nohup_name=nohup.out
 

examples/dtln/step_1_prepare_data.py

@@ -33,13 +33,13 @@ def get_args():
     parser.add_argument("--train_dataset", default="train.jsonl", type=str)
     parser.add_argument("--valid_dataset", default="valid.jsonl", type=str)
 
-    parser.add_argument("--duration", default=4.0, type=float)
+    parser.add_argument("--duration", default=2.0, type=float)
     parser.add_argument("--min_snr_db", default=-10, type=float)
     parser.add_argument("--max_snr_db", default=20, type=float)
 
     parser.add_argument("--target_sample_rate", default=8000, type=int)
 
-    parser.add_argument("--max_count", default=10000, type=int)
+    parser.add_argument("--max_count", default=-1, type=int)
 
     args = parser.parse_args()
     return args

main.py

@@ -72,22 +72,22 @@ def shell(cmd: str):
 
 
 denoise_engines = {
-    "dtln-nx-dns3": {
+    "dtln-256-nx-dns3": {
         "infer_cls": InferenceDTLN,
         "kwargs": {
-            "pretrained_model_path_or_zip_file": (project_path / "trained_models/dtln-nx-dns3.zip").as_posix()
+            "pretrained_model_path_or_zip_file": (project_path / "trained_models/dtln-256-nx-dns3.zip").as_posix()
         }
     },
-    "dfnet2-nx-dns3": {
-        "infer_cls": InferenceDfNet2,
+    "dtln-512-nx-dns3": {
+        "infer_cls": InferenceDTLN,
         "kwargs": {
-            "pretrained_model_path_or_zip_file": (project_path / "trained_models/dfnet2-nx-dns3.zip").as_posix()
+            "pretrained_model_path_or_zip_file": (project_path / "trained_models/dtln-512-nx-dns3.zip").as_posix()
         }
     },
-    "dfnet-nx-dns3": {
-        "infer_cls": InferenceDfNet,
+    "dfnet2-nx-dns3": {
+        "infer_cls": InferenceDfNet2,
         "kwargs": {
-            "pretrained_model_path_or_zip_file": (project_path / "trained_models/dfnet-nx-dns3.zip").as_posix()
+            "pretrained_model_path_or_zip_file": (project_path / "trained_models/dfnet2-nx-dns3.zip").as_posix()
         }
     },
     "frcrn-dns3": {
@@ -114,13 +114,11 @@ def load_denoise_model(infer_cls, **kwargs):
 
 def generate_spectrogram(signal: np.ndarray, sample_rate: int = 8000, title: str = "Spectrogram"):
     mag = np.abs(librosa.stft(signal))
-    mag_db = librosa.amplitude_to_db(mag, ref=np.max)
+    # mag_db = librosa.amplitude_to_db(mag, ref=np.max)
+    mag_db = librosa.amplitude_to_db(mag, ref=20)
 
-    # 保存为临时图片文件
-    plt.figure(figsize=(10, 3))
+    plt.figure(figsize=(10, 4))
     librosa.display.specshow(mag_db, sr=sample_rate)
-    # librosa.display.specshow(mag_db, sr=sample_rate, x_axis='time', y_axis='log')
-    # plt.colorbar(format='%+2.0f dB')
     plt.title(title)
 
     temp_file = tempfile.NamedTemporaryFile(suffix=".png", delete=False)

toolbox/torchaudio/models/dfnet2/modeling_dfnet2.py

@@ -1464,7 +1464,14 @@ def main():
     import time
     # torch.set_num_threads(1)
 
-    config = DfNet2Config()
+    config = DfNet2Config(
+        # nfft=512,
+        # win_size=200,
+        # hop_size=80,
+        nfft=512,
+        win_size=512,
+        hop_size=128,
+    )
     model = DfNet2PretrainedModel(config=config)
     model.eval()
 
@@ -1473,7 +1480,8 @@ def main():
     duration = num_samples / config.sample_rate
 
     begin = time.time()
-    est_spec, est_wav, est_mask, lsnr = model.forward(noisy)
+    with torch.no_grad():
+        est_spec, est_wav, est_mask, lsnr = model.forward(noisy)
     time_cost = time.time() - begin
     print(f"time_cost: {time_cost:.4f}, audio_duration: {duration:.4f}, fpr: {time_cost / duration:.4f}")
 
@@ -1485,6 +1493,9 @@ def main():
     waveform = est_wav
     print(f"waveform.shape: {waveform.shape}, waveform.dtype: {waveform.dtype}")
     print(waveform[:, :, 300: 302])
+    print(waveform[:, :, 1000: 1002])
+    print(waveform[:, :, 8000: 8002])
+    print(waveform[:, :, 14000: 14002])
     print(waveform[:, :, 15680: 15682])
     print(waveform[:, :, 15760: 15762])
     print(waveform[:, :, 15840: 15842])
@@ -1497,6 +1508,9 @@ def main():
     waveform = waveform[:, :, (config.df_lookahead*config.hop_size):]
     print(f"waveform.shape: {waveform.shape}, waveform.dtype: {waveform.dtype}")
     print(waveform[:, :, 300: 302])
+    print(waveform[:, :, 1000: 1002])
+    print(waveform[:, :, 8000: 8002])
+    print(waveform[:, :, 14000: 14002])
     print(waveform[:, :, 15680: 15682])
     print(waveform[:, :, 15760: 15762])
     print(waveform[:, :, 15840: 15842])

toolbox/torchaudio/models/dfnet2/yaml/{config.yaml → config-200.yaml}

@@ -1,4 +1,4 @@
-model_name: "dfnet"
+model_name: "dfnet2"
 
 # spec
 sample_rate: 8000
@@ -68,7 +68,7 @@ clip_grad_norm: 10.0
 seed: 1234
 
 num_workers: 8
-batch_size: 64
+batch_size: 96
 eval_steps: 10000
 
 # runtime

toolbox/torchaudio/models/dfnet2/yaml/config-512.yaml

@@ -0,0 +1,75 @@
+model_name: "dfnet"
+
+# spec
+sample_rate: 8000
+nfft: 512
+win_size: 512
+hop_size: 128
+
+spec_bins: 256
+erb_bins: 32
+min_freq_bins_for_erb: 2
+use_ema_norm: true
+
+# model
+conv_channels: 64
+conv_kernel_size_input:
+  - 3
+  - 3
+conv_kernel_size_inner:
+  - 1
+  - 3
+convt_kernel_size_inner:
+  - 1
+  - 3
+
+embedding_hidden_size: 256
+encoder_combine_op: "concat"
+
+encoder_emb_skip_op: "none"
+encoder_emb_linear_groups: 16
+encoder_emb_hidden_size: 256
+
+encoder_linear_groups: 32
+
+decoder_emb_num_layers: 3
+decoder_emb_skip_op: "none"
+decoder_emb_linear_groups: 16
+decoder_emb_hidden_size: 256
+
+df_decoder_hidden_size: 256
+df_num_layers: 2
+df_order: 5
+df_bins: 96
+df_gru_skip: "grouped_linear"
+df_decoder_linear_groups: 16
+df_pathway_kernel_size_t: 5
+df_lookahead: 2
+
+# lsnr
+n_frame: 3
+lsnr_max: 30
+lsnr_min: -15
+norm_tau: 1.
+
+# data
+min_snr_db: -10
+max_snr_db: 20
+
+# train
+lr: 0.001
+lr_scheduler: "CosineAnnealingLR"
+lr_scheduler_kwargs:
+  T_max: 250000
+  eta_min: 0.0001
+
+max_epochs: 100
+clip_grad_norm: 10.0
+seed: 1234
+
+num_workers: 8
+batch_size: 96
+eval_steps: 10000
+
+# runtime
+use_post_filter: true