first commit

.dockerignore

@@ -0,0 +1,5 @@
+
+.git/
+.idea/
+
+/examples/

.gitattributes

@@ -0,0 +1,35 @@
+*.7z filter=lfs diff=lfs merge=lfs -text
+*.arrow filter=lfs diff=lfs merge=lfs -text
+*.bin filter=lfs diff=lfs merge=lfs -text
+*.bz2 filter=lfs diff=lfs merge=lfs -text
+*.ckpt filter=lfs diff=lfs merge=lfs -text
+*.ftz filter=lfs diff=lfs merge=lfs -text
+*.gz filter=lfs diff=lfs merge=lfs -text
+*.h5 filter=lfs diff=lfs merge=lfs -text
+*.joblib filter=lfs diff=lfs merge=lfs -text
+*.lfs.* filter=lfs diff=lfs merge=lfs -text
+*.mlmodel filter=lfs diff=lfs merge=lfs -text
+*.model filter=lfs diff=lfs merge=lfs -text
+*.msgpack filter=lfs diff=lfs merge=lfs -text
+*.npy filter=lfs diff=lfs merge=lfs -text
+*.npz filter=lfs diff=lfs merge=lfs -text
+*.onnx filter=lfs diff=lfs merge=lfs -text
+*.ot filter=lfs diff=lfs merge=lfs -text
+*.parquet filter=lfs diff=lfs merge=lfs -text
+*.pb filter=lfs diff=lfs merge=lfs -text
+*.pickle filter=lfs diff=lfs merge=lfs -text
+*.pkl filter=lfs diff=lfs merge=lfs -text
+*.pt filter=lfs diff=lfs merge=lfs -text
+*.pth filter=lfs diff=lfs merge=lfs -text
+*.rar filter=lfs diff=lfs merge=lfs -text
+*.safetensors filter=lfs diff=lfs merge=lfs -text
+saved_model/**/* filter=lfs diff=lfs merge=lfs -text
+*.tar.* filter=lfs diff=lfs merge=lfs -text
+*.tar filter=lfs diff=lfs merge=lfs -text
+*.tflite filter=lfs diff=lfs merge=lfs -text
+*.tgz filter=lfs diff=lfs merge=lfs -text
+*.wasm filter=lfs diff=lfs merge=lfs -text
+*.xz filter=lfs diff=lfs merge=lfs -text
+*.zip filter=lfs diff=lfs merge=lfs -text
+*.zst filter=lfs diff=lfs merge=lfs -text
+*tfevents* filter=lfs diff=lfs merge=lfs -text

.gitignore

@@ -0,0 +1,25 @@
+
+.gradio/
+.git/
+.idea/
+
+**/evaluation_audio/
+**/file_dir/
+**/flagged/
+**/log/
+**/logs/
+**/__pycache__/
+
+/data/
+/docs/
+/dotenv/
+/hub_datasets/
+/script/
+/thirdparty/
+/trained_models/
+/temp/
+
+**/*.wav
+**/*.xlsx
+
+requirements-python-3-9-9.txt

Dockerfile

@@ -0,0 +1,24 @@
+FROM python:3.12
+
+WORKDIR /code
+
+COPY . /code
+
+RUN apt-get update
+RUN apt-get install -y ffmpeg build-essential
+
+RUN pip install --upgrade pip
+RUN pip install --no-cache-dir --upgrade -r /code/requirements.txt
+
+RUN useradd -m -u 1000 user
+
+USER user
+
+ENV HOME=/home/user \
+    PATH=/home/user/.local/bin:$PATH
+
+WORKDIR $HOME/app
+
+COPY --chown=user . $HOME/app
+
+CMD ["python3", "main.py"]

README.md

@@ -0,0 +1,129 @@
+---
+title: NX Denoise
+emoji: 🐢
+colorFrom: purple
+colorTo: blue
+sdk: docker
+pinned: false
+license: apache-2.0
+---
+
+Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+## NX Denoise
+
+
+### datasets
+
+```text
+
+AISHELL (15G)
+https://openslr.trmal.net/resources/33/
+
+AISHELL-3 (19G)
+http://www.openslr.org/93/
+
+DNS3
+https://github.com/microsoft/DNS-Challenge/blob/master/download-dns-challenge-3.sh
+噪音数据来源于 DEMAND, FreeSound, AudioSet. 
+
+MS-SNSD
+https://github.com/microsoft/MS-SNSD
+噪音数据来源于 DEMAND, FreeSound. 
+
+MUSAN
+https://www.openslr.org/17/
+其中包含 music, noise, speech. 
+music 是一些纯音乐, noise 包含 free-sound, sound-bible, sound-bible部分也许可以做为补充部分. 
+总的来说, 有用的不部不多, 可能噪音数据仍然需要自己收集为主, 更加可靠. 
+
+CHiME-4
+https://www.chimechallenge.org/challenges/chime4/download.html
+
+freesound
+https://freesound.org/
+
+AudioSet
+https://research.google.com/audioset/index.html
+```
+
+
+### ### 创建训练容器
+
+```text
+在容器中训练模型，需要能够从容器中访问到 GPU，参考：
+https://hub.docker.com/r/ollama/ollama
+
+docker run -itd \
+--name nx_denoise \
+--network host \
+--gpus all \
+--privileged \
+--ipc=host \
+-v /data/tianxing/HuggingDatasets/nx_noise/data:/data/tianxing/HuggingDatasets/nx_noise/data \
+-v /data/tianxing/PycharmProjects/nx_denoise:/data/tianxing/PycharmProjects/nx_denoise \
+python:3.12
+
+
+查看GPU
+nvidia-smi
+watch -n 1 -d nvidia-smi
+
+
+```
+
+```text
+在容器中访问 GPU
+
+参考：
+https://blog.csdn.net/footless_bird/article/details/136291344
+步骤：
+# 安装 
+yum install -y nvidia-container-toolkit
+
+# 编辑文件 /etc/docker/daemon.json
+cat /etc/docker/daemon.json
+{
+    "data-root": "/data/lib/docker",
+    "default-runtime": "nvidia",
+    "runtimes": {
+        "nvidia": {
+            "path": "/usr/bin/nvidia-container-runtime",
+            "runtimeArgs": []
+        }
+    },
+    "registry-mirrors": [
+        "https://docker.m.daocloud.io", 
+        "https://dockerproxy.com", 
+        "https://docker.mirrors.ustc.edu.cn", 
+        "https://docker.nju.edu.cn"
+    ]
+}
+
+# 重启 docker
+systemctl restart docker
+systemctl daemon-reload
+
+# 测试容器内能否访问 GPU. 
+docker run --gpus all python:3.12-slim nvidia-smi
+
+# 通过这种方式启动容器, 在容器中, 可以查看到 GPU. 但是容器中没有 GPU驱动 nvidia-smi 不工作. 
+docker run -it --privileged python:3.12-slim /bin/bash
+apt update
+apt install -y pciutils
+lspci | grep -i nvidia
+#00:08.0 3D controller: NVIDIA Corporation TU104GL [Tesla T4] (rev a1)
+
+# 网上看的是这种启动容器的方式, 但是进去后仍然是 nvidia-smi 不工作.
+docker run \
+--device /dev/nvidia0:/dev/nvidia0 \
+--device /dev/nvidiactl:/dev/nvidiactl \
+--device /dev/nvidia-uvm:/dev/nvidia-uvm \
+-v /usr/local/nvidia:/usr/local/nvidia \
+-it --privileged python:3.12-slim /bin/bash
+
+
+# 这种方式进入容器, nvidia-smi 可以工作. 应该关键是 --gpus all 参数. 
+docker run -itd --gpus all --name open_unsloth python:3.12-slim /bin/bash
+docker run -itd --gpus all --name Qwen2-7B-Instruct python:3.12-slim /bin/bash
+
+```

examples/clean_unet/run.sh

@@ -0,0 +1,181 @@
+#!/usr/bin/env bash
+
+: <<'END'
+
+
+sh run.sh --stage 2 --stop_stage 2 --system_version windows --file_folder_name file_dir \
+--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 2 --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"
+
+sh run.sh --stage 2 --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"
+
+
+END
+
+
+# params
+system_version="windows";
+verbose=true;
+stage=0 # start from 0 if you need to start from data preparation
+stop_stage=9
+
+work_dir="$(pwd)"
+file_folder_name=file_folder_name
+final_model_name=final_model_name
+config_file="yaml/config.yaml"
+limit=10
+
+noise_dir=/data/tianxing/HuggingDatasets/nx_noise/data/noise
+speech_dir=/data/tianxing/HuggingDatasets/aishell/data_aishell/wav/train
+
+max_count=10000000
+
+nohup_name=nohup.out
+
+# model params
+batch_size=64
+max_epochs=200
+save_top_k=10
+patience=5
+
+
+# parse options
+while true; do
+  [ -z "${1:-}" ] && break;  # break if there are no arguments
+  case "$1" in
+    --*) name=$(echo "$1" | sed s/^--// | sed s/-/_/g);
+      eval '[ -z "${'"$name"'+xxx}" ]' && echo "$0: invalid option $1" 1>&2 && exit 1;
+      old_value="(eval echo \\$$name)";
+      if [ "${old_value}" == "true" ] || [ "${old_value}" == "false" ]; then
+        was_bool=true;
+      else
+        was_bool=false;
+      fi
+
+      # Set the variable to the right value-- the escaped quotes make it work if
+      # the option had spaces, like --cmd "queue.pl -sync y"
+      eval "${name}=\"$2\"";
+
+      # Check that Boolean-valued arguments are really Boolean.
+      if $was_bool && [[ "$2" != "true" && "$2" != "false" ]]; then
+        echo "$0: expected \"true\" or \"false\": $1 $2" 1>&2
+        exit 1;
+      fi
+      shift 2;
+      ;;
+
+    *) break;
+  esac
+done
+
+file_dir="${work_dir}/${file_folder_name}"
+final_model_dir="${work_dir}/../../trained_models/${final_model_name}";
+evaluation_audio_dir="${file_dir}/evaluation_audio"
+
+dataset="${file_dir}/dataset.xlsx"
+train_dataset="${file_dir}/train.xlsx"
+valid_dataset="${file_dir}/valid.xlsx"
+
+$verbose && echo "system_version: ${system_version}"
+$verbose && echo "file_folder_name: ${file_folder_name}"
+
+if [ $system_version == "windows" ]; then
+  alias python3='D:/Users/tianx/PycharmProjects/virtualenv/nx_denoise/Scripts/python.exe'
+elif [ $system_version == "centos" ] || [ $system_version == "ubuntu" ]; then
+  #source /data/local/bin/nx_denoise/bin/activate
+  alias python3='/data/local/bin/nx_denoise/bin/python3'
+fi
+
+
+if [ ${stage} -le 1 ] && [ ${stop_stage} -ge 1 ]; then
+  $verbose && echo "stage 1: prepare data"
+  cd "${work_dir}" || exit 1
+  python3 step_1_prepare_data.py \
+  --file_dir "${file_dir}" \
+  --noise_dir "${noise_dir}" \
+  --speech_dir "${speech_dir}" \
+  --train_dataset "${train_dataset}" \
+  --valid_dataset "${valid_dataset}" \
+  --max_count "${max_count}" \
+
+fi
+
+
+if [ ${stage} -le 2 ] && [ ${stop_stage} -ge 2 ]; then
+  $verbose && echo "stage 2: train model"
+  cd "${work_dir}" || exit 1
+  python3 step_2_train_model.py \
+  --train_dataset "${train_dataset}" \
+  --valid_dataset "${valid_dataset}" \
+  --serialization_dir "${file_dir}" \
+  --config_file "${config_file}" \
+
+fi
+
+
+if [ ${stage} -le 3 ] && [ ${stop_stage} -ge 3 ]; then
+  $verbose && echo "stage 3: test model"
+  cd "${work_dir}" || exit 1
+  python3 step_3_evaluation.py \
+  --valid_dataset "${valid_dataset}" \
+  --model_dir "${file_dir}/best" \
+  --evaluation_audio_dir "${evaluation_audio_dir}" \
+  --limit "${limit}" \
+
+fi
+
+
+if [ ${stage} -le 4 ] && [ ${stop_stage} -ge 4 ]; then
+  $verbose && echo "stage 4: export model"
+  cd "${work_dir}" || exit 1
+  python3 step_5_export_models.py \
+  --vocabulary_dir "${vocabulary_dir}" \
+  --model_dir "${file_dir}/best" \
+  --serialization_dir "${file_dir}" \
+
+fi
+
+
+if [ ${stage} -le 5 ] && [ ${stop_stage} -ge 5 ]; then
+  $verbose && echo "stage 5: collect files"
+  cd "${work_dir}" || exit 1
+
+  mkdir -p ${final_model_dir}
+
+  cp "${file_dir}/best"/* "${final_model_dir}"
+  cp -r "${file_dir}/vocabulary" "${final_model_dir}"
+
+  cp "${file_dir}/evaluation.xlsx" "${final_model_dir}/evaluation.xlsx"
+
+  cp "${file_dir}/trace_model.zip" "${final_model_dir}/trace_model.zip"
+  cp "${file_dir}/trace_quant_model.zip" "${final_model_dir}/trace_quant_model.zip"
+  cp "${file_dir}/script_model.zip" "${final_model_dir}/script_model.zip"
+  cp "${file_dir}/script_quant_model.zip" "${final_model_dir}/script_quant_model.zip"
+
+  cd "${final_model_dir}/.." || exit 1;
+
+  if [ -e "${final_model_name}.zip" ]; then
+    rm -rf "${final_model_name}_backup.zip"
+    mv "${final_model_name}.zip" "${final_model_name}_backup.zip"
+  fi
+
+  zip -r "${final_model_name}.zip" "${final_model_name}"
+  rm -rf "${final_model_name}"
+
+fi
+
+
+if [ ${stage} -le 6 ] && [ ${stop_stage} -ge 6 ]; then
+  $verbose && echo "stage 6: clear file_dir"
+  cd "${work_dir}" || exit 1
+
+  rm -rf "${file_dir}";
+
+fi

examples/clean_unet/step_1_prepare_data.py

@@ -0,0 +1,201 @@
+#!/usr/bin/python3
+# -*- coding: utf-8 -*-
+import argparse
+import os
+from pathlib import Path
+import random
+import sys
+import shutil
+
+pwd = os.path.abspath(os.path.dirname(__file__))
+sys.path.append(os.path.join(pwd, "../../"))
+
+import pandas as pd
+from scipy.io import wavfile
+from tqdm import tqdm
+import librosa
+
+from project_settings import project_path
+
+
+def get_args():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--file_dir", default="./", type=str)
+
+    parser.add_argument(
+        "--noise_dir",
+        default=r"E:\Users\tianx\HuggingDatasets\nx_noise\data\noise",
+        type=str
+    )
+    parser.add_argument(
+        "--speech_dir",
+        default=r"E:\programmer\asr_datasets\aishell\data_aishell\wav\train",
+        type=str
+    )
+
+    parser.add_argument("--train_dataset", default="train.xlsx", type=str)
+    parser.add_argument("--valid_dataset", default="valid.xlsx", type=str)
+
+    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)
+
+    args = parser.parse_args()
+    return args
+
+
+def filename_generator(data_dir: str):
+    data_dir = Path(data_dir)
+    for filename in data_dir.glob("**/*.wav"):
+        yield filename.as_posix()
+
+
+def target_second_signal_generator(data_dir: str, duration: int = 2, sample_rate: int = 8000):
+    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
+
+
+def get_dataset(args):
+    file_dir = Path(args.file_dir)
+    file_dir.mkdir(exist_ok=True)
+
+    noise_dir = Path(args.noise_dir)
+    speech_dir = Path(args.speech_dir)
+
+    noise_generator = target_second_signal_generator(
+        noise_dir.as_posix(),
+        duration=args.duration,
+        sample_rate=args.target_sample_rate
+    )
+    speech_generator = target_second_signal_generator(
+        speech_dir.as_posix(),
+        duration=args.duration,
+        sample_rate=args.target_sample_rate
+    )
+
+    dataset = list()
+
+    count = 0
+    process_bar = tqdm(desc="build dataset excel")
+    for noise, speech in zip(noise_generator, speech_generator):
+        if count >= args.max_count:
+            break
+
+        noise_filename = noise["filename"]
+        noise_raw_duration = noise["raw_duration"]
+        noise_offset = noise["offset"]
+        noise_duration = noise["duration"]
+
+        speech_filename = speech["filename"]
+        speech_raw_duration = speech["raw_duration"]
+        speech_offset = speech["offset"]
+        speech_duration = speech["duration"]
+
+        random1 = random.random()
+        random2 = random.random()
+
+        row = {
+            "noise_filename": noise_filename,
+            "noise_raw_duration": noise_raw_duration,
+            "noise_offset": noise_offset,
+            "noise_duration": noise_duration,
+
+            "speech_filename": speech_filename,
+            "speech_raw_duration": speech_raw_duration,
+            "speech_offset": speech_offset,
+            "speech_duration": speech_duration,
+
+            "snr_db": random.uniform(args.min_snr_db, args.max_snr_db),
+
+            "random1": random1,
+            "random2": random2,
+            "flag": "TRAIN" if random2 < 0.8 else "TEST",
+        }
+        dataset.append(row)
+        count += 1
+        duration_seconds = count * args.duration
+        duration_hours = duration_seconds / 3600
+
+        process_bar.update(n=1)
+        process_bar.set_postfix({
+            # "duration_seconds": round(duration_seconds, 4),
+            "duration_hours": round(duration_hours, 4),
+
+        })
+
+    dataset = pd.DataFrame(dataset)
+    dataset = dataset.sort_values(by=["random1"], ascending=False)
+    dataset.to_excel(
+        file_dir / "dataset.xlsx",
+        index=False,
+    )
+    return
+
+
+
+def split_dataset(args):
+    """分割训练集, 测试集"""
+    file_dir = Path(args.file_dir)
+    file_dir.mkdir(exist_ok=True)
+
+    df = pd.read_excel(file_dir / "dataset.xlsx")
+
+    train = list()
+    test = list()
+
+    for i, row in df.iterrows():
+        flag = row["flag"]
+        if flag == "TRAIN":
+            train.append(row)
+        else:
+            test.append(row)
+
+    train = pd.DataFrame(train)
+    train.to_excel(
+        args.train_dataset,
+        index=False,
+        # encoding="utf_8_sig"
+    )
+    test = pd.DataFrame(test)
+    test.to_excel(
+        args.valid_dataset,
+        index=False,
+        # encoding="utf_8_sig"
+    )
+
+    return
+
+
+def main():
+    args = get_args()
+
+    get_dataset(args)
+    split_dataset(args)
+    return
+
+
+if __name__ == "__main__":
+    main()

examples/clean_unet/step_2_train_model.py

@@ -0,0 +1,419 @@
+#!/usr/bin/python3
+# -*- coding: utf-8 -*-
+"""
+https://github.com/NVIDIA/CleanUNet/blob/main/train.py
+
+https://github.com/NVIDIA/CleanUNet/blob/main/configs/DNS-large-full.json
+"""
+import argparse
+import json
+import logging
+from logging.handlers import TimedRotatingFileHandler
+import os
+import platform
+from pathlib import Path
+import random
+import sys
+import shutil
+from typing import List
+
+pwd = os.path.abspath(os.path.dirname(__file__))
+sys.path.append(os.path.join(pwd, "../../"))
+
+import numpy as np
+import torch
+import torch.nn as nn
+from torch.nn import functional as F
+from torch.utils.data.dataloader import DataLoader
+from tqdm import tqdm
+
+from toolbox.torch.utils.data.dataset.denoise_excel_dataset import DenoiseExcelDataset
+from toolbox.torchaudio.models.clean_unet.configuration_clean_unet import CleanUNetConfig
+from toolbox.torchaudio.models.clean_unet.modeling_clean_unet import CleanUNetPretrainedModel
+from toolbox.torchaudio.models.clean_unet.training import LinearWarmupCosineDecay
+from toolbox.torchaudio.models.clean_unet.loss import MultiResolutionSTFTLoss
+from toolbox.torchaudio.models.clean_unet.metrics import run_pesq_score
+
+torch.autograd.set_detect_anomaly(True)
+
+
+def get_args():
+    parser = argparse.ArgumentParser()
+    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("--batch_size", default=64, type=int)
+    parser.add_argument("--learning_rate", default=2e-4, type=float)
+    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)
+
+    args = parser.parse_args()
+    return args
+
+
+def logging_config(file_dir: str):
+    fmt = "%(asctime)s - %(name)s - %(levelname)s  %(filename)s:%(lineno)d >  %(message)s"
+
+    logging.basicConfig(format=fmt,
+                        datefmt="%m/%d/%Y %H:%M:%S",
+                        level=logging.INFO)
+    file_handler = TimedRotatingFileHandler(
+        filename=os.path.join(file_dir, "main.log"),
+        encoding="utf-8",
+        when="D",
+        interval=1,
+        backupCount=7
+    )
+    file_handler.setLevel(logging.INFO)
+    file_handler.setFormatter(logging.Formatter(fmt))
+    logger = logging.getLogger(__name__)
+    logger.addHandler(file_handler)
+
+    return logger
+
+
+class CollateFunction(object):
+    def __init__(self):
+        pass
+
+    def __call__(self, batch: List[dict]):
+        clean_audios = list()
+        noisy_audios = list()
+
+        for sample in batch:
+            # 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"]
+
+            clean_audios.append(clean_audio)
+            noisy_audios.append(noisy_audio)
+
+        clean_audios = torch.stack(clean_audios)
+        noisy_audios = torch.stack(noisy_audios)
+
+        # assert
+        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()
+
+
+def main():
+    args = get_args()
+
+    config = CleanUNetConfig.from_pretrained(
+        pretrained_model_name_or_path=args.config_file,
+    )
+
+    serialization_dir = Path(args.serialization_dir)
+    serialization_dir.mkdir(parents=True, exist_ok=True)
+
+    logger = logging_config(serialization_dir)
+
+    random.seed(args.seed)
+    np.random.seed(args.seed)
+    torch.manual_seed(args.seed)
+    logger.info(f"set seed: {args.seed}")
+
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    n_gpu = torch.cuda.device_count()
+    logger.info(f"GPU available count: {n_gpu}; device: {device}")
+
+    # datasets
+    train_dataset = DenoiseExcelDataset(
+        excel_file=args.train_dataset,
+        expected_sample_rate=8000,
+        max_wave_value=32768.0,
+    )
+    valid_dataset = DenoiseExcelDataset(
+        excel_file=args.valid_dataset,
+        expected_sample_rate=8000,
+        max_wave_value=32768.0,
+    )
+    train_data_loader = DataLoader(
+        dataset=train_dataset,
+        batch_size=args.batch_size,
+        shuffle=True,
+        sampler=None,
+        # Linux 系统中可以使用多个子进程加载数据, 而在 Windows 系统中不能.
+        num_workers=0 if platform.system() == "Windows" else os.cpu_count() // 2,
+        collate_fn=collate_fn,
+        pin_memory=False,
+        # prefetch_factor=64,
+    )
+    valid_data_loader = DataLoader(
+        dataset=valid_dataset,
+        batch_size=args.batch_size,
+        shuffle=True,
+        sampler=None,
+        # Linux 系统中可以使用多个子进程加载数据, 而在 Windows 系统中不能.
+        num_workers=0 if platform.system() == "Windows" else os.cpu_count() // 2,
+        collate_fn=collate_fn,
+        pin_memory=False,
+        # prefetch_factor=64,
+    )
+
+    # models
+    logger.info(f"prepare models. config_file: {args.config_file}")
+    model = CleanUNetPretrainedModel(config).to(device)
+
+    # optimizer
+    logger.info("prepare optimizer, lr_scheduler, loss_fn, categorical_accuracy")
+    optimizer = torch.optim.AdamW(model.parameters(), args.learning_rate)
+
+    # resume training
+    last_epoch = -1
+    for epoch_i in serialization_dir.glob("epoch-*"):
+        epoch_i = Path(epoch_i)
+        epoch_idx = epoch_i.stem.split("-")[1]
+        epoch_idx = int(epoch_idx)
+        if epoch_idx > last_epoch:
+            last_epoch = epoch_idx
+
+    if last_epoch != -1:
+        logger.info(f"resume from epoch-{last_epoch}.")
+        model_pt = serialization_dir / f"epoch-{last_epoch}/model.pt"
+        optimizer_pth = serialization_dir / f"epoch-{last_epoch}/optimizer.pth"
+
+        logger.info(f"load state dict for model.")
+        with open(model_pt.as_posix(), "rb") as f:
+            state_dict = torch.load(f, map_location="cpu", weights_only=True)
+        model.load_state_dict(state_dict, strict=True)
+
+        logger.info(f"load state dict for optimizer.")
+        with open(optimizer_pth.as_posix(), "rb") as f:
+            state_dict = torch.load(f, map_location="cpu", weights_only=True)
+        optimizer.load_state_dict(state_dict)
+
+    lr_scheduler = LinearWarmupCosineDecay(
+        optimizer,
+        lr_max=args.learning_rate,
+        n_iter=250000,
+        iteration=250000,
+        divider=25,
+        warmup_proportion=0.05,
+        phase=("linear", "cosine"),
+    )
+
+    # ae_loss_fn = nn.MSELoss(reduction="mean")
+    ae_loss_fn = nn.L1Loss(reduction="mean").to(device)
+
+    mr_stft_loss_fn = MultiResolutionSTFTLoss(
+        fft_sizes=[256, 512, 1024],
+        hop_sizes=[25, 50, 120],
+        win_lengths=[120, 240, 600],
+        sc_lambda=0.5,
+        mag_lambda=0.5,
+        band="full"
+    ).to(device)
+
+    # training loop
+
+    # state
+    average_pesq_score = 10000000000
+    average_loss = 10000000000
+    average_ae_loss = 10000000000
+    average_sc_loss = 10000000000
+    average_mag_loss = 10000000000
+
+    model_list = list()
+    best_idx_epoch = None
+    best_metric = None
+    patience_count = 0
+
+    logger.info("training")
+    for idx_epoch in range(max(0, last_epoch+1), args.max_epochs):
+        # train
+        model.train()
+
+        total_pesq_score = 0.
+        total_loss = 0.
+        total_ae_loss = 0.
+        total_sc_loss = 0.
+        total_mag_loss = 0.
+        total_batches = 0.
+
+        progress_bar = tqdm(
+            total=len(train_data_loader),
+            desc="Training; epoch: {}".format(idx_epoch),
+        )
+        for batch in train_data_loader:
+            clean_audios, noisy_audios = batch
+            clean_audios = clean_audios.to(device)
+            noisy_audios = noisy_audios.to(device)
+
+            enhanced_audios = model.forward(noisy_audios)
+            enhanced_audios = torch.squeeze(enhanced_audios, dim=1)
+
+            ae_loss = ae_loss_fn(enhanced_audios, clean_audios)
+            sc_loss, mag_loss = mr_stft_loss_fn(enhanced_audios, clean_audios)
+
+            loss = ae_loss + sc_loss + mag_loss
+
+            enhanced_audios_list_r = list(enhanced_audios.detach().cpu().numpy())
+            clean_audios_list_r = list(clean_audios.detach().cpu().numpy())
+            pesq_score = run_pesq_score(clean_audios_list_r, enhanced_audios_list_r, sample_rate=8000, mode="nb")
+
+            optimizer.zero_grad()
+            loss.backward()
+            optimizer.step()
+            lr_scheduler.step()
+
+            total_pesq_score += pesq_score
+            total_loss += loss.item()
+            total_ae_loss += ae_loss.item()
+            total_sc_loss += sc_loss.item()
+            total_mag_loss += mag_loss.item()
+            total_batches += 1
+
+            average_pesq_score = round(total_pesq_score / total_batches, 4)
+            average_loss = round(total_loss / total_batches, 4)
+            average_ae_loss = round(total_ae_loss / total_batches, 4)
+            average_sc_loss = round(total_sc_loss / total_batches, 4)
+            average_mag_loss = round(total_mag_loss / total_batches, 4)
+
+            progress_bar.update(1)
+            progress_bar.set_postfix({
+                "pesq_score": average_pesq_score,
+                "loss": average_loss,
+                "ae_loss": average_ae_loss,
+                "sc_loss": average_sc_loss,
+                "mag_loss": average_mag_loss,
+            })
+
+        # evaluation
+        model.eval()
+
+        torch.cuda.empty_cache()
+
+        total_pesq_score = 0.
+        total_loss = 0.
+        total_ae_loss = 0.
+        total_sc_loss = 0.
+        total_mag_loss = 0.
+        total_batches = 0.
+
+        progress_bar = tqdm(
+            total=len(valid_data_loader),
+            desc="Evaluation; epoch: {}".format(idx_epoch),
+        )
+        with torch.no_grad():
+            for batch in valid_data_loader:
+                clean_audios, noisy_audios = batch
+                clean_audios = clean_audios.to(device)
+                noisy_audios = noisy_audios.to(device)
+
+                enhanced_audios = model.forward(noisy_audios)
+                enhanced_audios = torch.squeeze(enhanced_audios, dim=1)
+
+                ae_loss = ae_loss_fn(enhanced_audios, clean_audios)
+                sc_loss, mag_loss = mr_stft_loss_fn(enhanced_audios, clean_audios)
+
+                loss = ae_loss + sc_loss + mag_loss
+
+                enhanced_audios_list_r = list(enhanced_audios.detach().cpu().numpy())
+                clean_audios_list_r = list(clean_audios.detach().cpu().numpy())
+                pesq_score = run_pesq_score(clean_audios_list_r, enhanced_audios_list_r, sample_rate=8000, mode="nb")
+
+                total_pesq_score += pesq_score
+                total_loss += loss.item()
+                total_ae_loss += ae_loss.item()
+                total_sc_loss += sc_loss.item()
+                total_mag_loss += mag_loss.item()
+                total_batches += 1
+
+                average_pesq_score = round(total_pesq_score / total_batches, 4)
+                average_loss = round(total_loss / total_batches, 4)
+                average_ae_loss = round(total_ae_loss / total_batches, 4)
+                average_sc_loss = round(total_sc_loss / total_batches, 4)
+                average_mag_loss = round(total_mag_loss / total_batches, 4)
+
+                progress_bar.update(1)
+                progress_bar.set_postfix({
+                    "pesq_score": average_pesq_score,
+                    "loss": average_loss,
+                    "ae_loss": average_ae_loss,
+                    "sc_loss": average_sc_loss,
+                    "mag_loss": average_mag_loss,
+                })
+
+        # scheduler
+        lr_scheduler.step()
+
+        # save path
+        epoch_dir = serialization_dir / "epoch-{}".format(idx_epoch)
+        epoch_dir.mkdir(parents=True, exist_ok=False)
+
+        # save models
+        model.save_pretrained(epoch_dir.as_posix())
+
+        model_list.append(epoch_dir)
+        if len(model_list) >= args.num_serialized_models_to_keep:
+            model_to_delete: Path = model_list.pop(0)
+            shutil.rmtree(model_to_delete.as_posix())
+
+        # save optim
+        torch.save(optimizer.state_dict(), (epoch_dir / "optimizer.pth").as_posix())
+
+        # save metric
+        if best_metric is None:
+            best_idx_epoch = idx_epoch
+            best_metric = average_pesq_score
+        elif average_pesq_score > best_metric:
+            # great is better.
+            best_idx_epoch = idx_epoch
+            best_metric = average_pesq_score
+        else:
+            pass
+
+        metrics = {
+            "idx_epoch": idx_epoch,
+            "best_idx_epoch": best_idx_epoch,
+
+            "pesq_score": average_pesq_score,
+            "loss": average_loss,
+            "ae_loss": average_ae_loss,
+            "sc_loss": average_sc_loss,
+            "mag_loss": average_mag_loss,
+
+        }
+        metrics_filename = epoch_dir / "metrics_epoch.json"
+        with open(metrics_filename, "w", encoding="utf-8") as f:
+            json.dump(metrics, f, indent=4, ensure_ascii=False)
+
+        # save best
+        best_dir = serialization_dir / "best"
+        if best_idx_epoch == idx_epoch:
+            if best_dir.exists():
+                shutil.rmtree(best_dir)
+            shutil.copytree(epoch_dir, best_dir)
+
+        # early stop
+        early_stop_flag = False
+        if best_idx_epoch == idx_epoch:
+            patience_count = 0
+        else:
+            patience_count += 1
+        if patience_count >= args.patience:
+            early_stop_flag = True
+
+        # early stop
+        if early_stop_flag:
+            break
+
+    return
+
+
+if __name__ == "__main__":
+    main()

examples/clean_unet/step_3_evaluation.py

@@ -0,0 +1,6 @@
+#!/usr/bin/python3
+# -*- coding: utf-8 -*-
+
+
+if __name__ == '__main__':
+    pass

examples/clean_unet/yaml/config.yaml

@@ -0,0 +1,14 @@
+model_name: "clean_unet"
+
+channels_input: 1
+channels_output: 1
+channels_h: 64
+max_h: 768
+encoder_n_layers: 8
+kernel_size: 4
+stride: 2
+tsfm_n_layers: 5
+tsfm_n_head: 8
+tsfm_d_model: 512
+tsfm_d_inner: 2048
+

examples/conv_tasnet/run.sh

@@ -0,0 +1,154 @@
+#!/usr/bin/env bash
+
+: <<'END'
+
+
+sh run.sh --stage 2 --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
+
+
+END
+
+
+# params
+system_version="windows";
+verbose=true;
+stage=0 # start from 0 if you need to start from data preparation
+stop_stage=9
+
+work_dir="$(pwd)"
+file_folder_name=file_folder_name
+final_model_name=final_model_name
+config_file="yaml/config.yaml"
+limit=10
+
+noise_dir=/data/tianxing/HuggingDatasets/nx_noise/data/noise
+speech_dir=/data/tianxing/HuggingDatasets/aishell/data_aishell/wav/train
+
+max_count=10000000
+
+nohup_name=nohup.out
+
+# model params
+batch_size=64
+max_epochs=200
+save_top_k=10
+patience=5
+
+
+# parse options
+while true; do
+  [ -z "${1:-}" ] && break;  # break if there are no arguments
+  case "$1" in
+    --*) name=$(echo "$1" | sed s/^--// | sed s/-/_/g);
+      eval '[ -z "${'"$name"'+xxx}" ]' && echo "$0: invalid option $1" 1>&2 && exit 1;
+      old_value="(eval echo \\$$name)";
+      if [ "${old_value}" == "true" ] || [ "${old_value}" == "false" ]; then
+        was_bool=true;
+      else
+        was_bool=false;
+      fi
+
+      # Set the variable to the right value-- the escaped quotes make it work if
+      # the option had spaces, like --cmd "queue.pl -sync y"
+      eval "${name}=\"$2\"";
+
+      # Check that Boolean-valued arguments are really Boolean.
+      if $was_bool && [[ "$2" != "true" && "$2" != "false" ]]; then
+        echo "$0: expected \"true\" or \"false\": $1 $2" 1>&2
+        exit 1;
+      fi
+      shift 2;
+      ;;
+
+    *) break;
+  esac
+done
+
+file_dir="${work_dir}/${file_folder_name}"
+final_model_dir="${work_dir}/../../trained_models/${final_model_name}";
+evaluation_audio_dir="${file_dir}/evaluation_audio"
+
+train_dataset="${file_dir}/train.jsonl"
+valid_dataset="${file_dir}/valid.jsonl"
+
+$verbose && echo "system_version: ${system_version}"
+$verbose && echo "file_folder_name: ${file_folder_name}"
+
+if [ $system_version == "windows" ]; then
+  alias python3='D:/Users/tianx/PycharmProjects/virtualenv/nx_denoise/Scripts/python.exe'
+elif [ $system_version == "centos" ] || [ $system_version == "ubuntu" ]; then
+  #source /data/local/bin/nx_denoise/bin/activate
+  alias python3='/data/local/bin/nx_denoise/bin/python3'
+fi
+
+
+if [ ${stage} -le 1 ] && [ ${stop_stage} -ge 1 ]; then
+  $verbose && echo "stage 1: prepare data"
+  cd "${work_dir}" || exit 1
+  python3 step_1_prepare_data.py \
+  --file_dir "${file_dir}" \
+  --noise_dir "${noise_dir}" \
+  --speech_dir "${speech_dir}" \
+  --train_dataset "${train_dataset}" \
+  --valid_dataset "${valid_dataset}" \
+  --max_count "${max_count}" \
+
+fi
+
+
+if [ ${stage} -le 2 ] && [ ${stop_stage} -ge 2 ]; then
+  $verbose && echo "stage 2: train model"
+  cd "${work_dir}" || exit 1
+  python3 step_2_train_model.py \
+  --train_dataset "${train_dataset}" \
+  --valid_dataset "${valid_dataset}" \
+  --serialization_dir "${file_dir}" \
+  --config_file "${config_file}" \
+
+fi
+
+
+if [ ${stage} -le 3 ] && [ ${stop_stage} -ge 3 ]; then
+  $verbose && echo "stage 3: test model"
+  cd "${work_dir}" || exit 1
+  python3 step_3_evaluation.py \
+  --valid_dataset "${valid_dataset}" \
+  --model_dir "${file_dir}/best" \
+  --evaluation_audio_dir "${evaluation_audio_dir}" \
+  --limit "${limit}" \
+
+fi
+
+
+if [ ${stage} -le 4 ] && [ ${stop_stage} -ge 4 ]; then
+  $verbose && echo "stage 4: collect files"
+  cd "${work_dir}" || exit 1
+
+  mkdir -p ${final_model_dir}
+
+  cp "${file_dir}/best"/* "${final_model_dir}"
+  cp -r "${file_dir}/evaluation_audio" "${final_model_dir}"
+
+  cd "${final_model_dir}/.." || exit 1;
+
+  if [ -e "${final_model_name}.zip" ]; then
+    rm -rf "${final_model_name}_backup.zip"
+    mv "${final_model_name}.zip" "${final_model_name}_backup.zip"
+  fi
+
+  zip -r "${final_model_name}.zip" "${final_model_name}"
+  rm -rf "${final_model_name}"
+
+fi
+
+
+if [ ${stage} -le 5 ] && [ ${stop_stage} -ge 5 ]; then
+  $verbose && echo "stage 5: clear file_dir"
+  cd "${work_dir}" || exit 1
+
+  rm -rf "${file_dir}";
+
+fi

examples/conv_tasnet/step_1_prepare_data.py

@@ -0,0 +1,164 @@
+#!/usr/bin/python3
+# -*- coding: utf-8 -*-
+import argparse
+import json
+import os
+from pathlib import Path
+import random
+import sys
+
+pwd = os.path.abspath(os.path.dirname(__file__))
+sys.path.append(os.path.join(pwd, "../../"))
+
+import librosa
+import numpy as np
+from tqdm import tqdm
+
+
+def get_args():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--file_dir", default="./", type=str)
+
+    parser.add_argument(
+        "--noise_dir",
+        default=r"E:\Users\tianx\HuggingDatasets\nx_noise\data\noise",
+        type=str
+    )
+    parser.add_argument(
+        "--speech_dir",
+        default=r"E:\programmer\asr_datasets\aishell\data_aishell\wav\train",
+        type=str
+    )
+
+    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("--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)
+
+    args = parser.parse_args()
+    return args
+
+
+def filename_generator(data_dir: str):
+    data_dir = Path(data_dir)
+    for filename in data_dir.glob("**/*.wav"):
+        yield filename.as_posix()
+
+
+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 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):
+                if np.sum(signal[begin: begin+win_size]) == 0:
+                    continue
+                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 main():
+    args = get_args()
+
+    file_dir = Path(args.file_dir)
+    file_dir.mkdir(exist_ok=True)
+
+    noise_dir = Path(args.noise_dir)
+    speech_dir = Path(args.speech_dir)
+
+    noise_generator = target_second_signal_generator(
+        noise_dir.as_posix(),
+        duration=args.duration,
+        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,
+        max_epoch=1,
+    )
+
+    dataset = list()
+
+    count = 0
+    process_bar = tqdm(desc="build dataset excel")
+    with open(args.train_dataset, "w", encoding="utf-8") as ftrain, open(args.valid_dataset, "w", encoding="utf-8") as fvalid:
+        for noise, speech in zip(noise_generator, speech_generator):
+            if count >= args.max_count > 0:
+                break
+
+            noise_filename = noise["filename"]
+            noise_raw_duration = noise["raw_duration"]
+            noise_offset = noise["offset"]
+            noise_duration = noise["duration"]
+
+            speech_filename = speech["filename"]
+            speech_raw_duration = speech["raw_duration"]
+            speech_offset = speech["offset"]
+            speech_duration = speech["duration"]
+
+            random1 = random.random()
+            random2 = random.random()
+
+            row = {
+                "count": count,
+
+                "noise_filename": noise_filename,
+                "noise_raw_duration": noise_raw_duration,
+                "noise_offset": noise_offset,
+                "noise_duration": noise_duration,
+
+                "speech_filename": speech_filename,
+                "speech_raw_duration": speech_raw_duration,
+                "speech_offset": speech_offset,
+                "speech_duration": speech_duration,
+
+                "snr_db": random.uniform(args.min_snr_db, args.max_snr_db),
+
+                "random1": random1,
+            }
+            row = json.dumps(row, ensure_ascii=False)
+            if random2 < (1 / 300 / 1):
+                fvalid.write(f"{row}\n")
+            else:
+                ftrain.write(f"{row}\n")
+
+            count += 1
+            duration_seconds = count * args.duration
+            duration_hours = duration_seconds / 3600
+
+            process_bar.update(n=1)
+            process_bar.set_postfix({
+                # "duration_seconds": round(duration_seconds, 4),
+                "duration_hours": round(duration_hours, 4),
+
+            })
+
+    return
+
+
+if __name__ == "__main__":
+    main()

examples/conv_tasnet/step_2_train_model.py

@@ -0,0 +1,509 @@
+#!/usr/bin/python3
+# -*- coding: utf-8 -*-
+"""
+https://github.com/kaituoxu/Conv-TasNet/tree/master/src
+
+一般场景：
+
+目标 SI-SNR ≥ 10 dB，适用于电话通信、基础语音助手等。
+
+高要求场景（如医疗助听、语音识别）：
+需 SI-SNR ≥ 14 dB，并配合 PESQ ≥ 3.0 和 STOI ≥ 0.851812。
+
+DeepFilterNet2 模型在 DNS4 数据集，超过500小时的音频上训练了 100 个 epoch。
+https://arxiv.org/abs/2205.05474
+
+"""
+import argparse
+import json
+import logging
+from logging.handlers import TimedRotatingFileHandler
+import os
+import platform
+from pathlib import Path
+import random
+import sys
+import shutil
+from typing import List
+
+pwd = os.path.abspath(os.path.dirname(__file__))
+sys.path.append(os.path.join(pwd, "../../"))
+
+import numpy as np
+import torch
+import torch.nn as nn
+from torch.nn import functional as F
+from torch.utils.data.dataloader import DataLoader
+from tqdm import tqdm
+
+from toolbox.torch.utils.data.dataset.denoise_jsonl_dataset import DenoiseJsonlDataset
+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, MultiResolutionSTFTLoss
+from toolbox.torchaudio.losses.perceptual import NegSTOILoss, PesqLoss
+from toolbox.torchaudio.metrics.pesq import run_pesq_score
+
+
+def get_args():
+    parser = argparse.ArgumentParser()
+    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=200, type=int)
+
+    parser.add_argument("--batch_size", default=8, type=int)
+    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=1234, type=int)
+
+    parser.add_argument("--config_file", default="config.yaml", type=str)
+
+    args = parser.parse_args()
+    return args
+
+
+def logging_config(file_dir: str):
+    fmt = "%(asctime)s - %(name)s - %(levelname)s  %(filename)s:%(lineno)d >  %(message)s"
+
+    logging.basicConfig(format=fmt,
+                        datefmt="%m/%d/%Y %H:%M:%S",
+                        level=logging.INFO)
+    file_handler = TimedRotatingFileHandler(
+        filename=os.path.join(file_dir, "main.log"),
+        encoding="utf-8",
+        when="D",
+        interval=1,
+        backupCount=7
+    )
+    file_handler.setLevel(logging.INFO)
+    file_handler.setFormatter(logging.Formatter(fmt))
+    logger = logging.getLogger(__name__)
+    logger.addHandler(file_handler)
+
+    return logger
+
+
+class CollateFunction(object):
+    def __init__(self):
+        pass
+
+    def __call__(self, batch: List[dict]):
+        clean_audios = list()
+        noisy_audios = list()
+
+        for sample in batch:
+            # 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"]
+
+            clean_audios.append(clean_audio)
+            noisy_audios.append(noisy_audio)
+
+        clean_audios = torch.stack(clean_audios)
+        noisy_audios = torch.stack(noisy_audios)
+
+        # assert
+        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()
+
+
+def main():
+    args = get_args()
+
+    config = ConvTasNetConfig.from_pretrained(
+        pretrained_model_name_or_path=args.config_file,
+    )
+
+    serialization_dir = Path(args.serialization_dir)
+    serialization_dir.mkdir(parents=True, exist_ok=True)
+
+    logger = logging_config(serialization_dir)
+
+    random.seed(args.seed)
+    np.random.seed(args.seed)
+    torch.manual_seed(args.seed)
+    logger.info(f"set seed: {args.seed}")
+
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    n_gpu = torch.cuda.device_count()
+    logger.info(f"GPU available count: {n_gpu}; device: {device}")
+
+    # datasets
+    train_dataset = DenoiseJsonlDataset(
+        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=225000,
+    )
+    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,
+    )
+    train_data_loader = DataLoader(
+        dataset=train_dataset,
+        batch_size=args.batch_size,
+        # shuffle=True,
+        sampler=None,
+        # Linux 系统中可以使用多个子进程加载数据, 而在 Windows 系统中不能.
+        num_workers=0 if platform.system() == "Windows" else os.cpu_count() // 2,
+        collate_fn=collate_fn,
+        pin_memory=False,
+        prefetch_factor=2,
+    )
+    valid_data_loader = DataLoader(
+        dataset=valid_dataset,
+        batch_size=args.batch_size,
+        # shuffle=True,
+        sampler=None,
+        # Linux 系统中可以使用多个子进程加载数据, 而在 Windows 系统中不能.
+        num_workers=0 if platform.system() == "Windows" else os.cpu_count() // 2,
+        collate_fn=collate_fn,
+        pin_memory=False,
+        prefetch_factor=2,
+    )
+
+    # models
+    logger.info(f"prepare models. config_file: {args.config_file}")
+    model = ConvTasNetPretrainedModel(config).to(device)
+    model.to(device)
+    model.train()
+
+    # optimizer
+    logger.info("prepare optimizer, lr_scheduler, loss_fn, categorical_accuracy")
+    optimizer = torch.optim.AdamW(model.parameters(), config.lr)
+
+    # resume training
+    last_step_idx = -1
+    last_epoch = -1
+    for step_idx_str in serialization_dir.glob("steps-*"):
+        step_idx_str = Path(step_idx_str)
+        step_idx = step_idx_str.stem.split("-")[1]
+        step_idx = int(step_idx)
+        if step_idx > last_step_idx:
+            last_step_idx = step_idx
+    last_epoch = 1
+
+    if last_step_idx != -1:
+        logger.info(f"resume from steps-{last_step_idx}.")
+        model_pt = serialization_dir / f"steps-{last_step_idx}/model.pt"
+        optimizer_pth = serialization_dir / f"steps-{last_step_idx}/optimizer.pth"
+
+        logger.info(f"load state dict for model.")
+        with open(model_pt.as_posix(), "rb") as f:
+            state_dict = torch.load(f, map_location="cpu", weights_only=True)
+        model.load_state_dict(state_dict, strict=True)
+
+        logger.info(f"load state dict for optimizer.")
+        with open(optimizer_pth.as_posix(), "rb") as f:
+            state_dict = torch.load(f, map_location="cpu", weights_only=True)
+        optimizer.load_state_dict(state_dict)
+
+    if config.lr_scheduler == "CosineAnnealingLR":
+        lr_scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(
+            optimizer,
+            last_epoch=last_epoch,
+            # T_max=10 * config.eval_steps,
+            # eta_min=0.01 * config.lr,
+            **config.lr_scheduler_kwargs,
+        )
+    elif config.lr_scheduler == "MultiStepLR":
+        lr_scheduler = torch.optim.lr_scheduler.MultiStepLR(
+            optimizer,
+            last_epoch=last_epoch,
+            milestones=[10000, 20000, 30000, 40000, 50000], gamma=0.5
+        )
+    else:
+        raise AssertionError(f"invalid lr_scheduler: {config.lr_scheduler}")
+
+    ae_loss_fn = nn.L1Loss(reduction="mean").to(device)
+    neg_si_snr_loss_fn = NegativeSISNRLoss(reduction="mean").to(device)
+    neg_stoi_loss_fn = NegSTOILoss(sample_rate=config.sample_rate, 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],
+        factor_sc=1.5,
+        factor_mag=1.0,
+        reduction="mean"
+    ).to(device)
+    pesq_loss_fn = PesqLoss(0.5, sample_rate=config.sample_rate).to(device)
+
+    # training loop
+
+    # state
+    average_pesq_score = 1000000000
+    average_loss = 1000000000
+    average_ae_loss = 1000000000
+    average_neg_si_snr_loss = 1000000000
+    average_neg_stoi_loss = 1000000000
+
+    model_list = list()
+    best_epoch_idx = None
+    best_step_idx = None
+    best_metric = None
+    patience_count = 0
+
+    step_idx = 0 if last_step_idx == -1 else last_step_idx
+
+    logger.info("training")
+    for epoch_idx in range(max(0, last_epoch+1), args.max_epochs):
+        # train
+        model.train()
+
+        total_pesq_score = 0.
+        total_loss = 0.
+        total_ae_loss = 0.
+        total_neg_si_snr_loss = 0.
+        total_neg_stoi_loss = 0.
+        total_mr_stft_loss = 0.
+        total_pesq_loss = 0.
+        total_batches = 0.
+
+        progress_bar_train = tqdm(
+            initial=step_idx,
+            desc="Training; epoch-{}".format(epoch_idx),
+        )
+        for train_batch in train_data_loader:
+            clean_audios, noisy_audios = train_batch
+            clean_audios: torch.Tensor = clean_audios.to(device)
+            noisy_audios: torch.Tensor = noisy_audios.to(device)
+
+            denoise_audios = model.forward(noisy_audios)
+            denoise_audios = torch.squeeze(denoise_audios, dim=1)
+
+            if torch.any(torch.isnan(denoise_audios)) or torch.any(torch.isinf(denoise_audios)):
+                raise AssertionError("nan or inf in denoise_audios")
+
+            ae_loss = ae_loss_fn.forward(denoise_audios, clean_audios)
+            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)
+            mr_stft_loss = mr_stft_loss_fn.forward(denoise_audios, clean_audios)
+            pesq_loss = pesq_loss_fn.forward(denoise_audios, clean_audios)
+
+            # loss = 0.25 * ae_loss + 0.25 * neg_si_snr_loss
+            # loss = 0.25 * ae_loss + 0.25 * neg_si_snr_loss + 0.25 * neg_stoi_loss + 0.25 * mr_stft_loss
+            # loss = 1.0 * ae_loss + 0.8 * neg_si_snr_loss + 0.5 * mr_stft_loss + 0.3 * neg_stoi_loss
+            # loss = 1.0 * ae_loss + 0.8 * neg_si_snr_loss + 0.7 * mr_stft_loss + 0.5 * neg_stoi_loss
+            # loss = 2.0 * mr_stft_loss + 0.8 * ae_loss + 0.7 * neg_si_snr_loss + 0.5 * neg_stoi_loss
+            # loss = 1.0 * ae_loss + 0.8 * neg_si_snr_loss + 0.7 * mr_stft_loss + 0.5 * neg_stoi_loss + 0.5 * pesq_loss
+            loss = 0.1 * ae_loss + 0.1 * neg_si_snr_loss + 1.0 * mr_stft_loss + 0.2 * neg_stoi_loss + 0.2 * pesq_loss
+            if torch.any(torch.isnan(loss)) or torch.any(torch.isinf(loss)):
+                logger.info(f"find nan or inf in loss.")
+                continue
+
+            denoise_audios_list_r = list(denoise_audios.detach().cpu().numpy())
+            clean_audios_list_r = list(clean_audios.detach().cpu().numpy())
+            pesq_score = run_pesq_score(clean_audios_list_r, denoise_audios_list_r, sample_rate=config.sample_rate, mode="nb")
+
+            optimizer.zero_grad()
+            loss.backward()
+            optimizer.step()
+            lr_scheduler.step()
+
+            total_pesq_score += pesq_score
+            total_loss += loss.item()
+            total_ae_loss += ae_loss.item()
+            total_neg_si_snr_loss += neg_si_snr_loss.item()
+            total_neg_stoi_loss += neg_stoi_loss.item()
+            total_mr_stft_loss += mr_stft_loss.item()
+            total_pesq_loss += pesq_loss.item()
+            total_batches += 1
+
+            average_pesq_score = round(total_pesq_score / total_batches, 4)
+            average_loss = round(total_loss / total_batches, 4)
+            average_ae_loss = round(total_ae_loss / total_batches, 4)
+            average_neg_si_snr_loss = round(total_neg_si_snr_loss / total_batches, 4)
+            average_neg_stoi_loss = round(total_neg_stoi_loss / total_batches, 4)
+            average_mr_stft_loss = round(total_mr_stft_loss / total_batches, 4)
+            average_pesq_loss = round(total_pesq_loss / total_batches, 4)
+
+            progress_bar_train.update(1)
+            progress_bar_train.set_postfix({
+                "lr": lr_scheduler.get_last_lr()[0],
+                "pesq_score": average_pesq_score,
+                "loss": average_loss,
+                "ae_loss": average_ae_loss,
+                "neg_si_snr_loss": average_neg_si_snr_loss,
+                "neg_stoi_loss": average_neg_stoi_loss,
+                "mr_stft_loss": average_mr_stft_loss,
+                "pesq_loss": average_pesq_loss,
+            })
+
+            # evaluation
+            step_idx += 1
+            if step_idx % config.eval_steps == 0:
+                model.eval()
+                with torch.no_grad():
+                    torch.cuda.empty_cache()
+
+                    total_pesq_score = 0.
+                    total_loss = 0.
+                    total_ae_loss = 0.
+                    total_neg_si_snr_loss = 0.
+                    total_neg_stoi_loss = 0.
+                    total_mr_stft_loss = 0.
+                    total_pesq_loss = 0.
+                    total_batches = 0.
+
+                    progress_bar_train.close()
+                    progress_bar_eval = tqdm(
+                        desc="Evaluation; steps-{}k".format(int(step_idx/1000)),
+                    )
+                    for eval_batch in valid_data_loader:
+                        clean_audios, noisy_audios = eval_batch
+                        clean_audios = clean_audios.to(device)
+                        noisy_audios = noisy_audios.to(device)
+
+                        denoise_audios = model.forward(noisy_audios)
+                        denoise_audios = torch.squeeze(denoise_audios, dim=1)
+
+                        ae_loss = ae_loss_fn.forward(denoise_audios, clean_audios)
+                        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)
+                        mr_stft_loss = mr_stft_loss_fn.forward(denoise_audios, clean_audios)
+                        pesq_loss = pesq_loss_fn.forward(denoise_audios, clean_audios)
+
+                        # loss = 0.25 * ae_loss + 0.25 * neg_si_snr_loss
+                        # loss = 0.25 * ae_loss + 0.25 * neg_si_snr_loss + 0.25 * neg_stoi_loss + 0.25 * mr_stft_loss
+                        # loss = 1.0 * ae_loss + 0.8 * neg_si_snr_loss + 0.5 * mr_stft_loss + 0.3 * neg_stoi_loss
+                        # loss = 1.0 * ae_loss + 0.8 * neg_si_snr_loss + 0.7 * mr_stft_loss + 0.5 * neg_stoi_loss
+                        # loss = 2.0 * mr_stft_loss + 0.8 * ae_loss + 0.7 * neg_si_snr_loss + 0.5 * neg_stoi_loss
+                        # loss = 1.0 * ae_loss + 0.8 * neg_si_snr_loss + 0.7 * mr_stft_loss + 0.5 * neg_stoi_loss + 0.5 * pesq_loss
+                        loss = 0.1 * ae_loss + 0.1 * neg_si_snr_loss + 1.0 * mr_stft_loss + 0.2 * neg_stoi_loss + 0.2 * pesq_loss
+                        if torch.any(torch.isnan(loss)) or torch.any(torch.isinf(loss)):
+                            logger.info(f"find nan or inf in loss.")
+                            continue
+
+                        denoise_audios_list_r = list(denoise_audios.detach().cpu().numpy())
+                        clean_audios_list_r = list(clean_audios.detach().cpu().numpy())
+                        pesq_score = run_pesq_score(clean_audios_list_r, denoise_audios_list_r, sample_rate=config.sample_rate, mode="nb")
+
+                        total_pesq_score += pesq_score
+                        total_loss += loss.item()
+                        total_ae_loss += ae_loss.item()
+                        total_neg_si_snr_loss += neg_si_snr_loss.item()
+                        total_neg_stoi_loss += neg_stoi_loss.item()
+                        total_mr_stft_loss += mr_stft_loss.item()
+                        total_pesq_loss += pesq_loss.item()
+                        total_batches += 1
+
+                        average_pesq_score = round(total_pesq_score / total_batches, 4)
+                        average_loss = round(total_loss / total_batches, 4)
+                        average_ae_loss = round(total_ae_loss / total_batches, 4)
+                        average_neg_si_snr_loss = round(total_neg_si_snr_loss / total_batches, 4)
+                        average_neg_stoi_loss = round(total_neg_stoi_loss / total_batches, 4)
+                        average_mr_stft_loss = round(total_mr_stft_loss / total_batches, 4)
+                        average_pesq_loss = round(total_pesq_loss / total_batches, 4)
+
+                        progress_bar_eval.update(1)
+                        progress_bar_eval.set_postfix({
+                            "lr": lr_scheduler.get_last_lr()[0],
+                            "pesq_score": average_pesq_score,
+                            "loss": average_loss,
+                            "ae_loss": average_ae_loss,
+                            "neg_si_snr_loss": average_neg_si_snr_loss,
+                            "neg_stoi_loss": average_neg_stoi_loss,
+                            "mr_stft_loss": average_mr_stft_loss,
+                            "pesq_loss": average_pesq_loss,
+                        })
+
+                    total_pesq_score = 0.
+                    total_loss = 0.
+                    total_ae_loss = 0.
+                    total_neg_si_snr_loss = 0.
+                    total_neg_stoi_loss = 0.
+                    total_mr_stft_loss = 0.
+                    total_pesq_loss = 0.
+                    total_batches = 0.
+
+                    progress_bar_eval.close()
+                    progress_bar_train = tqdm(
+                        initial=progress_bar_train.n,
+                        postfix=progress_bar_train.postfix,
+                        desc=progress_bar_train.desc,
+                    )
+
+                    # save path
+                    save_dir = serialization_dir / "steps-{}".format(step_idx)
+                    save_dir.mkdir(parents=True, exist_ok=False)
+
+                    # save models
+                    model.save_pretrained(save_dir.as_posix())
+
+                    model_list.append(save_dir)
+                    if len(model_list) >= args.num_serialized_models_to_keep:
+                        model_to_delete: Path = model_list.pop(0)
+                        shutil.rmtree(model_to_delete.as_posix())
+
+                    # save optim
+                    torch.save(optimizer.state_dict(), (save_dir / "optimizer.pth").as_posix())
+
+                    # save metric
+                    if best_metric is None:
+                        best_epoch_idx = epoch_idx
+                        best_step_idx = step_idx
+                        best_metric = average_pesq_score
+                    elif average_pesq_score > best_metric:
+                        # great is better.
+                        best_epoch_idx = epoch_idx
+                        best_step_idx = step_idx
+                        best_metric = average_pesq_score
+                    else:
+                        pass
+
+                    metrics = {
+                        "epoch_idx": epoch_idx,
+                        "best_epoch_idx": best_epoch_idx,
+                        "best_step_idx": best_step_idx,
+                        "pesq_score": average_pesq_score,
+                        "loss": average_loss,
+                        "ae_loss": average_ae_loss,
+                        "neg_si_snr_loss": average_neg_si_snr_loss,
+                        "neg_stoi_loss": average_neg_stoi_loss,
+                        "mr_stft_loss": average_mr_stft_loss,
+                        "pesq_loss": average_pesq_loss,
+                    }
+                    metrics_filename = save_dir / "metrics_epoch.json"
+                    with open(metrics_filename, "w", encoding="utf-8") as f:
+                        json.dump(metrics, f, indent=4, ensure_ascii=False)
+
+                    # save best
+                    best_dir = serialization_dir / "best"
+                    if best_epoch_idx == epoch_idx and best_step_idx == step_idx:
+                        if best_dir.exists():
+                            shutil.rmtree(best_dir)
+                        shutil.copytree(save_dir, best_dir)
+
+                    # early stop
+                    early_stop_flag = False
+                    if best_epoch_idx == epoch_idx and best_step_idx == step_idx:
+                        patience_count = 0
+                    else:
+                        patience_count += 1
+                    if patience_count >= args.patience:
+                        early_stop_flag = True
+
+                    # early stop
+                    if early_stop_flag:
+                        break
+                model.train()
+
+    return
+
+
+if __name__ == "__main__":
+    main()

examples/conv_tasnet/yaml/config.yaml

@@ -0,0 +1,28 @@
+model_name: "conv_tasnet"
+
+sample_rate: 8000
+segment_size: 4
+
+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
+
+norm_type: "gLN"
+causal: false
+mask_nonlinear: "relu"
+
+min_snr_db: -10
+max_snr_db: 20
+
+lr: 0.005
+lr_scheduler: "CosineAnnealingLR"
+lr_scheduler_kwargs:
+  T_max: 250000
+  eta_min: 0.00005
+
+eval_steps: 25000

examples/data_preprocess/dns_challenge_to_8k/process_dns_challenge_clean_emotional_speech.py

@@ -0,0 +1,90 @@
+#!/usr/bin/python3
+# -*- coding: utf-8 -*-
+"""
+https://github.com/microsoft/DNS-Challenge/blob/master/download-dns-challenge-3.sh
+
+1.2G
+wget https://dns3public.blob.core.windows.net/dns3archive/DEMAND.tar.bz2
+
+14G
+wget https://dns3public.blob.core.windows.net/dns3archive/datasets/datasets.noise.tar.bz2
+
+38G
+wget https://dns3public.blob.core.windows.net/dns3archive/datasets/datasets.clean.read_speech.tar.bz2
+
+247M
+wget https://dns3public.blob.core.windows.net/dns3archive/datasets/datasets.clean.emotional_speech.tar.bz2
+
+
+"""
+import argparse
+import os
+from pathlib import Path
+import sys
+
+import numpy as np
+from tqdm import tqdm
+
+pwd = os.path.abspath(os.path.dirname(__file__))
+sys.path.append(os.path.join(pwd, "../../"))
+
+import librosa
+from scipy.io import wavfile
+
+
+def get_args():
+    parser = argparse.ArgumentParser()
+
+    parser.add_argument(
+        "--data_dir",
+        default=r"E:\programmer\asr_datasets\dns-challenge\datasets.clean.emotional_speech\datasets\clean\emotional_speech",
+        type=str
+    )
+    parser.add_argument(
+        "--output_dir",
+        default=r"E:\programmer\asr_datasets\denoise\dns-clean-emotional-speech-8k",
+        type=str
+    )
+    parser.add_argument("--sample_rate", default=8000, type=int)
+    args = parser.parse_args()
+    return args
+
+
+def main():
+    args = get_args()
+
+    data_dir = Path(args.data_dir)
+    output_dir = Path(args.output_dir)
+    output_dir.mkdir(parents=True, exist_ok=True)
+
+    # finished_set
+    finished_set = set()
+    for filename in tqdm(output_dir.glob("**/*.wav")):
+        name = filename.stem
+        finished_set.add(name)
+    print(f"finished_set count: {len(finished_set)}")
+
+    for filename in tqdm(data_dir.glob("**/*.wav")):
+        label = filename.parts[-2]
+        name = filename.stem
+        # print(f"filename: {filename.as_posix()}")
+        if name in finished_set:
+            continue
+
+        signal, _ = librosa.load(filename.as_posix(), sr=args.sample_rate)
+
+        signal = signal * (1 << 15)
+        signal = np.array(signal, dtype=np.int16)
+
+        to_file = output_dir / f"{label}/{name}.wav"
+        to_file.parent.mkdir(parents=True, exist_ok=True)
+        wavfile.write(
+            to_file.as_posix(),
+            rate=args.sample_rate,
+            data=signal,
+        )
+    return
+
+
+if __name__ == "__main__":
+    main()

examples/data_preprocess/dns_challenge_to_8k/process_dns_challenge_clean_read_speech.py

@@ -0,0 +1,129 @@
+#!/usr/bin/python3
+# -*- coding: utf-8 -*-
+"""
+https://github.com/microsoft/DNS-Challenge/blob/master/download-dns-challenge-3.sh
+
+1.2G
+wget https://dns3public.blob.core.windows.net/dns3archive/DEMAND.tar.bz2
+
+14G
+wget https://dns3public.blob.core.windows.net/dns3archive/datasets/datasets.noise.tar.bz2
+
+38G
+wget https://dns3public.blob.core.windows.net/dns3archive/datasets/datasets.clean.read_speech.tar.bz2
+
+12G
+wget https://dns3public.blob.core.windows.net/dns3archive/datasets/datasets.clean.french_data.tar.bz2
+
+43G
+wget https://dns3public.blob.core.windows.net/dns3archive/datasets/datasets.clean.german_speech.tar.bz2
+
+7.9G
+wget https://dns3public.blob.core.windows.net/dns3archive/datasets/datasets.clean.italian_speech.tar.bz2
+
+12G
+wget https://dns3public.blob.core.windows.net/dns3archive/datasets/datasets.clean.mandarin_speech.tar.bz2
+
+3.1G
+wget https://dns3public.blob.core.windows.net/dns3archive/datasets/datasets.clean.russian_speech.tar.bz2
+
+9.7G
+wget https://dns3public.blob.core.windows.net/dns3archive/datasets/datasets.clean.spanish_speech.tar.bz2
+
+617M
+wget https://dns3public.blob.core.windows.net/dns3archive/datasets/datasets.clean.singing_voice.tar.bz2
+
+"""
+import argparse
+import os
+from pathlib import Path
+import sys
+
+import numpy as np
+from tqdm import tqdm
+
+pwd = os.path.abspath(os.path.dirname(__file__))
+sys.path.append(os.path.join(pwd, "../../"))
+
+import librosa
+from scipy.io import wavfile
+
+
+def get_args():
+    parser = argparse.ArgumentParser()
+
+    parser.add_argument(
+        "--data_dir",
+        # 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.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.spanish_speech\datasets\clean\spanish_speech",
+        type=str
+    )
+    parser.add_argument(
+        "--output_dir",
+        # 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-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-spanish-speech-8k",
+        type=str
+    )
+    parser.add_argument("--sample_rate", default=8000, type=int)
+    args = parser.parse_args()
+    return args
+
+
+def main():
+    args = get_args()
+
+    data_dir = Path(args.data_dir)
+    output_dir = Path(args.output_dir)
+    output_dir.mkdir(parents=True, exist_ok=True)
+
+    # finished_set
+    finished_set = set()
+    for filename in tqdm(output_dir.glob("**/*.wav")):
+        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")):
+        relative_name = filename.relative_to(data_dir)
+        relative_name_ = relative_name.as_posix()
+        if relative_name_ in finished_set:
+            continue
+        finished_set.add(relative_name_)
+
+        try:
+            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)
+
+        to_file = output_dir / relative_name.as_posix()
+        to_file.parent.mkdir(parents=True, exist_ok=True)
+        wavfile.write(
+            to_file.as_posix(),
+            rate=args.sample_rate,
+            data=signal,
+        )
+    return
+
+
+if __name__ == "__main__":
+    main()

examples/data_preprocess/dns_challenge_to_8k/process_dns_challenge_demand.py

@@ -0,0 +1,71 @@
+#!/usr/bin/python3
+# -*- coding: utf-8 -*-
+"""
+https://github.com/microsoft/DNS-Challenge/blob/master/download-dns-challenge-3.sh
+
+1.2G
+wget https://dns3public.blob.core.windows.net/dns3archive/DEMAND.tar.bz2
+
+"""
+import argparse
+import os
+from pathlib import Path
+import random
+import sys
+import shutil
+
+import numpy as np
+
+pwd = os.path.abspath(os.path.dirname(__file__))
+sys.path.append(os.path.join(pwd, "../../"))
+
+import librosa
+from scipy.io import wavfile
+
+
+def get_args():
+    parser = argparse.ArgumentParser()
+
+    parser.add_argument(
+        "--data_dir",
+        default=r"E:\programmer\asr_datasets\dns-challenge\DEMAND\demand",
+        type=str
+    )
+    parser.add_argument(
+        "--output_dir",
+        default=r"E:\programmer\asr_datasets\denoise\demand-8k",
+        type=str
+    )
+    parser.add_argument("--sample_rate", default=8000, type=int)
+    args = parser.parse_args()
+    return args
+
+
+def main():
+    args = get_args()
+
+    data_dir = Path(args.data_dir)
+    output_dir = Path(args.output_dir)
+    output_dir.mkdir(parents=True, exist_ok=False)
+
+    for filename in data_dir.glob("**/ch01.wav"):
+        label = filename.parts[-2]
+        name = filename.stem
+
+        signal, _ = librosa.load(filename.as_posix(), sr=args.sample_rate)
+
+        signal = signal * (1 << 15)
+        signal = np.array(signal, dtype=np.int16)
+
+        to_file = output_dir / f"{label}/{name}.wav"
+        to_file.parent.mkdir(parents=True, exist_ok=True)
+        wavfile.write(
+            to_file.as_posix(),
+            rate=args.sample_rate,
+            data=signal,
+        )
+    return
+
+
+if __name__ == '__main__':
+    main()

examples/data_preprocess/dns_challenge_to_8k/process_dns_challenge_impulse_responses.py

@@ -0,0 +1,93 @@
+#!/usr/bin/python3
+# -*- coding: utf-8 -*-
+"""
+https://github.com/microsoft/DNS-Challenge/blob/master/download-dns-challenge-3.sh
+
+1.2G
+wget https://dns3public.blob.core.windows.net/dns3archive/DEMAND.tar.bz2
+
+14G
+wget https://dns3public.blob.core.windows.net/dns3archive/datasets/datasets.noise.tar.bz2
+
+38G
+wget https://dns3public.blob.core.windows.net/dns3archive/datasets/datasets.clean.read_speech.tar.bz2
+
+247M
+wget https://dns3public.blob.core.windows.net/dns3archive/datasets/datasets.clean.emotional_speech.tar.bz2
+
+240M
+wget https://dns3public.blob.core.windows.net/dns3archive/datasets/datasets.impulse_responses.tar.bz2
+
+
+"""
+import argparse
+import os
+from pathlib import Path
+import sys
+
+import numpy as np
+from tqdm import tqdm
+
+pwd = os.path.abspath(os.path.dirname(__file__))
+sys.path.append(os.path.join(pwd, "../../"))
+
+import librosa
+from scipy.io import wavfile
+
+
+def get_args():
+    parser = argparse.ArgumentParser()
+
+    parser.add_argument(
+        "--data_dir",
+        default=r"E:\programmer\asr_datasets\dns-challenge\datasets.clean.emotional_speech\datasets\clean\emotional_speech",
+        type=str
+    )
+    parser.add_argument(
+        "--output_dir",
+        default=r"E:\programmer\asr_datasets\denoise\dns-clean-emotional-speech-8k",
+        type=str
+    )
+    parser.add_argument("--sample_rate", default=8000, type=int)
+    args = parser.parse_args()
+    return args
+
+
+def main():
+    args = get_args()
+
+    data_dir = Path(args.data_dir)
+    output_dir = Path(args.output_dir)
+    output_dir.mkdir(parents=True, exist_ok=True)
+
+    # finished_set
+    finished_set = set()
+    for filename in tqdm(output_dir.glob("**/*.wav")):
+        name = filename.stem
+        finished_set.add(name)
+    print(f"finished_set count: {len(finished_set)}")
+
+    for filename in tqdm(data_dir.glob("**/*.wav")):
+        label = filename.parts[-2]
+        name = filename.stem
+        # print(f"filename: {filename.as_posix()}")
+        if name in finished_set:
+            continue
+
+        signal, _ = librosa.load(filename.as_posix(), sr=args.sample_rate)
+
+        signal = signal * (1 << 15)
+        signal = np.array(signal, dtype=np.int16)
+
+        to_file = output_dir / f"{label}/{name}.wav"
+        to_file.parent.mkdir(parents=True, exist_ok=True)
+        wavfile.write(
+            to_file.as_posix(),
+            rate=args.sample_rate,
+            data=signal,
+        )
+    return
+
+
+if __name__ == "__main__":
+    main()

examples/data_preprocess/dns_challenge_to_8k/process_dns_challenge_noise.py

@@ -0,0 +1,77 @@
+#!/usr/bin/python3
+# -*- coding: utf-8 -*-
+"""
+https://github.com/microsoft/DNS-Challenge/blob/master/download-dns-challenge-3.sh
+
+1.2G
+wget https://dns3public.blob.core.windows.net/dns3archive/DEMAND.tar.bz2
+
+14G
+wget https://dns3public.blob.core.windows.net/dns3archive/datasets/datasets.noise.tar.bz2
+
+38G
+wget https://dns3public.blob.core.windows.net/dns3archive/datasets/datasets.clean.read_speech.tar.bz2
+
+"""
+import argparse
+import os
+from pathlib import Path
+import sys
+
+import numpy as np
+from tqdm import tqdm
+
+pwd = os.path.abspath(os.path.dirname(__file__))
+sys.path.append(os.path.join(pwd, "../../"))
+
+import librosa
+from scipy.io import wavfile
+
+
+def get_args():
+    parser = argparse.ArgumentParser()
+
+    parser.add_argument(
+        "--data_dir",
+        default=r"E:\programmer\asr_datasets\dns-challenge\datasets.noise\datasets",
+        type=str
+    )
+    parser.add_argument(
+        "--output_dir",
+        default=r"E:\programmer\asr_datasets\denoise\dns-noise-8k",
+        type=str
+    )
+    parser.add_argument("--sample_rate", default=8000, type=int)
+    args = parser.parse_args()
+    return args
+
+
+def main():
+    args = get_args()
+
+    data_dir = Path(args.data_dir)
+    output_dir = Path(args.output_dir)
+    output_dir.mkdir(parents=True, exist_ok=True)
+
+    for filename in tqdm(data_dir.glob("**/*.wav")):
+        label = filename.parts[-2]
+        name = filename.stem
+        # print(f"filename: {filename.as_posix()}")
+
+        signal, _ = librosa.load(filename.as_posix(), sr=args.sample_rate)
+
+        signal = signal * (1 << 15)
+        signal = np.array(signal, dtype=np.int16)
+
+        to_file = output_dir / f"{label}/{name}.wav"
+        to_file.parent.mkdir(parents=True, exist_ok=True)
+        wavfile.write(
+            to_file.as_posix(),
+            rate=args.sample_rate,
+            data=signal,
+        )
+    return
+
+
+if __name__ == '__main__':
+    main()

examples/data_preprocess/dns_challenge_to_8k/process_musan.py

@@ -0,0 +1,8 @@
+#!/usr/bin/python3
+# -*- coding: utf-8 -*-
+"""
+https://www.openslr.org/17/
+"""
+
+if __name__ == '__main__':
+    pass

examples/data_preprocess/ms_snsd_to_8k/process_ms_snsd.py

@@ -0,0 +1,70 @@
+#!/usr/bin/python3
+# -*- coding: utf-8 -*-
+"""
+MS-SNSD
+https://github.com/microsoft/MS-SNSD
+"""
+import argparse
+import os
+from pathlib import Path
+import sys
+
+import numpy as np
+from tqdm import tqdm
+
+pwd = os.path.abspath(os.path.dirname(__file__))
+sys.path.append(os.path.join(pwd, "../../"))
+
+import librosa
+from scipy.io import wavfile
+
+
+def get_args():
+    parser = argparse.ArgumentParser()
+
+    parser.add_argument(
+        "--data_dir",
+        default=r"E:\programmer\asr_datasets\MS-SNSD",
+        type=str
+    )
+    parser.add_argument(
+        "--output_dir",
+        default=r"E:\programmer\asr_datasets\denoise\ms-snsd-noise-8k",
+        type=str
+    )
+    parser.add_argument("--sample_rate", default=8000, type=int)
+    args = parser.parse_args()
+    return args
+
+
+def main():
+    args = get_args()
+
+    data_dir = Path(args.data_dir)
+    output_dir = Path(args.output_dir)
+    output_dir.mkdir(parents=True, exist_ok=True)
+
+    for filename in tqdm(data_dir.glob("**/*.wav")):
+        label = filename.parts[-2]
+        name = filename.stem
+
+        if label not in ["noise_train", "noise_test", "clean_train", "clean_test"]:
+            continue
+
+        signal, _ = librosa.load(filename.as_posix(), sr=args.sample_rate)
+
+        signal = signal * (1 << 15)
+        signal = np.array(signal, dtype=np.int16)
+
+        to_file = output_dir / f"{label}/{name}.wav"
+        to_file.parent.mkdir(parents=True, exist_ok=True)
+        wavfile.write(
+            to_file.as_posix(),
+            rate=args.sample_rate,
+            data=signal,
+        )
+    return
+
+
+if __name__ == "__main__":
+    main()

examples/dfnet/run.sh

@@ -0,0 +1,156 @@
+#!/usr/bin/env bash
+
+: <<'END'
+
+sh run.sh --stage 2 --stop_stage 2 --system_version windows --file_folder_name file_dir --final_model_name dfnet-nx-speech \
+--noise_dir "E:/Users/tianx/HuggingDatasets/nx_noise/data/noise" \
+--speech_dir "E:/Users/tianx/HuggingDatasets/nx_noise/data/speech"
+
+sh run.sh --stage 2 --stop_stage 2 --system_version centos --file_folder_name file_dir --final_model_name dfnet-nx-dns3 \
+--noise_dir "/data/tianxing/HuggingDatasets/nx_noise/data/noise" \
+--speech_dir "/data/tianxing/HuggingDatasets/nx_noise/data/speech/dns3-speech"
+
+
+END
+
+
+# params
+system_version="windows";
+verbose=true;
+stage=0 # start from 0 if you need to start from data preparation
+stop_stage=9
+
+work_dir="$(pwd)"
+file_folder_name=file_folder_name
+final_model_name=final_model_name
+config_file="yaml/config.yaml"
+limit=10
+
+noise_dir=/data/tianxing/HuggingDatasets/nx_noise/data/noise
+speech_dir=/data/tianxing/HuggingDatasets/aishell/data_aishell/wav/train
+
+max_count=10000000
+
+nohup_name=nohup.out
+
+# model params
+batch_size=64
+max_epochs=200
+save_top_k=10
+patience=5
+
+
+# parse options
+while true; do
+  [ -z "${1:-}" ] && break;  # break if there are no arguments
+  case "$1" in
+    --*) name=$(echo "$1" | sed s/^--// | sed s/-/_/g);
+      eval '[ -z "${'"$name"'+xxx}" ]' && echo "$0: invalid option $1" 1>&2 && exit 1;
+      old_value="(eval echo \\$$name)";
+      if [ "${old_value}" == "true" ] || [ "${old_value}" == "false" ]; then
+        was_bool=true;
+      else
+        was_bool=false;
+      fi
+
+      # Set the variable to the right value-- the escaped quotes make it work if
+      # the option had spaces, like --cmd "queue.pl -sync y"
+      eval "${name}=\"$2\"";
+
+      # Check that Boolean-valued arguments are really Boolean.
+      if $was_bool && [[ "$2" != "true" && "$2" != "false" ]]; then
+        echo "$0: expected \"true\" or \"false\": $1 $2" 1>&2
+        exit 1;
+      fi
+      shift 2;
+      ;;
+
+    *) break;
+  esac
+done
+
+file_dir="${work_dir}/${file_folder_name}"
+final_model_dir="${work_dir}/../../trained_models/${final_model_name}";
+evaluation_audio_dir="${file_dir}/evaluation_audio"
+
+train_dataset="${file_dir}/train.jsonl"
+valid_dataset="${file_dir}/valid.jsonl"
+
+$verbose && echo "system_version: ${system_version}"
+$verbose && echo "file_folder_name: ${file_folder_name}"
+
+if [ $system_version == "windows" ]; then
+  alias python3='D:/Users/tianx/PycharmProjects/virtualenv/nx_denoise/Scripts/python.exe'
+elif [ $system_version == "centos" ] || [ $system_version == "ubuntu" ]; then
+  #source /data/local/bin/nx_denoise/bin/activate
+  alias python3='/data/local/bin/nx_denoise/bin/python3'
+fi
+
+
+if [ ${stage} -le 1 ] && [ ${stop_stage} -ge 1 ]; then
+  $verbose && echo "stage 1: prepare data"
+  cd "${work_dir}" || exit 1
+  python3 step_1_prepare_data.py \
+  --file_dir "${file_dir}" \
+  --noise_dir "${noise_dir}" \
+  --speech_dir "${speech_dir}" \
+  --train_dataset "${train_dataset}" \
+  --valid_dataset "${valid_dataset}" \
+  --max_count "${max_count}" \
+
+fi
+
+
+if [ ${stage} -le 2 ] && [ ${stop_stage} -ge 2 ]; then
+  $verbose && echo "stage 2: train model"
+  cd "${work_dir}" || exit 1
+  python3 step_2_train_model.py \
+  --train_dataset "${train_dataset}" \
+  --valid_dataset "${valid_dataset}" \
+  --serialization_dir "${file_dir}" \
+  --config_file "${config_file}" \
+
+fi
+
+
+if [ ${stage} -le 3 ] && [ ${stop_stage} -ge 3 ]; then
+  $verbose && echo "stage 3: test model"
+  cd "${work_dir}" || exit 1
+  python3 step_3_evaluation.py \
+  --valid_dataset "${valid_dataset}" \
+  --model_dir "${file_dir}/best" \
+  --evaluation_audio_dir "${evaluation_audio_dir}" \
+  --limit "${limit}" \
+
+fi
+
+
+if [ ${stage} -le 4 ] && [ ${stop_stage} -ge 4 ]; then
+  $verbose && echo "stage 4: collect files"
+  cd "${work_dir}" || exit 1
+
+  mkdir -p ${final_model_dir}
+
+  cp "${file_dir}/best"/* "${final_model_dir}"
+  cp -r "${file_dir}/evaluation_audio" "${final_model_dir}"
+
+  cd "${final_model_dir}/.." || exit 1;
+
+  if [ -e "${final_model_name}.zip" ]; then
+    rm -rf "${final_model_name}_backup.zip"
+    mv "${final_model_name}.zip" "${final_model_name}_backup.zip"
+  fi
+
+  zip -r "${final_model_name}.zip" "${final_model_name}"
+  rm -rf "${final_model_name}"
+
+fi
+
+
+if [ ${stage} -le 5 ] && [ ${stop_stage} -ge 5 ]; then
+  $verbose && echo "stage 5: clear file_dir"
+  cd "${work_dir}" || exit 1
+
+  rm -rf "${file_dir}";
+
+fi

examples/dfnet/step_1_prepare_data.py

@@ -0,0 +1,164 @@
+#!/usr/bin/python3
+# -*- coding: utf-8 -*-
+import argparse
+import json
+import os
+from pathlib import Path
+import random
+import sys
+
+pwd = os.path.abspath(os.path.dirname(__file__))
+sys.path.append(os.path.join(pwd, "../../"))
+
+import librosa
+import numpy as np
+from tqdm import tqdm
+
+
+def get_args():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--file_dir", default="./", type=str)
+
+    parser.add_argument(
+        "--noise_dir",
+        default=r"E:\Users\tianx\HuggingDatasets\nx_noise\data\noise",
+        type=str
+    )
+    parser.add_argument(
+        "--speech_dir",
+        default=r"E:\programmer\asr_datasets\aishell\data_aishell\wav\train",
+        type=str
+    )
+
+    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("--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)
+
+    args = parser.parse_args()
+    return args
+
+
+def filename_generator(data_dir: str):
+    data_dir = Path(data_dir)
+    for filename in data_dir.glob("**/*.wav"):
+        yield filename.as_posix()
+
+
+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 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):
+                if np.sum(signal[begin: begin+win_size]) == 0:
+                    continue
+                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 main():
+    args = get_args()
+
+    file_dir = Path(args.file_dir)
+    file_dir.mkdir(exist_ok=True)
+
+    noise_dir = Path(args.noise_dir)
+    speech_dir = Path(args.speech_dir)
+
+    noise_generator = target_second_signal_generator(
+        noise_dir.as_posix(),
+        duration=args.duration,
+        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,
+        max_epoch=1,
+    )
+
+    dataset = list()
+
+    count = 0
+    process_bar = tqdm(desc="build dataset jsonl")
+    with open(args.train_dataset, "w", encoding="utf-8") as ftrain, open(args.valid_dataset, "w", encoding="utf-8") as fvalid:
+        for noise, speech in zip(noise_generator, speech_generator):
+            if count >= args.max_count > 0:
+                break
+
+            noise_filename = noise["filename"]
+            noise_raw_duration = noise["raw_duration"]
+            noise_offset = noise["offset"]
+            noise_duration = noise["duration"]
+
+            speech_filename = speech["filename"]
+            speech_raw_duration = speech["raw_duration"]
+            speech_offset = speech["offset"]
+            speech_duration = speech["duration"]
+
+            random1 = random.random()
+            random2 = random.random()
+
+            row = {
+                "count": count,
+
+                "noise_filename": noise_filename,
+                "noise_raw_duration": noise_raw_duration,
+                "noise_offset": noise_offset,
+                "noise_duration": noise_duration,
+
+                "speech_filename": speech_filename,
+                "speech_raw_duration": speech_raw_duration,
+                "speech_offset": speech_offset,
+                "speech_duration": speech_duration,
+
+                "snr_db": random.uniform(args.min_snr_db, args.max_snr_db),
+
+                "random1": random1,
+            }
+            row = json.dumps(row, ensure_ascii=False)
+            if random2 < (1 / 300 / 1):
+                fvalid.write(f"{row}\n")
+            else:
+                ftrain.write(f"{row}\n")
+
+            count += 1
+            duration_seconds = count * args.duration
+            duration_hours = duration_seconds / 3600
+
+            process_bar.update(n=1)
+            process_bar.set_postfix({
+                # "duration_seconds": round(duration_seconds, 4),
+                "duration_hours": round(duration_hours, 4),
+
+            })
+
+    return
+
+
+if __name__ == "__main__":
+    main()

examples/dfnet/step_2_train_model.py

@@ -0,0 +1,461 @@
+#!/usr/bin/python3
+# -*- coding: utf-8 -*-
+"""
+https://github.com/Rikorose/DeepFilterNet
+"""
+import argparse
+import json
+import logging
+from logging.handlers import TimedRotatingFileHandler
+import os
+import platform
+from pathlib import Path
+import random
+import sys
+import shutil
+from typing import List
+
+from fontTools.varLib.plot import stops
+
+pwd = os.path.abspath(os.path.dirname(__file__))
+sys.path.append(os.path.join(pwd, "../../"))
+
+import numpy as np
+import torch
+import torch.nn as nn
+from torch.nn import functional as F
+from torch.utils.data.dataloader import DataLoader
+from tqdm import tqdm
+
+from toolbox.torch.utils.data.dataset.denoise_jsonl_dataset import DenoiseJsonlDataset
+from toolbox.torchaudio.losses.snr import NegativeSISNRLoss
+from toolbox.torchaudio.losses.spectral import LSDLoss, MultiResolutionSTFTLoss
+from toolbox.torchaudio.metrics.pesq import run_pesq_score
+from toolbox.torchaudio.models.dfnet.configuration_dfnet import DfNetConfig
+from toolbox.torchaudio.models.dfnet.modeling_dfnet import DfNet, DfNetPretrainedModel
+
+
+def get_args():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--train_dataset", default="train.jsonl", type=str)
+    parser.add_argument("--valid_dataset", default="valid.jsonl", type=str)
+
+    parser.add_argument("--num_serialized_models_to_keep", default=15, type=int)
+    parser.add_argument("--patience", default=10, type=int)
+    parser.add_argument("--serialization_dir", default="serialization_dir", type=str)
+
+    parser.add_argument("--config_file", default="config.yaml", type=str)
+
+    args = parser.parse_args()
+    return args
+
+
+def logging_config(file_dir: str):
+    fmt = "%(asctime)s - %(name)s - %(levelname)s  %(filename)s:%(lineno)d >  %(message)s"
+
+    logging.basicConfig(format=fmt,
+                        datefmt="%m/%d/%Y %H:%M:%S",
+                        level=logging.INFO)
+    file_handler = TimedRotatingFileHandler(
+        filename=os.path.join(file_dir, "main.log"),
+        encoding="utf-8",
+        when="D",
+        interval=1,
+        backupCount=7
+    )
+    file_handler.setLevel(logging.INFO)
+    file_handler.setFormatter(logging.Formatter(fmt))
+    logger = logging.getLogger(__name__)
+    logger.addHandler(file_handler)
+
+    return logger
+
+
+class CollateFunction(object):
+    def __init__(self):
+        pass
+
+    def __call__(self, batch: List[dict]):
+        clean_audios = list()
+        noisy_audios = list()
+        snr_db_list = list()
+
+        for sample in batch:
+            # 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"]
+
+            clean_audios.append(clean_audio)
+            noisy_audios.append(noisy_audio)
+
+        clean_audios = torch.stack(clean_audios)
+        noisy_audios = torch.stack(noisy_audios)
+
+        # assert
+        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()
+
+
+def main():
+    args = get_args()
+
+    config = DfNetConfig.from_pretrained(
+        pretrained_model_name_or_path=args.config_file,
+    )
+
+    serialization_dir = Path(args.serialization_dir)
+    serialization_dir.mkdir(parents=True, exist_ok=True)
+
+    logger = logging_config(serialization_dir)
+
+    random.seed(config.seed)
+    np.random.seed(config.seed)
+    torch.manual_seed(config.seed)
+    logger.info(f"set seed: {config.seed}")
+
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    n_gpu = torch.cuda.device_count()
+    logger.info(f"GPU available count: {n_gpu}; device: {device}")
+
+    # datasets
+    train_dataset = DenoiseJsonlDataset(
+        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=225000,
+    )
+    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,
+    )
+    train_data_loader = DataLoader(
+        dataset=train_dataset,
+        batch_size=config.batch_size,
+        # shuffle=True,
+        sampler=None,
+        # Linux 系统中可以使用多个子进程加载数据, 而在 Windows 系统中不能.
+        num_workers=0 if platform.system() == "Windows" else os.cpu_count() // 2,
+        collate_fn=collate_fn,
+        pin_memory=False,
+        prefetch_factor=None if platform.system() == "Windows" else 2,
+    )
+    valid_data_loader = DataLoader(
+        dataset=valid_dataset,
+        batch_size=config.batch_size,
+        # shuffle=True,
+        sampler=None,
+        # Linux 系统中可以使用多个子进程加载数据, 而在 Windows 系统中不能.
+        num_workers=0 if platform.system() == "Windows" else os.cpu_count() // 2,
+        collate_fn=collate_fn,
+        pin_memory=False,
+        prefetch_factor=None if platform.system() == "Windows" else 2,
+    )
+
+    # models
+    logger.info(f"prepare models. config_file: {args.config_file}")
+    model = DfNetPretrainedModel(config).to(device)
+    model.to(device)
+    model.train()
+
+    # optimizer
+    logger.info("prepare optimizer, lr_scheduler, loss_fn, evaluation_metric")
+    optimizer = torch.optim.AdamW(model.parameters(), config.lr)
+
+    # resume training
+    last_step_idx = -1
+    last_epoch = -1
+    for step_idx_str in serialization_dir.glob("steps-*"):
+        step_idx_str = Path(step_idx_str)
+        step_idx = step_idx_str.stem.split("-")[1]
+        step_idx = int(step_idx)
+        if step_idx > last_step_idx:
+            last_step_idx = step_idx
+    # last_epoch = 1
+
+    if last_step_idx != -1:
+        logger.info(f"resume from steps-{last_step_idx}.")
+        model_pt = serialization_dir / f"steps-{last_step_idx}/model.pt"
+
+        logger.info(f"load state dict for model.")
+        with open(model_pt.as_posix(), "rb") as f:
+            state_dict = torch.load(f, map_location="cpu", weights_only=True)
+        model.load_state_dict(state_dict, strict=True)
+
+    if config.lr_scheduler == "CosineAnnealingLR":
+        lr_scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(
+            optimizer,
+            last_epoch=last_epoch,
+            # T_max=10 * config.eval_steps,
+            # eta_min=0.01 * config.lr,
+            **config.lr_scheduler_kwargs,
+        )
+    elif config.lr_scheduler == "MultiStepLR":
+        lr_scheduler = torch.optim.lr_scheduler.MultiStepLR(
+            optimizer,
+            last_epoch=last_epoch,
+            milestones=[10000, 20000, 30000, 40000, 50000], gamma=0.5
+        )
+    else:
+        raise AssertionError(f"invalid lr_scheduler: {config.lr_scheduler}")
+
+    neg_si_snr_loss_fn = NegativeSISNRLoss(reduction="mean").to(device)
+    mr_stft_loss_fn = MultiResolutionSTFTLoss(
+        fft_size_list=[256, 512, 1024],
+        win_size_list=[256, 512, 1024],
+        hop_size_list=[128, 256, 512],
+        factor_sc=1.5,
+        factor_mag=1.0,
+        reduction="mean"
+    ).to(device)
+
+    # training loop
+
+    # state
+    average_pesq_score = 1000000000
+    average_loss = 1000000000
+    average_mr_stft_loss = 1000000000
+    average_neg_si_snr_loss = 1000000000
+    average_mask_loss = 1000000000
+    average_lsnr_loss = 1000000000
+
+    model_list = list()
+    best_epoch_idx = None
+    best_step_idx = None
+    best_metric = None
+    patience_count = 0
+
+    step_idx = 0 if last_step_idx == -1 else last_step_idx
+
+    logger.info("training")
+    early_stop_flag = False
+    for epoch_idx in range(max(0, last_epoch+1), config.max_epochs):
+        if early_stop_flag:
+            break
+
+        # train
+        model.train()
+
+        total_pesq_score = 0.
+        total_loss = 0.
+        total_mr_stft_loss = 0.
+        total_neg_si_snr_loss = 0.
+        total_mask_loss = 0.
+        total_lsnr_loss = 0.
+        total_batches = 0.
+
+        progress_bar_train = tqdm(
+            initial=step_idx,
+            desc="Training; epoch-{}".format(epoch_idx),
+        )
+        for train_batch in train_data_loader:
+            clean_audios, noisy_audios = train_batch
+            clean_audios: torch.Tensor = clean_audios.to(device)
+            noisy_audios: torch.Tensor = noisy_audios.to(device)
+
+            est_spec, est_wav, est_mask, lsnr = model.forward(noisy_audios)
+
+            mr_stft_loss = mr_stft_loss_fn.forward(est_wav, clean_audios)
+            neg_si_snr_loss = neg_si_snr_loss_fn.forward(est_wav, clean_audios)
+            mask_loss = model.mask_loss_fn(est_mask, clean_audios, noisy_audios)
+            lsnr_loss = model.lsnr_loss_fn(lsnr, clean_audios, noisy_audios)
+
+            loss = 1.0 * mr_stft_loss + 1.0 * neg_si_snr_loss + 1.0 * mask_loss + 0.3 * lsnr_loss
+            if torch.any(torch.isnan(loss)) or torch.any(torch.isinf(loss)):
+                logger.info(f"find nan or inf in loss.")
+                continue
+
+            denoise_audios_list_r = list(est_wav.detach().cpu().numpy())
+            clean_audios_list_r = list(clean_audios.detach().cpu().numpy())
+            pesq_score = run_pesq_score(clean_audios_list_r, denoise_audios_list_r, sample_rate=config.sample_rate, mode="nb")
+
+            optimizer.zero_grad()
+            loss.backward()
+            torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=config.clip_grad_norm)
+            optimizer.step()
+            lr_scheduler.step()
+
+            total_pesq_score += pesq_score
+            total_loss += loss.item()
+            total_mr_stft_loss += mr_stft_loss.item()
+            total_neg_si_snr_loss += neg_si_snr_loss.item()
+            total_mask_loss += mask_loss.item()
+            total_lsnr_loss += lsnr_loss.item()
+            total_batches += 1
+
+            average_pesq_score = round(total_pesq_score / total_batches, 4)
+            average_loss = round(total_loss / total_batches, 4)
+            average_mr_stft_loss = round(total_mr_stft_loss / total_batches, 4)
+            average_neg_si_snr_loss = round(total_neg_si_snr_loss / total_batches, 4)
+            average_mask_loss = round(total_mask_loss / total_batches, 4)
+            average_lsnr_loss = round(total_lsnr_loss / total_batches, 4)
+
+            progress_bar_train.update(1)
+            progress_bar_train.set_postfix({
+                "lr": lr_scheduler.get_last_lr()[0],
+                "pesq_score": average_pesq_score,
+                "loss": average_loss,
+                "mr_stft_loss": average_mr_stft_loss,
+                "neg_si_snr_loss": average_neg_si_snr_loss,
+                "mask_loss": average_mask_loss,
+                "lsnr_loss": average_lsnr_loss,
+            })
+
+            # evaluation
+            step_idx += 1
+            if step_idx % config.eval_steps == 0:
+                model.eval()
+                with torch.no_grad():
+                    torch.cuda.empty_cache()
+
+                    total_pesq_score = 0.
+                    total_loss = 0.
+                    total_mr_stft_loss = 0.
+                    total_neg_si_snr_loss = 0.
+                    total_mask_loss = 0.
+                    total_lsnr_loss = 0.
+                    total_batches = 0.
+
+                    progress_bar_train.close()
+                    progress_bar_eval = tqdm(
+                        desc="Evaluation; steps-{}k".format(int(step_idx/1000)),
+                    )
+                    for eval_batch in valid_data_loader:
+                        clean_audios, noisy_audios = eval_batch
+                        clean_audios: torch.Tensor = clean_audios.to(device)
+                        noisy_audios: torch.Tensor = noisy_audios.to(device)
+
+                        est_spec, est_wav, est_mask, lsnr = model.forward(noisy_audios)
+
+                        mr_stft_loss = mr_stft_loss_fn.forward(est_wav, clean_audios)
+                        neg_si_snr_loss = neg_si_snr_loss_fn.forward(est_wav, clean_audios)
+                        mask_loss = model.mask_loss_fn(est_mask, clean_audios, noisy_audios)
+                        lsnr_loss = model.lsnr_loss_fn(lsnr, clean_audios, noisy_audios)
+
+                        loss = 1.0 * mr_stft_loss + 1.0 * neg_si_snr_loss + 1.0 * mask_loss + 0.3 * lsnr_loss
+                        if torch.any(torch.isnan(loss)) or torch.any(torch.isinf(loss)):
+                            logger.info(f"find nan or inf in loss.")
+                            continue
+
+                        denoise_audios_list_r = list(est_wav.detach().cpu().numpy())
+                        clean_audios_list_r = list(clean_audios.detach().cpu().numpy())
+                        pesq_score = run_pesq_score(clean_audios_list_r, denoise_audios_list_r, sample_rate=config.sample_rate, mode="nb")
+
+                        total_pesq_score += pesq_score
+                        total_loss += loss.item()
+                        total_mr_stft_loss += mr_stft_loss.item()
+                        total_neg_si_snr_loss += neg_si_snr_loss.item()
+                        total_mask_loss += mask_loss.item()
+                        total_lsnr_loss += lsnr_loss.item()
+                        total_batches += 1
+
+                        average_pesq_score = round(total_pesq_score / total_batches, 4)
+                        average_loss = round(total_loss / total_batches, 4)
+                        average_mr_stft_loss = round(total_mr_stft_loss / total_batches, 4)
+                        average_neg_si_snr_loss = round(total_neg_si_snr_loss / total_batches, 4)
+                        average_mask_loss = round(total_mask_loss / total_batches, 4)
+                        average_lsnr_loss = round(total_lsnr_loss / total_batches, 4)
+
+                        progress_bar_eval.update(1)
+                        progress_bar_eval.set_postfix({
+                            "lr": lr_scheduler.get_last_lr()[0],
+                            "pesq_score": average_pesq_score,
+                            "loss": average_loss,
+                            "mr_stft_loss": average_mr_stft_loss,
+                            "neg_si_snr_loss": average_neg_si_snr_loss,
+                            "mask_loss": average_mask_loss,
+                            "lsnr_loss": average_lsnr_loss,
+                        })
+
+                    total_pesq_score = 0.
+                    total_loss = 0.
+                    total_mr_stft_loss = 0.
+                    total_neg_si_snr_loss = 0.
+                    total_mask_loss = 0.
+                    total_lsnr_loss = 0.
+                    total_batches = 0.
+
+                    progress_bar_eval.close()
+                    progress_bar_train = tqdm(
+                        initial=progress_bar_train.n,
+                        postfix=progress_bar_train.postfix,
+                        desc=progress_bar_train.desc,
+                    )
+
+                    # save path
+                    save_dir = serialization_dir / "steps-{}".format(step_idx)
+                    save_dir.mkdir(parents=True, exist_ok=False)
+
+                    # save models
+                    model.save_pretrained(save_dir.as_posix())
+
+                    model_list.append(save_dir)
+                    if len(model_list) >= args.num_serialized_models_to_keep:
+                        model_to_delete: Path = model_list.pop(0)
+                        shutil.rmtree(model_to_delete.as_posix())
+
+                    # save metric
+                    if best_metric is None:
+                        best_epoch_idx = epoch_idx
+                        best_step_idx = step_idx
+                        best_metric = average_pesq_score
+                    elif average_pesq_score >= best_metric:
+                        # great is better.
+                        best_epoch_idx = epoch_idx
+                        best_step_idx = step_idx
+                        best_metric = average_pesq_score
+                    else:
+                        pass
+
+                    metrics = {
+                        "epoch_idx": epoch_idx,
+                        "best_epoch_idx": best_epoch_idx,
+                        "best_step_idx": best_step_idx,
+                        "pesq_score": average_pesq_score,
+                        "loss": average_loss,
+                        "mr_stft_loss": average_mr_stft_loss,
+                        "neg_si_snr_loss": average_neg_si_snr_loss,
+                        "mask_loss": average_mask_loss,
+                        "lsnr_loss": average_lsnr_loss,
+                    }
+                    metrics_filename = save_dir / "metrics_epoch.json"
+                    with open(metrics_filename, "w", encoding="utf-8") as f:
+                        json.dump(metrics, f, indent=4, ensure_ascii=False)
+
+                    # save best
+                    best_dir = serialization_dir / "best"
+                    if best_epoch_idx == epoch_idx and best_step_idx == step_idx:
+                        if best_dir.exists():
+                            shutil.rmtree(best_dir)
+                        shutil.copytree(save_dir, best_dir)
+
+                    # early stop
+                    early_stop_flag = False
+                    if best_epoch_idx == epoch_idx and best_step_idx == step_idx:
+                        patience_count = 0
+                    else:
+                        patience_count += 1
+                    if patience_count >= args.patience:
+                        early_stop_flag = True
+
+                    # early stop
+                    if early_stop_flag:
+                        break
+                model.train()
+
+    return
+
+
+if __name__ == "__main__":
+    main()

examples/dfnet/yaml/config.yaml

@@ -0,0 +1,74 @@
+model_name: "dfnet"
+
+# spec
+sample_rate: 8000
+nfft: 512
+win_size: 200
+hop_size: 80
+
+spec_bins: 256
+
+# model
+conv_channels: 64
+conv_kernel_size_input:
+  - 3
+  - 3
+conv_kernel_size_inner:
+  - 1
+  - 3
+conv_lookahead: 0
+
+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: 64
+eval_steps: 10000
+
+# runtime
+use_post_filter: true

examples/dfnet2/run.sh

@@ -0,0 +1,164 @@
+#!/usr/bin/env bash
+
+: <<'END'
+
+sh run.sh --stage 2 --stop_stage 2 --system_version windows --file_folder_name file_dir --final_model_name dfnet-nx-speech \
+--noise_dir "E:/Users/tianx/HuggingDatasets/nx_noise/data/noise" \
+--speech_dir "E:/Users/tianx/HuggingDatasets/nx_noise/data/speech"
+
+sh run.sh --stage 2 --stop_stage 2 --system_version centos --file_folder_name file_dir --final_model_name dfnet2-nx-dns3 \
+--noise_dir "/data/tianxing/HuggingDatasets/nx_noise/data/noise" \
+--speech_dir "/data/tianxing/HuggingDatasets/nx_noise/data/speech/dns3-speech"
+
+sh run.sh --stage 2 --stop_stage 2 --system_version centos --file_folder_name file_dir --final_model_name dfnet2-nx2 \
+--noise_dir "/data/tianxing/HuggingDatasets/nx_noise/data/noise/nx-noise" \
+--speech_dir "/data/tianxing/HuggingDatasets/nx_noise/data/speech/nx-speech2"
+
+sh run.sh --stage 2 --stop_stage 2 --system_version centos --file_folder_name dfnet2-nx2-dns3 --final_model_name dfnet2-nx2-dns3 \
+--noise_dir "/data/tianxing/HuggingDatasets/nx_noise/data/noise/" \
+--speech_dir "/data/tianxing/HuggingDatasets/nx_noise/data/speech/"
+
+
+END
+
+
+# params
+system_version="windows";
+verbose=true;
+stage=0 # start from 0 if you need to start from data preparation
+stop_stage=9
+
+work_dir="$(pwd)"
+file_folder_name=file_folder_name
+final_model_name=final_model_name
+config_file="yaml/config.yaml"
+limit=10
+
+noise_dir=/data/tianxing/HuggingDatasets/nx_noise/data/noise
+speech_dir=/data/tianxing/HuggingDatasets/aishell/data_aishell/wav/train
+
+max_count=-1
+
+nohup_name=nohup.out
+
+# model params
+batch_size=64
+max_epochs=200
+save_top_k=10
+patience=5
+
+
+# parse options
+while true; do
+  [ -z "${1:-}" ] && break;  # break if there are no arguments
+  case "$1" in
+    --*) name=$(echo "$1" | sed s/^--// | sed s/-/_/g);
+      eval '[ -z "${'"$name"'+xxx}" ]' && echo "$0: invalid option $1" 1>&2 && exit 1;
+      old_value="(eval echo \\$$name)";
+      if [ "${old_value}" == "true" ] || [ "${old_value}" == "false" ]; then
+        was_bool=true;
+      else
+        was_bool=false;
+      fi
+
+      # Set the variable to the right value-- the escaped quotes make it work if
+      # the option had spaces, like --cmd "queue.pl -sync y"
+      eval "${name}=\"$2\"";
+
+      # Check that Boolean-valued arguments are really Boolean.
+      if $was_bool && [[ "$2" != "true" && "$2" != "false" ]]; then
+        echo "$0: expected \"true\" or \"false\": $1 $2" 1>&2
+        exit 1;
+      fi
+      shift 2;
+      ;;
+
+    *) break;
+  esac
+done
+
+file_dir="${work_dir}/${file_folder_name}"
+final_model_dir="${work_dir}/../../trained_models/${final_model_name}";
+evaluation_audio_dir="${file_dir}/evaluation_audio"
+
+train_dataset="${file_dir}/train.jsonl"
+valid_dataset="${file_dir}/valid.jsonl"
+
+$verbose && echo "system_version: ${system_version}"
+$verbose && echo "file_folder_name: ${file_folder_name}"
+
+if [ $system_version == "windows" ]; then
+  alias python3='D:/Users/tianx/PycharmProjects/virtualenv/nx_denoise/Scripts/python.exe'
+elif [ $system_version == "centos" ] || [ $system_version == "ubuntu" ]; then
+  #source /data/local/bin/nx_denoise/bin/activate
+  alias python3='/data/local/bin/nx_denoise/bin/python3'
+fi
+
+
+if [ ${stage} -le 1 ] && [ ${stop_stage} -ge 1 ]; then
+  $verbose && echo "stage 1: prepare data"
+  cd "${work_dir}" || exit 1
+  python3 step_1_prepare_data.py \
+  --file_dir "${file_dir}" \
+  --noise_dir "${noise_dir}" \
+  --speech_dir "${speech_dir}" \
+  --train_dataset "${train_dataset}" \
+  --valid_dataset "${valid_dataset}" \
+  --max_count "${max_count}" \
+
+fi
+
+
+if [ ${stage} -le 2 ] && [ ${stop_stage} -ge 2 ]; then
+  $verbose && echo "stage 2: train model"
+  cd "${work_dir}" || exit 1
+  python3 step_2_train_model.py \
+  --train_dataset "${train_dataset}" \
+  --valid_dataset "${valid_dataset}" \
+  --serialization_dir "${file_dir}" \
+  --config_file "${config_file}" \
+
+fi
+
+
+if [ ${stage} -le 3 ] && [ ${stop_stage} -ge 3 ]; then
+  $verbose && echo "stage 3: test model"
+  cd "${work_dir}" || exit 1
+  python3 step_3_evaluation.py \
+  --valid_dataset "${valid_dataset}" \
+  --model_dir "${file_dir}/best" \
+  --evaluation_audio_dir "${evaluation_audio_dir}" \
+  --limit "${limit}" \
+
+fi
+
+
+if [ ${stage} -le 4 ] && [ ${stop_stage} -ge 4 ]; then
+  $verbose && echo "stage 4: collect files"
+  cd "${work_dir}" || exit 1
+
+  mkdir -p ${final_model_dir}
+
+  cp "${file_dir}/best"/* "${final_model_dir}"
+  cp -r "${file_dir}/evaluation_audio" "${final_model_dir}"
+
+  cd "${final_model_dir}/.." || exit 1;
+
+  if [ -e "${final_model_name}.zip" ]; then
+    rm -rf "${final_model_name}_backup.zip"
+    mv "${final_model_name}.zip" "${final_model_name}_backup.zip"
+  fi
+
+  zip -r "${final_model_name}.zip" "${final_model_name}"
+  rm -rf "${final_model_name}"
+
+fi
+
+
+if [ ${stage} -le 5 ] && [ ${stop_stage} -ge 5 ]; then
+  $verbose && echo "stage 5: clear file_dir"
+  cd "${work_dir}" || exit 1
+
+  rm -rf "${file_dir}";
+
+fi

examples/dfnet2/step_1_prepare_data.py

@@ -0,0 +1,164 @@
+#!/usr/bin/python3
+# -*- coding: utf-8 -*-
+import argparse
+import json
+import os
+from pathlib import Path
+import random
+import sys
+
+pwd = os.path.abspath(os.path.dirname(__file__))
+sys.path.append(os.path.join(pwd, "../../"))
+
+import librosa
+import numpy as np
+from tqdm import tqdm
+
+
+def get_args():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--file_dir", default="./", type=str)
+
+    parser.add_argument(
+        "--noise_dir",
+        default=r"E:\Users\tianx\HuggingDatasets\nx_noise\data\noise",
+        type=str
+    )
+    parser.add_argument(
+        "--speech_dir",
+        default=r"E:\programmer\asr_datasets\aishell\data_aishell\wav\train",
+        type=str
+    )
+
+    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=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=-1, type=int)
+
+    args = parser.parse_args()
+    return args
+
+
+def filename_generator(data_dir: str):
+    data_dir = Path(data_dir)
+    for filename in data_dir.glob("**/*.wav"):
+        yield filename.as_posix()
+
+
+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 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):
+                if np.sum(signal[begin: begin+win_size]) == 0:
+                    continue
+                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 main():
+    args = get_args()
+
+    file_dir = Path(args.file_dir)
+    file_dir.mkdir(exist_ok=True)
+
+    noise_dir = Path(args.noise_dir)
+    speech_dir = Path(args.speech_dir)
+
+    noise_generator = target_second_signal_generator(
+        noise_dir.as_posix(),
+        duration=args.duration,
+        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,
+        max_epoch=1,
+    )
+
+    dataset = list()
+
+    count = 0
+    process_bar = tqdm(desc="build dataset jsonl")
+    with open(args.train_dataset, "w", encoding="utf-8") as ftrain, open(args.valid_dataset, "w", encoding="utf-8") as fvalid:
+        for noise, speech in zip(noise_generator, speech_generator):
+            if count >= args.max_count > 0:
+                break
+
+            noise_filename = noise["filename"]
+            noise_raw_duration = noise["raw_duration"]
+            noise_offset = noise["offset"]
+            noise_duration = noise["duration"]
+
+            speech_filename = speech["filename"]
+            speech_raw_duration = speech["raw_duration"]
+            speech_offset = speech["offset"]
+            speech_duration = speech["duration"]
+
+            random1 = random.random()
+            random2 = random.random()
+
+            row = {
+                "count": count,
+
+                "noise_filename": noise_filename,
+                "noise_raw_duration": noise_raw_duration,
+                "noise_offset": noise_offset,
+                "noise_duration": noise_duration,
+
+                "speech_filename": speech_filename,
+                "speech_raw_duration": speech_raw_duration,
+                "speech_offset": speech_offset,
+                "speech_duration": speech_duration,
+
+                "snr_db": random.uniform(args.min_snr_db, args.max_snr_db),
+
+                "random1": random1,
+            }
+            row = json.dumps(row, ensure_ascii=False)
+            if random2 < (1 / 300 / 1):
+                fvalid.write(f"{row}\n")
+            else:
+                ftrain.write(f"{row}\n")
+
+            count += 1
+            duration_seconds = count * args.duration
+            duration_hours = duration_seconds / 3600
+
+            process_bar.update(n=1)
+            process_bar.set_postfix({
+                # "duration_seconds": round(duration_seconds, 4),
+                "duration_hours": round(duration_hours, 4),
+
+            })
+
+    return
+
+
+if __name__ == "__main__":
+    main()

examples/dfnet2/step_2_train_model.py

@@ -0,0 +1,469 @@
+#!/usr/bin/python3
+# -*- coding: utf-8 -*-
+"""
+https://github.com/Rikorose/DeepFilterNet
+"""
+import argparse
+import json
+import logging
+from logging.handlers import TimedRotatingFileHandler
+import os
+import platform
+from pathlib import Path
+import random
+import sys
+import shutil
+from typing import List
+
+from fontTools.varLib.plot import stops
+
+pwd = os.path.abspath(os.path.dirname(__file__))
+sys.path.append(os.path.join(pwd, "../../"))
+
+import numpy as np
+import torch
+import torch.nn as nn
+from torch.nn import functional as F
+from torch.utils.data.dataloader import DataLoader
+from tqdm import tqdm
+
+from toolbox.torch.utils.data.dataset.denoise_jsonl_dataset import DenoiseJsonlDataset
+from toolbox.torchaudio.losses.snr import NegativeSISNRLoss
+from toolbox.torchaudio.losses.spectral import LSDLoss, MultiResolutionSTFTLoss
+from toolbox.torchaudio.metrics.pesq import run_pesq_score
+from toolbox.torchaudio.models.dfnet2.configuration_dfnet2 import DfNet2Config
+from toolbox.torchaudio.models.dfnet2.modeling_dfnet2 import DfNet2, DfNet2PretrainedModel
+
+
+def get_args():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--train_dataset", default="train.jsonl", type=str)
+    parser.add_argument("--valid_dataset", default="valid.jsonl", type=str)
+
+    parser.add_argument("--num_serialized_models_to_keep", default=15, type=int)
+    parser.add_argument("--patience", default=30, type=int)
+    parser.add_argument("--serialization_dir", default="serialization_dir", type=str)
+
+    parser.add_argument("--config_file", default="config.yaml", type=str)
+
+    args = parser.parse_args()
+    return args
+
+
+def logging_config(file_dir: str):
+    fmt = "%(asctime)s - %(name)s - %(levelname)s  %(filename)s:%(lineno)d >  %(message)s"
+
+    logging.basicConfig(format=fmt,
+                        datefmt="%m/%d/%Y %H:%M:%S",
+                        level=logging.INFO)
+    file_handler = TimedRotatingFileHandler(
+        filename=os.path.join(file_dir, "main.log"),
+        encoding="utf-8",
+        when="D",
+        interval=1,
+        backupCount=7
+    )
+    file_handler.setLevel(logging.INFO)
+    file_handler.setFormatter(logging.Formatter(fmt))
+    logger = logging.getLogger(__name__)
+    logger.addHandler(file_handler)
+
+    return logger
+
+
+class CollateFunction(object):
+    def __init__(self):
+        pass
+
+    def __call__(self, batch: List[dict]):
+        clean_audios = list()
+        noisy_audios = list()
+        snr_db_list = list()
+
+        for sample in batch:
+            # 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"]
+
+            clean_audios.append(clean_audio)
+            noisy_audios.append(noisy_audio)
+
+        clean_audios = torch.stack(clean_audios)
+        noisy_audios = torch.stack(noisy_audios)
+
+        # assert
+        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()
+
+
+def main():
+    args = get_args()
+
+    config = DfNet2Config.from_pretrained(
+        pretrained_model_name_or_path=args.config_file,
+    )
+
+    serialization_dir = Path(args.serialization_dir)
+    serialization_dir.mkdir(parents=True, exist_ok=True)
+
+    logger = logging_config(serialization_dir)
+
+    random.seed(config.seed)
+    np.random.seed(config.seed)
+    torch.manual_seed(config.seed)
+    logger.info(f"set seed: {config.seed}")
+
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    n_gpu = torch.cuda.device_count()
+    logger.info(f"GPU available count: {n_gpu}; device: {device}")
+
+    # datasets
+    train_dataset = DenoiseJsonlDataset(
+        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=225000,
+    )
+    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,
+    )
+    train_data_loader = DataLoader(
+        dataset=train_dataset,
+        batch_size=config.batch_size,
+        # shuffle=True,
+        sampler=None,
+        # Linux 系统中可以使用多个子进程加载数据, 而在 Windows 系统中不能.
+        num_workers=0 if platform.system() == "Windows" else os.cpu_count() // 2,
+        collate_fn=collate_fn,
+        pin_memory=False,
+        prefetch_factor=None if platform.system() == "Windows" else 2,
+    )
+    valid_data_loader = DataLoader(
+        dataset=valid_dataset,
+        batch_size=config.batch_size,
+        # shuffle=True,
+        sampler=None,
+        # Linux 系统中可以使用多个子进程加载数据, 而在 Windows 系统中不能.
+        num_workers=0 if platform.system() == "Windows" else os.cpu_count() // 2,
+        collate_fn=collate_fn,
+        pin_memory=False,
+        prefetch_factor=None if platform.system() == "Windows" else 2,
+    )
+
+    # models
+    logger.info(f"prepare models. config_file: {args.config_file}")
+    model = DfNet2PretrainedModel(config).to(device)
+    model.to(device)
+    model.train()
+
+    # optimizer
+    logger.info("prepare optimizer, lr_scheduler, loss_fn, evaluation_metric")
+    optimizer = torch.optim.AdamW(model.parameters(), config.lr)
+
+    # resume training
+    last_step_idx = -1
+    last_epoch = -1
+    for step_idx_str in serialization_dir.glob("steps-*"):
+        step_idx_str = Path(step_idx_str)
+        step_idx = step_idx_str.stem.split("-")[1]
+        step_idx = int(step_idx)
+        if step_idx > last_step_idx:
+            last_step_idx = step_idx
+    # last_epoch = 1
+
+    if last_step_idx != -1:
+        logger.info(f"resume from steps-{last_step_idx}.")
+        model_pt = serialization_dir / f"steps-{last_step_idx}/model.pt"
+
+        logger.info(f"load state dict for model.")
+        with open(model_pt.as_posix(), "rb") as f:
+            state_dict = torch.load(f, map_location="cpu", weights_only=True)
+        model.load_state_dict(state_dict, strict=True)
+
+    if config.lr_scheduler == "CosineAnnealingLR":
+        lr_scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(
+            optimizer,
+            last_epoch=last_epoch,
+            # T_max=10 * config.eval_steps,
+            # eta_min=0.01 * config.lr,
+            **config.lr_scheduler_kwargs,
+        )
+    elif config.lr_scheduler == "MultiStepLR":
+        lr_scheduler = torch.optim.lr_scheduler.MultiStepLR(
+            optimizer,
+            last_epoch=last_epoch,
+            milestones=[10000, 20000, 30000, 40000, 50000], gamma=0.5
+        )
+    else:
+        raise AssertionError(f"invalid lr_scheduler: {config.lr_scheduler}")
+
+    neg_si_snr_loss_fn = NegativeSISNRLoss(reduction="mean").to(device)
+    mr_stft_loss_fn = MultiResolutionSTFTLoss(
+        fft_size_list=[256, 512, 1024],
+        win_size_list=[256, 512, 1024],
+        hop_size_list=[128, 256, 512],
+        factor_sc=1.5,
+        factor_mag=1.0,
+        reduction="mean"
+    ).to(device)
+
+    # training loop
+
+    # state
+    average_pesq_score = 1000000000
+    average_loss = 1000000000
+    average_mr_stft_loss = 1000000000
+    average_neg_si_snr_loss = 1000000000
+    average_mask_loss = 1000000000
+    average_lsnr_loss = 1000000000
+
+    model_list = list()
+    best_epoch_idx = None
+    best_step_idx = None
+    best_metric = None
+    patience_count = 0
+
+    step_idx = 0 if last_step_idx == -1 else last_step_idx
+
+    logger.info("training")
+    early_stop_flag = False
+    for epoch_idx in range(max(0, last_epoch+1), config.max_epochs):
+        if early_stop_flag:
+            break
+
+        # train
+        model.train()
+
+        total_pesq_score = 0.
+        total_loss = 0.
+        total_mr_stft_loss = 0.
+        total_neg_si_snr_loss = 0.
+        total_mask_loss = 0.
+        total_lsnr_loss = 0.
+        total_batches = 0.
+
+        progress_bar_train = tqdm(
+            initial=step_idx,
+            desc="Training; epoch-{}".format(epoch_idx),
+        )
+        for train_batch in train_data_loader:
+            clean_audios, noisy_audios = train_batch
+            clean_audios: torch.Tensor = clean_audios.to(device)
+            noisy_audios: torch.Tensor = noisy_audios.to(device)
+
+            est_spec, est_wav, est_mask, lsnr = model.forward(noisy_audios)
+            # est_wav shape: [b, 1, n_samples]
+            est_wav = torch.squeeze(est_wav, dim=1)
+            # est_wav shape: [b, n_samples]
+
+            mr_stft_loss = mr_stft_loss_fn.forward(est_wav, clean_audios)
+            neg_si_snr_loss = neg_si_snr_loss_fn.forward(est_wav, clean_audios)
+            mask_loss = model.mask_loss_fn(est_mask, clean_audios, noisy_audios)
+            lsnr_loss = model.lsnr_loss_fn(lsnr, clean_audios, noisy_audios)
+
+            loss = 1.0 * mr_stft_loss + 1.0 * neg_si_snr_loss + 1.0 * mask_loss + 0.01 * lsnr_loss
+            if torch.any(torch.isnan(loss)) or torch.any(torch.isinf(loss)):
+                logger.info(f"find nan or inf in loss. continue.")
+                continue
+
+            denoise_audios_list_r = list(est_wav.detach().cpu().numpy())
+            clean_audios_list_r = list(clean_audios.detach().cpu().numpy())
+            pesq_score = run_pesq_score(clean_audios_list_r, denoise_audios_list_r, sample_rate=config.sample_rate, mode="nb")
+
+            optimizer.zero_grad()
+            loss.backward()
+            torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=config.clip_grad_norm)
+            optimizer.step()
+            lr_scheduler.step()
+
+            total_pesq_score += pesq_score
+            total_loss += loss.item()
+            total_mr_stft_loss += mr_stft_loss.item()
+            total_neg_si_snr_loss += neg_si_snr_loss.item()
+            total_mask_loss += mask_loss.item()
+            total_lsnr_loss += lsnr_loss.item()
+            total_batches += 1
+
+            average_pesq_score = round(total_pesq_score / total_batches, 4)
+            average_loss = round(total_loss / total_batches, 4)
+            average_mr_stft_loss = round(total_mr_stft_loss / total_batches, 4)
+            average_neg_si_snr_loss = round(total_neg_si_snr_loss / total_batches, 4)
+            average_mask_loss = round(total_mask_loss / total_batches, 4)
+            average_lsnr_loss = round(total_lsnr_loss / total_batches, 4)
+
+            progress_bar_train.update(1)
+            progress_bar_train.set_postfix({
+                "lr": lr_scheduler.get_last_lr()[0],
+                "pesq_score": average_pesq_score,
+                "loss": average_loss,
+                "mr_stft_loss": average_mr_stft_loss,
+                "neg_si_snr_loss": average_neg_si_snr_loss,
+                "mask_loss": average_mask_loss,
+                "lsnr_loss": average_lsnr_loss,
+            })
+
+            # evaluation
+            step_idx += 1
+            if step_idx % config.eval_steps == 0:
+                with torch.no_grad():
+                    torch.cuda.empty_cache()
+
+                    model.eval()
+
+                    total_pesq_score = 0.
+                    total_loss = 0.
+                    total_mr_stft_loss = 0.
+                    total_neg_si_snr_loss = 0.
+                    total_mask_loss = 0.
+                    total_lsnr_loss = 0.
+                    total_batches = 0.
+
+                    progress_bar_train.close()
+                    progress_bar_eval = tqdm(
+                        desc="Evaluation; steps-{}k".format(int(step_idx/1000)),
+                    )
+                    for eval_batch in valid_data_loader:
+                        clean_audios, noisy_audios = eval_batch
+                        clean_audios: torch.Tensor = clean_audios.to(device)
+                        noisy_audios: torch.Tensor = noisy_audios.to(device)
+
+                        est_spec, est_wav, est_mask, lsnr = model.forward(noisy_audios)
+                        # est_wav shape: [b, 1, n_samples]
+                        est_wav = torch.squeeze(est_wav, dim=1)
+                        # est_wav shape: [b, n_samples]
+
+                        mr_stft_loss = mr_stft_loss_fn.forward(est_wav, clean_audios)
+                        neg_si_snr_loss = neg_si_snr_loss_fn.forward(est_wav, clean_audios)
+                        mask_loss = model.mask_loss_fn(est_mask, clean_audios, noisy_audios)
+                        lsnr_loss = model.lsnr_loss_fn(lsnr, clean_audios, noisy_audios)
+
+                        loss = 1.0 * mr_stft_loss + 1.0 * neg_si_snr_loss + 1.0 * mask_loss + 0.01 * lsnr_loss
+                        if torch.any(torch.isnan(loss)) or torch.any(torch.isinf(loss)):
+                            logger.info(f"find nan or inf in loss. continue.")
+                            continue
+
+                        denoise_audios_list_r = list(est_wav.detach().cpu().numpy())
+                        clean_audios_list_r = list(clean_audios.detach().cpu().numpy())
+                        pesq_score = run_pesq_score(clean_audios_list_r, denoise_audios_list_r, sample_rate=config.sample_rate, mode="nb")
+
+                        total_pesq_score += pesq_score
+                        total_loss += loss.item()
+                        total_mr_stft_loss += mr_stft_loss.item()
+                        total_neg_si_snr_loss += neg_si_snr_loss.item()
+                        total_mask_loss += mask_loss.item()
+                        total_lsnr_loss += lsnr_loss.item()
+                        total_batches += 1
+
+                        average_pesq_score = round(total_pesq_score / total_batches, 4)
+                        average_loss = round(total_loss / total_batches, 4)
+                        average_mr_stft_loss = round(total_mr_stft_loss / total_batches, 4)
+                        average_neg_si_snr_loss = round(total_neg_si_snr_loss / total_batches, 4)
+                        average_mask_loss = round(total_mask_loss / total_batches, 4)
+                        average_lsnr_loss = round(total_lsnr_loss / total_batches, 4)
+
+                        progress_bar_eval.update(1)
+                        progress_bar_eval.set_postfix({
+                            "lr": lr_scheduler.get_last_lr()[0],
+                            "pesq_score": average_pesq_score,
+                            "loss": average_loss,
+                            "mr_stft_loss": average_mr_stft_loss,
+                            "neg_si_snr_loss": average_neg_si_snr_loss,
+                            "mask_loss": average_mask_loss,
+                            "lsnr_loss": average_lsnr_loss,
+                        })
+
+                model.train()
+
+                total_pesq_score = 0.
+                total_loss = 0.
+                total_mr_stft_loss = 0.
+                total_neg_si_snr_loss = 0.
+                total_mask_loss = 0.
+                total_lsnr_loss = 0.
+                total_batches = 0.
+
+                progress_bar_eval.close()
+                progress_bar_train = tqdm(
+                    initial=progress_bar_train.n,
+                    postfix=progress_bar_train.postfix,
+                    desc=progress_bar_train.desc,
+                )
+
+                # save path
+                save_dir = serialization_dir / "steps-{}".format(step_idx)
+                save_dir.mkdir(parents=True, exist_ok=False)
+
+                # save models
+                model.save_pretrained(save_dir.as_posix())
+
+                model_list.append(save_dir)
+                if len(model_list) >= args.num_serialized_models_to_keep:
+                    model_to_delete: Path = model_list.pop(0)
+                    shutil.rmtree(model_to_delete.as_posix())
+
+                # save metric
+                if best_metric is None:
+                    best_epoch_idx = epoch_idx
+                    best_step_idx = step_idx
+                    best_metric = average_pesq_score
+                elif average_pesq_score >= best_metric:
+                    # great is better.
+                    best_epoch_idx = epoch_idx
+                    best_step_idx = step_idx
+                    best_metric = average_pesq_score
+                else:
+                    pass
+
+                metrics = {
+                    "epoch_idx": epoch_idx,
+                    "best_epoch_idx": best_epoch_idx,
+                    "best_step_idx": best_step_idx,
+                    "pesq_score": average_pesq_score,
+                    "loss": average_loss,
+                    "mr_stft_loss": average_mr_stft_loss,
+                    "neg_si_snr_loss": average_neg_si_snr_loss,
+                    "mask_loss": average_mask_loss,
+                    "lsnr_loss": average_lsnr_loss,
+                }
+                metrics_filename = save_dir / "metrics_epoch.json"
+                with open(metrics_filename, "w", encoding="utf-8") as f:
+                    json.dump(metrics, f, indent=4, ensure_ascii=False)
+
+                # save best
+                best_dir = serialization_dir / "best"
+                if best_epoch_idx == epoch_idx and best_step_idx == step_idx:
+                    if best_dir.exists():
+                        shutil.rmtree(best_dir)
+                    shutil.copytree(save_dir, best_dir)
+
+                # early stop
+                early_stop_flag = False
+                if best_epoch_idx == epoch_idx and best_step_idx == step_idx:
+                    patience_count = 0
+                else:
+                    patience_count += 1
+                if patience_count >= args.patience:
+                    early_stop_flag = True
+
+                # early stop
+                if early_stop_flag:
+                    break
+
+    return
+
+
+if __name__ == "__main__":
+    main()

examples/dfnet2/yaml/config.yaml

@@ -0,0 +1,75 @@
+model_name: "dfnet2"
+
+# spec
+sample_rate: 8000
+nfft: 512
+win_size: 200
+hop_size: 80
+
+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: -5
+max_snr_db: 40
+
+# 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

examples/dtln/run.sh

@@ -0,0 +1,171 @@
+#!/usr/bin/env bash
+
+: <<'END'
+
+sh run.sh --stage 2 --stop_stage 2 --system_version centos --file_folder_name file_dir-256 --final_model_name dtln-256-nx-dns3 \
+--config_file "yaml/config-256.yaml" \
+--noise_dir "/data/tianxing/HuggingDatasets/nx_noise/data/noise" \
+--speech_dir "/data/tianxing/HuggingDatasets/nx_noise/data/speech/dns3-speech"
+
+
+sh run.sh --stage 1 --stop_stage 2 --system_version centos --file_folder_name file_dir-512 --final_model_name dtln-512-nx-dns3 \
+--config_file "yaml/config-512.yaml" \
+--noise_dir "/data/tianxing/HuggingDatasets/nx_noise/data/noise" \
+--speech_dir "/data/tianxing/HuggingDatasets/nx_noise/data/speech/dns3-speech"
+
+
+sh run.sh --stage 2 --stop_stage 2 --system_version centos --file_folder_name dtnl-1024-nx2 --final_model_name dtln-1024-nx2 \
+--config_file "yaml/config-1024.yaml" \
+--noise_dir "/data/tianxing/HuggingDatasets/nx_noise/data/noise/nx-noise" \
+--speech_dir "/data/tianxing/HuggingDatasets/nx_noise/data/speech/nx-speech2"
+
+
+sh run.sh --stage 2 --stop_stage 2 --system_version centos --file_folder_name dtln-256-nx2-dns3 --final_model_name dtln-256-nx2-dns3 \
+--config_file "yaml/config-256.yaml" \
+--noise_dir "/data/tianxing/HuggingDatasets/nx_noise/data/noise" \
+--speech_dir "/data/tianxing/HuggingDatasets/nx_noise/data/speech"
+
+
+END
+
+
+# params
+system_version="windows";
+verbose=true;
+stage=0 # start from 0 if you need to start from data preparation
+stop_stage=9
+
+work_dir="$(pwd)"
+file_folder_name=file_folder_name
+final_model_name=final_model_name
+config_file="yaml/config.yaml"
+limit=10
+
+noise_dir=/data/tianxing/HuggingDatasets/nx_noise/data/noise
+speech_dir=/data/tianxing/HuggingDatasets/aishell/data_aishell/wav/train
+
+max_count=-1
+
+nohup_name=nohup.out
+
+# model params
+batch_size=64
+max_epochs=200
+save_top_k=10
+patience=5
+
+
+# parse options
+while true; do
+  [ -z "${1:-}" ] && break;  # break if there are no arguments
+  case "$1" in
+    --*) name=$(echo "$1" | sed s/^--// | sed s/-/_/g);
+      eval '[ -z "${'"$name"'+xxx}" ]' && echo "$0: invalid option $1" 1>&2 && exit 1;
+      old_value="(eval echo \\$$name)";
+      if [ "${old_value}" == "true" ] || [ "${old_value}" == "false" ]; then
+        was_bool=true;
+      else
+        was_bool=false;
+      fi
+
+      # Set the variable to the right value-- the escaped quotes make it work if
+      # the option had spaces, like --cmd "queue.pl -sync y"
+      eval "${name}=\"$2\"";
+
+      # Check that Boolean-valued arguments are really Boolean.
+      if $was_bool && [[ "$2" != "true" && "$2" != "false" ]]; then
+        echo "$0: expected \"true\" or \"false\": $1 $2" 1>&2
+        exit 1;
+      fi
+      shift 2;
+      ;;
+
+    *) break;
+  esac
+done
+
+file_dir="${work_dir}/${file_folder_name}"
+final_model_dir="${work_dir}/../../trained_models/${final_model_name}";
+evaluation_audio_dir="${file_dir}/evaluation_audio"
+
+train_dataset="${file_dir}/train.jsonl"
+valid_dataset="${file_dir}/valid.jsonl"
+
+$verbose && echo "system_version: ${system_version}"
+$verbose && echo "file_folder_name: ${file_folder_name}"
+
+if [ $system_version == "windows" ]; then
+  alias python3='D:/Users/tianx/PycharmProjects/virtualenv/nx_denoise/Scripts/python.exe'
+elif [ $system_version == "centos" ] || [ $system_version == "ubuntu" ]; then
+  #source /data/local/bin/nx_denoise/bin/activate
+  alias python3='/data/local/bin/nx_denoise/bin/python3'
+fi
+
+
+if [ ${stage} -le 1 ] && [ ${stop_stage} -ge 1 ]; then
+  $verbose && echo "stage 1: prepare data"
+  cd "${work_dir}" || exit 1
+  python3 step_1_prepare_data.py \
+  --file_dir "${file_dir}" \
+  --noise_dir "${noise_dir}" \
+  --speech_dir "${speech_dir}" \
+  --train_dataset "${train_dataset}" \
+  --valid_dataset "${valid_dataset}" \
+  --max_count "${max_count}" \
+
+fi
+
+
+if [ ${stage} -le 2 ] && [ ${stop_stage} -ge 2 ]; then
+  $verbose && echo "stage 2: train model"
+  cd "${work_dir}" || exit 1
+  python3 step_2_train_model.py \
+  --train_dataset "${train_dataset}" \
+  --valid_dataset "${valid_dataset}" \
+  --serialization_dir "${file_dir}" \
+  --config_file "${config_file}" \
+
+fi
+
+
+if [ ${stage} -le 3 ] && [ ${stop_stage} -ge 3 ]; then
+  $verbose && echo "stage 3: test model"
+  cd "${work_dir}" || exit 1
+  python3 step_3_evaluation.py \
+  --valid_dataset "${valid_dataset}" \
+  --model_dir "${file_dir}/best" \
+  --evaluation_audio_dir "${evaluation_audio_dir}" \
+  --limit "${limit}" \
+
+fi
+
+
+if [ ${stage} -le 4 ] && [ ${stop_stage} -ge 4 ]; then
+  $verbose && echo "stage 4: collect files"
+  cd "${work_dir}" || exit 1
+
+  mkdir -p ${final_model_dir}
+
+  cp "${file_dir}/best"/* "${final_model_dir}"
+  cp -r "${file_dir}/evaluation_audio" "${final_model_dir}"
+
+  cd "${final_model_dir}/.." || exit 1;
+
+  if [ -e "${final_model_name}.zip" ]; then
+    rm -rf "${final_model_name}_backup.zip"
+    mv "${final_model_name}.zip" "${final_model_name}_backup.zip"
+  fi
+
+  zip -r "${final_model_name}.zip" "${final_model_name}"
+  rm -rf "${final_model_name}"
+
+fi
+
+
+if [ ${stage} -le 5 ] && [ ${stop_stage} -ge 5 ]; then
+  $verbose && echo "stage 5: clear file_dir"
+  cd "${work_dir}" || exit 1
+
+  rm -rf "${file_dir}";
+
+fi

examples/dtln/step_1_prepare_data.py

@@ -0,0 +1,164 @@
+#!/usr/bin/python3
+# -*- coding: utf-8 -*-
+import argparse
+import json
+import os
+from pathlib import Path
+import random
+import sys
+
+pwd = os.path.abspath(os.path.dirname(__file__))
+sys.path.append(os.path.join(pwd, "../../"))
+
+import librosa
+import numpy as np
+from tqdm import tqdm
+
+
+def get_args():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--file_dir", default="./", type=str)
+
+    parser.add_argument(
+        "--noise_dir",
+        default=r"E:\Users\tianx\HuggingDatasets\nx_noise\data\noise",
+        type=str
+    )
+    parser.add_argument(
+        "--speech_dir",
+        default=r"E:\programmer\asr_datasets\aishell\data_aishell\wav\train",
+        type=str
+    )
+
+    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=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=-1, type=int)
+
+    args = parser.parse_args()
+    return args
+
+
+def filename_generator(data_dir: str):
+    data_dir = Path(data_dir)
+    for filename in data_dir.glob("**/*.wav"):
+        yield filename.as_posix()
+
+
+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 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):
+                if np.sum(signal[begin: begin+win_size]) == 0:
+                    continue
+                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 main():
+    args = get_args()
+
+    file_dir = Path(args.file_dir)
+    file_dir.mkdir(exist_ok=True)
+
+    noise_dir = Path(args.noise_dir)
+    speech_dir = Path(args.speech_dir)
+
+    noise_generator = target_second_signal_generator(
+        noise_dir.as_posix(),
+        duration=args.duration,
+        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,
+        max_epoch=1,
+    )
+
+    dataset = list()
+
+    count = 0
+    process_bar = tqdm(desc="build dataset jsonl")
+    with open(args.train_dataset, "w", encoding="utf-8") as ftrain, open(args.valid_dataset, "w", encoding="utf-8") as fvalid:
+        for noise, speech in zip(noise_generator, speech_generator):
+            if count >= args.max_count > 0:
+                break
+
+            noise_filename = noise["filename"]
+            noise_raw_duration = noise["raw_duration"]
+            noise_offset = noise["offset"]
+            noise_duration = noise["duration"]
+
+            speech_filename = speech["filename"]
+            speech_raw_duration = speech["raw_duration"]
+            speech_offset = speech["offset"]
+            speech_duration = speech["duration"]
+
+            random1 = random.random()
+            random2 = random.random()
+
+            row = {
+                "count": count,
+
+                "noise_filename": noise_filename,
+                "noise_raw_duration": noise_raw_duration,
+                "noise_offset": noise_offset,
+                "noise_duration": noise_duration,
+
+                "speech_filename": speech_filename,
+                "speech_raw_duration": speech_raw_duration,
+                "speech_offset": speech_offset,
+                "speech_duration": speech_duration,
+
+                "snr_db": random.uniform(args.min_snr_db, args.max_snr_db),
+
+                "random1": random1,
+            }
+            row = json.dumps(row, ensure_ascii=False)
+            if random2 < (1 / 300 / 1):
+                fvalid.write(f"{row}\n")
+            else:
+                ftrain.write(f"{row}\n")
+
+            count += 1
+            duration_seconds = count * args.duration
+            duration_hours = duration_seconds / 3600
+
+            process_bar.update(n=1)
+            process_bar.set_postfix({
+                # "duration_seconds": round(duration_seconds, 4),
+                "duration_hours": round(duration_hours, 4),
+
+            })
+
+    return
+
+
+if __name__ == "__main__":
+    main()

examples/dtln/step_2_train_model.py

@@ -0,0 +1,437 @@
+#!/usr/bin/python3
+# -*- coding: utf-8 -*-
+"""
+https://github.com/breizhn/DTLN
+
+"""
+import argparse
+import json
+import logging
+from logging.handlers import TimedRotatingFileHandler
+import os
+import platform
+from pathlib import Path
+import random
+import sys
+import shutil
+from typing import List
+
+pwd = os.path.abspath(os.path.dirname(__file__))
+sys.path.append(os.path.join(pwd, "../../"))
+
+import numpy as np
+import torch
+import torch.nn as nn
+from torch.nn import functional as F
+from torch.utils.data.dataloader import DataLoader
+from tqdm import tqdm
+
+from toolbox.torch.utils.data.dataset.denoise_jsonl_dataset import DenoiseJsonlDataset
+from toolbox.torchaudio.losses.snr import NegativeSISNRLoss
+from toolbox.torchaudio.losses.spectral import LSDLoss, MultiResolutionSTFTLoss
+from toolbox.torchaudio.metrics.pesq import run_pesq_score
+from toolbox.torchaudio.models.dtln.configuration_dtln import DTLNConfig
+from toolbox.torchaudio.models.dtln.modeling_dtln import DTLNModel, DTLNPretrainedModel
+
+
+def get_args():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--train_dataset", default="train.jsonl", type=str)
+    parser.add_argument("--valid_dataset", default="valid.jsonl", type=str)
+
+    parser.add_argument("--num_serialized_models_to_keep", default=15, type=int)
+    parser.add_argument("--patience", default=30, type=int)
+    parser.add_argument("--serialization_dir", default="serialization_dir", type=str)
+
+    parser.add_argument("--config_file", default="config.yaml", type=str)
+
+    args = parser.parse_args()
+    return args
+
+
+def logging_config(file_dir: str):
+    fmt = "%(asctime)s - %(name)s - %(levelname)s  %(filename)s:%(lineno)d >  %(message)s"
+
+    logging.basicConfig(format=fmt,
+                        datefmt="%m/%d/%Y %H:%M:%S",
+                        level=logging.INFO)
+    file_handler = TimedRotatingFileHandler(
+        filename=os.path.join(file_dir, "main.log"),
+        encoding="utf-8",
+        when="D",
+        interval=1,
+        backupCount=7
+    )
+    file_handler.setLevel(logging.INFO)
+    file_handler.setFormatter(logging.Formatter(fmt))
+    logger = logging.getLogger(__name__)
+    logger.addHandler(file_handler)
+
+    return logger
+
+
+class CollateFunction(object):
+    def __init__(self):
+        pass
+
+    def __call__(self, batch: List[dict]):
+        clean_audios = list()
+        noisy_audios = list()
+        snr_db_list = list()
+
+        for sample in batch:
+            # 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"]
+
+            clean_audios.append(clean_audio)
+            noisy_audios.append(noisy_audio)
+
+        clean_audios = torch.stack(clean_audios)
+        noisy_audios = torch.stack(noisy_audios)
+
+        # assert
+        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()
+
+
+def main():
+    args = get_args()
+
+    config = DTLNConfig.from_pretrained(
+        pretrained_model_name_or_path=args.config_file,
+    )
+
+    serialization_dir = Path(args.serialization_dir)
+    serialization_dir.mkdir(parents=True, exist_ok=True)
+
+    logger = logging_config(serialization_dir)
+
+    random.seed(config.seed)
+    np.random.seed(config.seed)
+    torch.manual_seed(config.seed)
+    logger.info(f"set seed: {config.seed}")
+
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    n_gpu = torch.cuda.device_count()
+    logger.info(f"GPU available count: {n_gpu}; device: {device}")
+
+    # datasets
+    train_dataset = DenoiseJsonlDataset(
+        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=225000,
+    )
+    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,
+    )
+    train_data_loader = DataLoader(
+        dataset=train_dataset,
+        batch_size=config.batch_size,
+        # shuffle=True,
+        sampler=None,
+        # Linux 系统中可以使用多个子进程加载数据, 而在 Windows 系统中不能.
+        num_workers=0 if platform.system() == "Windows" else os.cpu_count() // 2,
+        collate_fn=collate_fn,
+        pin_memory=False,
+        prefetch_factor=None if platform.system() == "Windows" else 2,
+    )
+    valid_data_loader = DataLoader(
+        dataset=valid_dataset,
+        batch_size=config.batch_size,
+        # shuffle=True,
+        sampler=None,
+        # Linux 系统中可以使用多个子进程加载数据, 而在 Windows 系统中不能.
+        num_workers=0 if platform.system() == "Windows" else os.cpu_count() // 2,
+        collate_fn=collate_fn,
+        pin_memory=False,
+        prefetch_factor=None if platform.system() == "Windows" else 2,
+    )
+
+    # models
+    logger.info(f"prepare models. config_file: {args.config_file}")
+    model = DTLNPretrainedModel(config).to(device)
+    model.to(device)
+    model.train()
+
+    # optimizer
+    logger.info("prepare optimizer, lr_scheduler, loss_fn, evaluation_metric")
+    optimizer = torch.optim.AdamW(model.parameters(), config.lr)
+
+    # resume training
+    last_step_idx = -1
+    last_epoch = -1
+    for step_idx_str in serialization_dir.glob("steps-*"):
+        step_idx_str = Path(step_idx_str)
+        step_idx = step_idx_str.stem.split("-")[1]
+        step_idx = int(step_idx)
+        if step_idx > last_step_idx:
+            last_step_idx = step_idx
+    # last_epoch = 1
+
+    if last_step_idx != -1:
+        logger.info(f"resume from steps-{last_step_idx}.")
+        model_pt = serialization_dir / f"steps-{last_step_idx}/model.pt"
+
+        logger.info(f"load state dict for model.")
+        with open(model_pt.as_posix(), "rb") as f:
+            state_dict = torch.load(f, map_location="cpu", weights_only=True)
+        model.load_state_dict(state_dict, strict=True)
+
+    if config.lr_scheduler == "CosineAnnealingLR":
+        lr_scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(
+            optimizer,
+            last_epoch=last_epoch,
+            # T_max=10 * config.eval_steps,
+            # eta_min=0.01 * config.lr,
+            **config.lr_scheduler_kwargs,
+        )
+    elif config.lr_scheduler == "MultiStepLR":
+        lr_scheduler = torch.optim.lr_scheduler.MultiStepLR(
+            optimizer,
+            last_epoch=last_epoch,
+            milestones=[10000, 20000, 30000, 40000, 50000], gamma=0.5
+        )
+    else:
+        raise AssertionError(f"invalid lr_scheduler: {config.lr_scheduler}")
+
+    neg_si_snr_loss_fn = NegativeSISNRLoss(reduction="mean").to(device)
+    mr_stft_loss_fn = MultiResolutionSTFTLoss(
+        fft_size_list=[256, 512, 1024],
+        win_size_list=[256, 512, 1024],
+        hop_size_list=[128, 256, 512],
+        factor_sc=1.5,
+        factor_mag=1.0,
+        reduction="mean"
+    ).to(device)
+
+    # training loop
+
+    # state
+    average_pesq_score = 1000000000
+    average_loss = 1000000000
+    average_mr_stft_loss = 1000000000
+    average_neg_si_snr_loss = 1000000000
+
+    model_list = list()
+    best_epoch_idx = None
+    best_step_idx = None
+    best_metric = None
+    patience_count = 0
+
+    step_idx = 0 if last_step_idx == -1 else last_step_idx
+
+    logger.info("training")
+    early_stop_flag = False
+    for epoch_idx in range(max(0, last_epoch+1), config.max_epochs):
+        if early_stop_flag:
+            break
+
+        # train
+        model.train()
+
+        total_pesq_score = 0.
+        total_loss = 0.
+        total_mr_stft_loss = 0.
+        total_neg_si_snr_loss = 0.
+        total_batches = 0.
+
+        progress_bar_train = tqdm(
+            initial=step_idx,
+            desc="Training; epoch-{}".format(epoch_idx),
+        )
+        for train_batch in train_data_loader:
+            clean_audios, noisy_audios = train_batch
+            clean_audios: torch.Tensor = clean_audios.to(device)
+            noisy_audios: torch.Tensor = noisy_audios.to(device)
+
+            denoise_audios = model.forward(noisy_audios)
+            denoise_audios = torch.squeeze(denoise_audios, dim=1)
+
+            mr_stft_loss = mr_stft_loss_fn.forward(denoise_audios, clean_audios)
+            neg_si_snr_loss = neg_si_snr_loss_fn.forward(denoise_audios, clean_audios)
+
+            loss = 1.0 * mr_stft_loss + 1.0 * neg_si_snr_loss
+            if torch.any(torch.isnan(loss)) or torch.any(torch.isinf(loss)):
+                logger.info(f"find nan or inf in loss.")
+                continue
+
+            denoise_audios_list_r = list(denoise_audios.detach().cpu().numpy())
+            clean_audios_list_r = list(clean_audios.detach().cpu().numpy())
+            pesq_score = run_pesq_score(clean_audios_list_r, denoise_audios_list_r, sample_rate=config.sample_rate, mode="nb")
+
+            optimizer.zero_grad()
+            loss.backward()
+            torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=config.clip_grad_norm)
+            optimizer.step()
+            lr_scheduler.step()
+
+            total_pesq_score += pesq_score
+            total_loss += loss.item()
+            total_mr_stft_loss += mr_stft_loss.item()
+            total_neg_si_snr_loss += neg_si_snr_loss.item()
+            total_batches += 1
+
+            average_pesq_score = round(total_pesq_score / total_batches, 4)
+            average_loss = round(total_loss / total_batches, 4)
+            average_mr_stft_loss = round(total_mr_stft_loss / total_batches, 4)
+            average_neg_si_snr_loss = round(total_neg_si_snr_loss / total_batches, 4)
+
+            progress_bar_train.update(1)
+            progress_bar_train.set_postfix({
+                "lr": lr_scheduler.get_last_lr()[0],
+                "pesq_score": average_pesq_score,
+                "loss": average_loss,
+                "mr_stft_loss": average_mr_stft_loss,
+                "neg_si_snr_loss": average_neg_si_snr_loss,
+            })
+
+            # evaluation
+            step_idx += 1
+            if step_idx % config.eval_steps == 0:
+                model.eval()
+                with torch.no_grad():
+                    torch.cuda.empty_cache()
+
+                    total_pesq_score = 0.
+                    total_loss = 0.
+                    total_mr_stft_loss = 0.
+                    total_neg_si_snr_loss = 0.
+                    total_batches = 0.
+
+                    progress_bar_train.close()
+                    progress_bar_eval = tqdm(
+                        desc="Evaluation; steps-{}k".format(int(step_idx/1000)),
+                    )
+                    for eval_batch in valid_data_loader:
+                        clean_audios, noisy_audios = eval_batch
+                        clean_audios: torch.Tensor = clean_audios.to(device)
+                        noisy_audios: torch.Tensor = noisy_audios.to(device)
+
+                        denoise_audios = model.forward(noisy_audios)
+                        denoise_audios = torch.squeeze(denoise_audios, dim=1)
+
+                        mr_stft_loss = mr_stft_loss_fn.forward(denoise_audios, clean_audios)
+                        neg_si_snr_loss = neg_si_snr_loss_fn.forward(denoise_audios, clean_audios)
+
+                        loss = 1.0 * mr_stft_loss + 1.0 * neg_si_snr_loss
+                        if torch.any(torch.isnan(loss)) or torch.any(torch.isinf(loss)):
+                            logger.info(f"find nan or inf in loss.")
+                            continue
+
+                        denoise_audios_list_r = list(denoise_audios.detach().cpu().numpy())
+                        clean_audios_list_r = list(clean_audios.detach().cpu().numpy())
+                        pesq_score = run_pesq_score(clean_audios_list_r, denoise_audios_list_r, sample_rate=config.sample_rate, mode="nb")
+
+                        total_pesq_score += pesq_score
+                        total_loss += loss.item()
+                        total_mr_stft_loss += mr_stft_loss.item()
+                        total_neg_si_snr_loss += neg_si_snr_loss.item()
+                        total_batches += 1
+
+                        average_pesq_score = round(total_pesq_score / total_batches, 4)
+                        average_loss = round(total_loss / total_batches, 4)
+                        average_mr_stft_loss = round(total_mr_stft_loss / total_batches, 4)
+                        average_neg_si_snr_loss = round(total_neg_si_snr_loss / total_batches, 4)
+
+                        progress_bar_eval.update(1)
+                        progress_bar_eval.set_postfix({
+                            "lr": lr_scheduler.get_last_lr()[0],
+                            "pesq_score": average_pesq_score,
+                            "loss": average_loss,
+                            "mr_stft_loss": average_mr_stft_loss,
+                            "neg_si_snr_loss": average_neg_si_snr_loss,
+
+                        })
+
+                    total_pesq_score = 0.
+                    total_loss = 0.
+                    total_mr_stft_loss = 0.
+                    total_neg_si_snr_loss = 0.
+                    total_batches = 0.
+
+                    progress_bar_eval.close()
+                    progress_bar_train = tqdm(
+                        initial=progress_bar_train.n,
+                        postfix=progress_bar_train.postfix,
+                        desc=progress_bar_train.desc,
+                    )
+
+                    # save path
+                    save_dir = serialization_dir / "steps-{}".format(step_idx)
+                    save_dir.mkdir(parents=True, exist_ok=False)
+
+                    # save models
+                    model.save_pretrained(save_dir.as_posix())
+
+                    model_list.append(save_dir)
+                    if len(model_list) >= args.num_serialized_models_to_keep:
+                        model_to_delete: Path = model_list.pop(0)
+                        shutil.rmtree(model_to_delete.as_posix())
+
+                    # save metric
+                    if best_metric is None:
+                        best_epoch_idx = epoch_idx
+                        best_step_idx = step_idx
+                        best_metric = average_pesq_score
+                    elif average_pesq_score >= best_metric:
+                        # great is better.
+                        best_epoch_idx = epoch_idx
+                        best_step_idx = step_idx
+                        best_metric = average_pesq_score
+                    else:
+                        pass
+
+                    metrics = {
+                        "epoch_idx": epoch_idx,
+                        "best_epoch_idx": best_epoch_idx,
+                        "best_step_idx": best_step_idx,
+                        "pesq_score": average_pesq_score,
+                        "loss": average_loss,
+                        "mr_stft_loss": average_mr_stft_loss,
+                        "neg_si_snr_loss": average_neg_si_snr_loss,
+                    }
+                    metrics_filename = save_dir / "metrics_epoch.json"
+                    with open(metrics_filename, "w", encoding="utf-8") as f:
+                        json.dump(metrics, f, indent=4, ensure_ascii=False)
+
+                    # save best
+                    best_dir = serialization_dir / "best"
+                    if best_epoch_idx == epoch_idx and best_step_idx == step_idx:
+                        if best_dir.exists():
+                            shutil.rmtree(best_dir)
+                        shutil.copytree(save_dir, best_dir)
+
+                    # early stop
+                    early_stop_flag = False
+                    if best_epoch_idx == epoch_idx and best_step_idx == step_idx:
+                        patience_count = 0
+                    else:
+                        patience_count += 1
+                    if patience_count >= args.patience:
+                        early_stop_flag = True
+
+                    # early stop
+                    if early_stop_flag:
+                        break
+                model.train()
+
+    return
+
+
+if __name__ == "__main__":
+    main()

examples/dtln/yaml/config-1024.yaml

@@ -0,0 +1,29 @@
+model_name: "DTLN"
+
+# spec
+sample_rate: 8000
+fft_size: 512
+hop_size: 128
+win_type: hann
+
+# data
+min_snr_db: -5
+max_snr_db: 25
+
+# model
+encoder_size: 1024
+
+# 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: 4
+batch_size: 64
+eval_steps: 15000

examples/dtln/yaml/config-256.yaml

@@ -0,0 +1,29 @@
+model_name: "DTLN"
+
+# spec
+sample_rate: 8000
+fft_size: 256
+hop_size: 128
+win_type: hann
+
+# data
+min_snr_db: -5
+max_snr_db: 25
+
+# model
+encoder_size: 256
+
+# 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: 4
+batch_size: 64
+eval_steps: 15000

examples/dtln/yaml/config-512.yaml

@@ -0,0 +1,29 @@
+model_name: "DTLN"
+
+# spec
+sample_rate: 8000
+fft_size: 512
+hop_size: 128
+win_type: hann
+
+# data
+min_snr_db: -5
+max_snr_db: 25
+
+# model
+encoder_size: 512
+
+# 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: 4
+batch_size: 64
+eval_steps: 15000

examples/dtln_mp3_to_wav/run.sh

@@ -0,0 +1,168 @@
+#!/usr/bin/env bash
+
+: <<'END'
+
+sh run.sh --stage 2 --stop_stage 2 --system_version centos --file_folder_name file_dir-256 --final_model_name dtln-256-nx-dns3 \
+--config_file "yaml/config-256.yaml" \
+--noise_dir "/data/tianxing/HuggingDatasets/nx_noise/data/noise" \
+--speech_dir "/data/tianxing/HuggingDatasets/nx_noise/data/speech/dns3-speech"
+
+
+sh run.sh --stage 1 --stop_stage 2 --system_version centos --file_folder_name file_dir-512 --final_model_name dtln-512-nx-dns3 \
+--config_file "yaml/config-512.yaml" \
+--noise_dir "/data/tianxing/HuggingDatasets/nx_noise/data/noise" \
+--speech_dir "/data/tianxing/HuggingDatasets/nx_noise/data/speech/dns3-speech"
+
+
+sh run.sh --stage 2 --stop_stage 2 --system_version centos --file_folder_name dtnl-1024-nx2 --final_model_name dtln-1024-nx2 \
+--config_file "yaml/config-1024.yaml" \
+--noise_dir "/data/tianxing/HuggingDatasets/nx_noise/data/noise/nx-noise" \
+--speech_dir "/data/tianxing/HuggingDatasets/nx_noise/data/speech/nx-speech2"
+
+
+bash run.sh --stage 1 --stop_stage 2 --system_version centos --file_folder_name dtln-256-nx2-dns3-mp3 --final_model_name dtln-256-nx2-dns3-mp3 \
+--config_file "yaml/config-256.yaml" \
+--audio_dir "/data/tianxing/HuggingDatasets/nx_noise/data" \
+
+
+END
+
+
+# params
+system_version="windows";
+verbose=true;
+stage=0 # start from 0 if you need to start from data preparation
+stop_stage=9
+
+work_dir="$(pwd)"
+file_folder_name=file_folder_name
+final_model_name=final_model_name
+config_file="yaml/config.yaml"
+limit=10
+
+audio_dir=/data/tianxing/HuggingDatasets/nx_noise/data
+
+max_count=-1
+
+nohup_name=nohup.out
+
+# model params
+batch_size=64
+max_epochs=200
+save_top_k=10
+patience=5
+
+
+# parse options
+while true; do
+  [ -z "${1:-}" ] && break;  # break if there are no arguments
+  case "$1" in
+    --*) name=$(echo "$1" | sed s/^--// | sed s/-/_/g);
+      eval '[ -z "${'"$name"'+xxx}" ]' && echo "$0: invalid option $1" 1>&2 && exit 1;
+      old_value="(eval echo \\$$name)";
+      if [ "${old_value}" == "true" ] || [ "${old_value}" == "false" ]; then
+        was_bool=true;
+      else
+        was_bool=false;
+      fi
+
+      # Set the variable to the right value-- the escaped quotes make it work if
+      # the option had spaces, like --cmd "queue.pl -sync y"
+      eval "${name}=\"$2\"";
+
+      # Check that Boolean-valued arguments are really Boolean.
+      if $was_bool && [[ "$2" != "true" && "$2" != "false" ]]; then
+        echo "$0: expected \"true\" or \"false\": $1 $2" 1>&2
+        exit 1;
+      fi
+      shift 2;
+      ;;
+
+    *) break;
+  esac
+done
+
+file_dir="${work_dir}/${file_folder_name}"
+final_model_dir="${work_dir}/../../trained_models/${final_model_name}";
+evaluation_audio_dir="${file_dir}/evaluation_audio"
+
+train_dataset="${file_dir}/train.jsonl"
+valid_dataset="${file_dir}/valid.jsonl"
+
+$verbose && echo "system_version: ${system_version}"
+$verbose && echo "file_folder_name: ${file_folder_name}"
+
+if [ $system_version == "windows" ]; then
+  alias python3='D:/Users/tianx/PycharmProjects/virtualenv/nx_denoise/Scripts/python.exe'
+elif [ $system_version == "centos" ] || [ $system_version == "ubuntu" ]; then
+  #source /data/local/bin/nx_denoise/bin/activate
+  alias python3='/data/local/bin/nx_denoise/bin/python3'
+fi
+
+
+if [ ${stage} -le 1 ] && [ ${stop_stage} -ge 1 ]; then
+  $verbose && echo "stage 1: prepare data"
+  cd "${work_dir}" || exit 1
+  python3 step_1_prepare_data.py \
+  --file_dir "${file_dir}" \
+  --audio_dir "${audio_dir}" \
+  --train_dataset "${train_dataset}" \
+  --valid_dataset "${valid_dataset}" \
+  --max_count "${max_count}" \
+
+fi
+
+
+if [ ${stage} -le 2 ] && [ ${stop_stage} -ge 2 ]; then
+  $verbose && echo "stage 2: train model"
+  cd "${work_dir}" || exit 1
+  python3 step_2_train_model.py \
+  --train_dataset "${train_dataset}" \
+  --valid_dataset "${valid_dataset}" \
+  --serialization_dir "${file_dir}" \
+  --config_file "${config_file}" \
+
+fi
+
+
+if [ ${stage} -le 3 ] && [ ${stop_stage} -ge 3 ]; then
+  $verbose && echo "stage 3: test model"
+  cd "${work_dir}" || exit 1
+  python3 step_3_evaluation.py \
+  --valid_dataset "${valid_dataset}" \
+  --model_dir "${file_dir}/best" \
+  --evaluation_audio_dir "${evaluation_audio_dir}" \
+  --limit "${limit}" \
+
+fi
+
+
+if [ ${stage} -le 4 ] && [ ${stop_stage} -ge 4 ]; then
+  $verbose && echo "stage 4: collect files"
+  cd "${work_dir}" || exit 1
+
+  mkdir -p ${final_model_dir}
+
+  cp "${file_dir}/best"/* "${final_model_dir}"
+  cp -r "${file_dir}/evaluation_audio" "${final_model_dir}"
+
+  cd "${final_model_dir}/.." || exit 1;
+
+  if [ -e "${final_model_name}.zip" ]; then
+    rm -rf "${final_model_name}_backup.zip"
+    mv "${final_model_name}.zip" "${final_model_name}_backup.zip"
+  fi
+
+  zip -r "${final_model_name}.zip" "${final_model_name}"
+  rm -rf "${final_model_name}"
+
+fi
+
+
+if [ ${stage} -le 5 ] && [ ${stop_stage} -ge 5 ]; then
+  $verbose && echo "stage 5: clear file_dir"
+  cd "${work_dir}" || exit 1
+
+  rm -rf "${file_dir}";
+
+fi

examples/dtln_mp3_to_wav/step_1_prepare_data.py

@@ -0,0 +1,127 @@
+#!/usr/bin/python3
+# -*- coding: utf-8 -*-
+import argparse
+import json
+import os
+from pathlib import Path
+import random
+import sys
+
+pwd = os.path.abspath(os.path.dirname(__file__))
+sys.path.append(os.path.join(pwd, "../../"))
+
+import librosa
+import numpy as np
+from tqdm import tqdm
+
+
+def get_args():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--file_dir", default="./", type=str)
+
+    parser.add_argument(
+        "--audio_dir",
+        default="E:/Users/tianx/HuggingDatasets/nx_noise/data/speech",
+        type=str
+    )
+
+    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("--target_sample_rate", default=8000, type=int)
+
+    parser.add_argument("--max_count", default=-1, type=int)
+
+    args = parser.parse_args()
+    return args
+
+
+def target_second_signal_generator(data_dir: str, duration: int = 2, sample_rate: int = 8000, max_epoch: int = 1):
+    data_dir = Path(data_dir)
+    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):
+                if np.sum(signal[begin: begin+win_size]) == 0:
+                    continue
+                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 main():
+    args = get_args()
+
+    file_dir = Path(args.file_dir)
+    file_dir.mkdir(exist_ok=True)
+
+    audio_dir = Path(args.audio_dir)
+
+    audio_generator = target_second_signal_generator(
+        audio_dir.as_posix(),
+        duration=args.duration,
+        sample_rate=args.target_sample_rate,
+        max_epoch=1,
+    )
+    count = 0
+    process_bar = tqdm(desc="build dataset jsonl")
+    with open(args.train_dataset, "w", encoding="utf-8") as ftrain, open(args.valid_dataset, "w", encoding="utf-8") as fvalid:
+        for audio in audio_generator:
+            if count >= args.max_count > 0:
+                break
+
+            filename = audio["filename"]
+            raw_duration = audio["raw_duration"]
+            offset = audio["offset"]
+            duration = audio["duration"]
+
+            random1 = random.random()
+            random2 = random.random()
+
+            row = {
+                "count": count,
+
+                "filename": filename,
+                "raw_duration": raw_duration,
+                "offset": offset,
+                "duration": duration,
+
+                "random1": random1,
+            }
+            row = json.dumps(row, ensure_ascii=False)
+            if random2 < (1 / 300):
+                fvalid.write(f"{row}\n")
+            else:
+                ftrain.write(f"{row}\n")
+
+            count += 1
+            duration_seconds = count * args.duration
+            duration_hours = duration_seconds / 3600
+
+            process_bar.update(n=1)
+            process_bar.set_postfix({
+                "duration_hours": round(duration_hours, 4),
+            })
+
+    return
+
+
+if __name__ == "__main__":
+    main()

examples/dtln_mp3_to_wav/step_2_train_model.py

@@ -0,0 +1,445 @@
+#!/usr/bin/python3
+# -*- coding: utf-8 -*-
+"""
+https://github.com/breizhn/DTLN
+
+"""
+import argparse
+import json
+import logging
+from logging.handlers import TimedRotatingFileHandler
+import os
+import platform
+from pathlib import Path
+import random
+import sys
+import shutil
+from typing import List
+
+pwd = os.path.abspath(os.path.dirname(__file__))
+sys.path.append(os.path.join(pwd, "../../"))
+
+import numpy as np
+import torch
+import torch.nn as nn
+from torch.nn import functional as F
+from torch.utils.data.dataloader import DataLoader
+from tqdm import tqdm
+
+from toolbox.torch.utils.data.dataset.mp3_to_wav_jsonl_dataset import Mp3ToWavJsonlDataset
+from toolbox.torchaudio.losses.snr import NegativeSISNRLoss
+from toolbox.torchaudio.losses.spectral import LSDLoss, MultiResolutionSTFTLoss
+from toolbox.torchaudio.metrics.pesq import run_pesq_score
+from toolbox.torchaudio.models.dtln.configuration_dtln import DTLNConfig
+from toolbox.torchaudio.models.dtln.modeling_dtln import DTLNModel, DTLNPretrainedModel
+
+
+def get_args():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--train_dataset", default="train.jsonl", type=str)
+    parser.add_argument("--valid_dataset", default="valid.jsonl", type=str)
+
+    parser.add_argument("--num_serialized_models_to_keep", default=15, type=int)
+    parser.add_argument("--patience", default=30, type=int)
+    parser.add_argument("--serialization_dir", default="serialization_dir", type=str)
+
+    parser.add_argument("--config_file", default="config.yaml", type=str)
+
+    args = parser.parse_args()
+    return args
+
+
+def logging_config(file_dir: str):
+    fmt = "%(asctime)s - %(name)s - %(levelname)s  %(filename)s:%(lineno)d >  %(message)s"
+
+    logging.basicConfig(format=fmt,
+                        datefmt="%m/%d/%Y %H:%M:%S",
+                        level=logging.INFO)
+    file_handler = TimedRotatingFileHandler(
+        filename=os.path.join(file_dir, "main.log"),
+        encoding="utf-8",
+        when="D",
+        interval=1,
+        backupCount=7
+    )
+    file_handler.setLevel(logging.INFO)
+    file_handler.setFormatter(logging.Formatter(fmt))
+    logger = logging.getLogger(__name__)
+    logger.addHandler(file_handler)
+
+    return logger
+
+
+class CollateFunction(object):
+    def __init__(self):
+        pass
+
+    def __call__(self, batch: List[dict]):
+        mp3_waveform_list = list()
+        wav_waveform_list = list()
+
+        for sample in batch:
+            mp3_waveform: torch.Tensor = sample["mp3_waveform"]
+            wav_waveform: torch.Tensor = sample["wav_waveform"]
+
+            mp3_waveform_list.append(mp3_waveform)
+            wav_waveform_list.append(wav_waveform)
+
+        mp3_waveform_list = torch.stack(mp3_waveform_list)
+        wav_waveform_list = torch.stack(wav_waveform_list)
+
+        # assert
+        if torch.any(torch.isnan(mp3_waveform_list)) or torch.any(torch.isinf(mp3_waveform_list)):
+            raise AssertionError("nan or inf in mp3_waveform_list")
+        if torch.any(torch.isnan(wav_waveform_list)) or torch.any(torch.isinf(wav_waveform_list)):
+            raise AssertionError("nan or inf in wav_waveform_list")
+
+        return mp3_waveform_list, wav_waveform_list
+
+
+collate_fn = CollateFunction()
+
+
+def main():
+    args = get_args()
+
+    config = DTLNConfig.from_pretrained(
+        pretrained_model_name_or_path=args.config_file,
+    )
+
+    serialization_dir = Path(args.serialization_dir)
+    serialization_dir.mkdir(parents=True, exist_ok=True)
+
+    logger = logging_config(serialization_dir)
+
+    random.seed(config.seed)
+    np.random.seed(config.seed)
+    torch.manual_seed(config.seed)
+    logger.info(f"set seed: {config.seed}")
+
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    n_gpu = torch.cuda.device_count()
+    logger.info(f"GPU available count: {n_gpu}; device: {device}")
+
+    # datasets
+    train_dataset = Mp3ToWavJsonlDataset(
+        jsonl_file=args.train_dataset,
+        expected_sample_rate=config.sample_rate,
+        max_wave_value=32768.0,
+        # skip=225000,
+    )
+    valid_dataset = Mp3ToWavJsonlDataset(
+        jsonl_file=args.valid_dataset,
+        expected_sample_rate=config.sample_rate,
+        max_wave_value=32768.0,
+    )
+    train_data_loader = DataLoader(
+        dataset=train_dataset,
+        batch_size=config.batch_size,
+        # shuffle=True,
+        sampler=None,
+        # Linux 系统中可以使用多个子进程加载数据, 而在 Windows 系统中不能.
+        num_workers=0 if platform.system() == "Windows" else os.cpu_count() // 2,
+        collate_fn=collate_fn,
+        pin_memory=False,
+        prefetch_factor=None if platform.system() == "Windows" else 2,
+    )
+    valid_data_loader = DataLoader(
+        dataset=valid_dataset,
+        batch_size=config.batch_size,
+        # shuffle=True,
+        sampler=None,
+        # Linux 系统中可以使用多个子进程加载数据, 而在 Windows 系统中不能.
+        num_workers=0 if platform.system() == "Windows" else os.cpu_count() // 2,
+        collate_fn=collate_fn,
+        pin_memory=False,
+        prefetch_factor=None if platform.system() == "Windows" else 2,
+    )
+
+    # models
+    logger.info(f"prepare models. config_file: {args.config_file}")
+    model = DTLNPretrainedModel(config).to(device)
+    model.to(device)
+    model.train()
+
+    # optimizer
+    logger.info("prepare optimizer, lr_scheduler, loss_fn, evaluation_metric")
+    optimizer = torch.optim.AdamW(model.parameters(), config.lr)
+
+    # resume training
+    last_step_idx = -1
+    last_epoch = -1
+    for step_idx_str in serialization_dir.glob("steps-*"):
+        step_idx_str = Path(step_idx_str)
+        step_idx = step_idx_str.stem.split("-")[1]
+        step_idx = int(step_idx)
+        if step_idx > last_step_idx:
+            last_step_idx = step_idx
+    # last_epoch = 1
+
+    if last_step_idx != -1:
+        logger.info(f"resume from steps-{last_step_idx}.")
+        model_pt = serialization_dir / f"steps-{last_step_idx}/model.pt"
+
+        logger.info(f"load state dict for model.")
+        with open(model_pt.as_posix(), "rb") as f:
+            state_dict = torch.load(f, map_location="cpu", weights_only=True)
+        model.load_state_dict(state_dict, strict=True)
+
+    if config.lr_scheduler == "CosineAnnealingLR":
+        lr_scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(
+            optimizer,
+            last_epoch=last_epoch,
+            # T_max=10 * config.eval_steps,
+            # eta_min=0.01 * config.lr,
+            **config.lr_scheduler_kwargs,
+        )
+    elif config.lr_scheduler == "MultiStepLR":
+        lr_scheduler = torch.optim.lr_scheduler.MultiStepLR(
+            optimizer,
+            last_epoch=last_epoch,
+            milestones=[10000, 20000, 30000, 40000, 50000], gamma=0.5
+        )
+    else:
+        raise AssertionError(f"invalid lr_scheduler: {config.lr_scheduler}")
+
+    neg_si_snr_loss_fn = NegativeSISNRLoss(reduction="mean").to(device)
+    mr_stft_loss_fn = MultiResolutionSTFTLoss(
+        fft_size_list=[256, 512, 1024],
+        win_size_list=[256, 512, 1024],
+        hop_size_list=[128, 256, 512],
+        factor_sc=1.5,
+        factor_mag=1.0,
+        reduction="mean"
+    ).to(device)
+    audio_l1_loss_fn = nn.L1Loss(reduction="mean")
+
+    # training loop
+
+    # state
+    average_pesq_score = 1000000000
+    average_loss = 1000000000
+    average_mr_stft_loss = 1000000000
+    average_audio_l1_loss = 1000000000
+    average_neg_si_snr_loss = 1000000000
+
+    model_list = list()
+    best_epoch_idx = None
+    best_step_idx = None
+    best_metric = None
+    patience_count = 0
+
+    step_idx = 0 if last_step_idx == -1 else last_step_idx
+
+    logger.info("training")
+    early_stop_flag = False
+    for epoch_idx in range(max(0, last_epoch+1), config.max_epochs):
+        if early_stop_flag:
+            break
+
+        # train
+        model.train()
+
+        total_pesq_score = 0.
+        total_loss = 0.
+        total_mr_stft_loss = 0.
+        total_audio_l1_loss = 0.
+        total_neg_si_snr_loss = 0.
+        total_batches = 0.
+
+        progress_bar_train = tqdm(
+            initial=step_idx,
+            desc="Training; epoch-{}".format(epoch_idx),
+        )
+        for train_batch in train_data_loader:
+            mp3_audios, wav_audios = train_batch
+            noisy_audios: torch.Tensor = mp3_audios.to(device)
+            clean_audios: torch.Tensor = wav_audios.to(device)
+
+            denoise_audios = model.forward(noisy_audios)
+            denoise_audios = torch.squeeze(denoise_audios, dim=1)
+
+            mr_stft_loss = mr_stft_loss_fn.forward(denoise_audios, clean_audios)
+            audio_l1_loss = audio_l1_loss_fn.forward(denoise_audios, clean_audios)
+            neg_si_snr_loss = neg_si_snr_loss_fn.forward(denoise_audios, clean_audios)
+
+            loss = 1.0 * mr_stft_loss + 1.0 * audio_l1_loss + 1.0 * neg_si_snr_loss
+            if torch.any(torch.isnan(loss)) or torch.any(torch.isinf(loss)):
+                logger.info(f"find nan or inf in loss.")
+                continue
+
+            denoise_audios_list_r = list(denoise_audios.detach().cpu().numpy())
+            clean_audios_list_r = list(clean_audios.detach().cpu().numpy())
+            pesq_score = run_pesq_score(clean_audios_list_r, denoise_audios_list_r, sample_rate=config.sample_rate, mode="nb")
+
+            optimizer.zero_grad()
+            loss.backward()
+            torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=config.clip_grad_norm)
+            optimizer.step()
+            lr_scheduler.step()
+
+            total_pesq_score += pesq_score
+            total_loss += loss.item()
+            total_mr_stft_loss += mr_stft_loss.item()
+            total_audio_l1_loss += audio_l1_loss.item()
+            total_neg_si_snr_loss += neg_si_snr_loss.item()
+            total_batches += 1
+
+            average_pesq_score = round(total_pesq_score / total_batches, 4)
+            average_loss = round(total_loss / total_batches, 4)
+            average_mr_stft_loss = round(total_mr_stft_loss / total_batches, 4)
+            average_audio_l1_loss = round(total_audio_l1_loss / total_batches, 4)
+            average_neg_si_snr_loss = round(total_neg_si_snr_loss / total_batches, 4)
+
+            progress_bar_train.update(1)
+            progress_bar_train.set_postfix({
+                "lr": lr_scheduler.get_last_lr()[0],
+                "pesq_score": average_pesq_score,
+                "loss": average_loss,
+                "mr_stft_loss": average_mr_stft_loss,
+                "audio_l1_loss": average_audio_l1_loss,
+                "neg_si_snr_loss": average_neg_si_snr_loss,
+            })
+
+            # evaluation
+            step_idx += 1
+            if step_idx % config.eval_steps == 0:
+                model.eval()
+                with torch.no_grad():
+                    torch.cuda.empty_cache()
+
+                    total_pesq_score = 0.
+                    total_loss = 0.
+                    total_mr_stft_loss = 0.
+                    total_audio_l1_loss = 0.
+                    total_neg_si_snr_loss = 0.
+                    total_batches = 0.
+
+                    progress_bar_train.close()
+                    progress_bar_eval = tqdm(
+                        desc="Evaluation; steps-{}k".format(int(step_idx/1000)),
+                    )
+                    for eval_batch in valid_data_loader:
+                        mp3_audios, wav_audios = eval_batch
+                        noisy_audios: torch.Tensor = mp3_audios.to(device)
+                        clean_audios: torch.Tensor = wav_audios.to(device)
+
+                        denoise_audios = model.forward(noisy_audios)
+                        denoise_audios = torch.squeeze(denoise_audios, dim=1)
+
+                        mr_stft_loss = mr_stft_loss_fn.forward(denoise_audios, clean_audios)
+                        audio_l1_loss = audio_l1_loss_fn.forward(denoise_audios, clean_audios)
+                        neg_si_snr_loss = neg_si_snr_loss_fn.forward(denoise_audios, clean_audios)
+
+                        loss = 1.0 * mr_stft_loss + 1.0 * audio_l1_loss + 1.0 * neg_si_snr_loss
+                        if torch.any(torch.isnan(loss)) or torch.any(torch.isinf(loss)):
+                            logger.info(f"find nan or inf in loss.")
+                            continue
+
+                        denoise_audios_list_r = list(denoise_audios.detach().cpu().numpy())
+                        clean_audios_list_r = list(clean_audios.detach().cpu().numpy())
+                        pesq_score = run_pesq_score(clean_audios_list_r, denoise_audios_list_r, sample_rate=config.sample_rate, mode="nb")
+
+                        total_pesq_score += pesq_score
+                        total_loss += loss.item()
+                        total_mr_stft_loss += mr_stft_loss.item()
+                        total_audio_l1_loss += audio_l1_loss.item()
+                        total_neg_si_snr_loss += neg_si_snr_loss.item()
+                        total_batches += 1
+
+                        average_pesq_score = round(total_pesq_score / total_batches, 4)
+                        average_loss = round(total_loss / total_batches, 4)
+                        average_mr_stft_loss = round(total_mr_stft_loss / total_batches, 4)
+                        average_audio_l1_loss = round(total_audio_l1_loss / total_batches, 4)
+                        average_neg_si_snr_loss = round(total_neg_si_snr_loss / total_batches, 4)
+
+                        progress_bar_eval.update(1)
+                        progress_bar_eval.set_postfix({
+                            "lr": lr_scheduler.get_last_lr()[0],
+                            "pesq_score": average_pesq_score,
+                            "loss": average_loss,
+                            "mr_stft_loss": average_mr_stft_loss,
+                            "audio_l1_loss": average_audio_l1_loss,
+                            "neg_si_snr_loss": average_neg_si_snr_loss,
+
+                        })
+
+                    total_pesq_score = 0.
+                    total_loss = 0.
+                    total_mr_stft_loss = 0.
+                    total_audio_l1_loss = 0.
+                    total_neg_si_snr_loss = 0.
+                    total_batches = 0.
+
+                    progress_bar_eval.close()
+                    progress_bar_train = tqdm(
+                        initial=progress_bar_train.n,
+                        postfix=progress_bar_train.postfix,
+                        desc=progress_bar_train.desc,
+                    )
+
+                    # save path
+                    save_dir = serialization_dir / "steps-{}".format(step_idx)
+                    save_dir.mkdir(parents=True, exist_ok=False)
+
+                    # save models
+                    model.save_pretrained(save_dir.as_posix())
+
+                    model_list.append(save_dir)
+                    if len(model_list) >= args.num_serialized_models_to_keep:
+                        model_to_delete: Path = model_list.pop(0)
+                        shutil.rmtree(model_to_delete.as_posix())
+
+                    # save metric
+                    if best_metric is None:
+                        best_epoch_idx = epoch_idx
+                        best_step_idx = step_idx
+                        best_metric = average_pesq_score
+                    elif average_pesq_score >= best_metric:
+                        # great is better.
+                        best_epoch_idx = epoch_idx
+                        best_step_idx = step_idx
+                        best_metric = average_pesq_score
+                    else:
+                        pass
+
+                    metrics = {
+                        "epoch_idx": epoch_idx,
+                        "best_epoch_idx": best_epoch_idx,
+                        "best_step_idx": best_step_idx,
+                        "pesq_score": average_pesq_score,
+                        "loss": average_loss,
+                        "mr_stft_loss": average_mr_stft_loss,
+                        "audio_l1_loss": average_audio_l1_loss,
+                        "neg_si_snr_loss": average_neg_si_snr_loss,
+                    }
+                    metrics_filename = save_dir / "metrics_epoch.json"
+                    with open(metrics_filename, "w", encoding="utf-8") as f:
+                        json.dump(metrics, f, indent=4, ensure_ascii=False)
+
+                    # save best
+                    best_dir = serialization_dir / "best"
+                    if best_epoch_idx == epoch_idx and best_step_idx == step_idx:
+                        if best_dir.exists():
+                            shutil.rmtree(best_dir)
+                        shutil.copytree(save_dir, best_dir)
+
+                    # early stop
+                    early_stop_flag = False
+                    if best_epoch_idx == epoch_idx and best_step_idx == step_idx:
+                        patience_count = 0
+                    else:
+                        patience_count += 1
+                    if patience_count >= args.patience:
+                        early_stop_flag = True
+
+                    # early stop
+                    if early_stop_flag:
+                        break
+                model.train()
+
+    return
+
+
+if __name__ == "__main__":
+    main()

examples/dtln_mp3_to_wav/yaml/config-1024.yaml

@@ -0,0 +1,29 @@
+model_name: "DTLN"
+
+# spec
+sample_rate: 8000
+fft_size: 512
+hop_size: 128
+win_type: hann
+
+# data
+min_snr_db: -5
+max_snr_db: 25
+
+# model
+encoder_size: 1024
+
+# 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: 4
+batch_size: 64
+eval_steps: 15000

examples/dtln_mp3_to_wav/yaml/config-256.yaml

@@ -0,0 +1,29 @@
+model_name: "DTLN"
+
+# spec
+sample_rate: 8000
+fft_size: 256
+hop_size: 128
+win_type: hann
+
+# data
+min_snr_db: -5
+max_snr_db: 25
+
+# model
+encoder_size: 256
+
+# 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: 4
+batch_size: 64
+eval_steps: 15000

examples/dtln_mp3_to_wav/yaml/config-512.yaml

@@ -0,0 +1,29 @@
+model_name: "DTLN"
+
+# spec
+sample_rate: 8000
+fft_size: 512
+hop_size: 128
+win_type: hann
+
+# data
+min_snr_db: -5
+max_snr_db: 25
+
+# model
+encoder_size: 512
+
+# 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: 4
+batch_size: 64
+eval_steps: 15000

examples/frcrn/run.sh

@@ -0,0 +1,159 @@
+#!/usr/bin/env bash
+
+: <<'END'
+
+
+sh run.sh --stage 1 --stop_stage 2 --system_version centos --file_folder_name file_dir --final_model_name frcrn-20-512-nx-dns3 \
+--config_file "yaml/config-10.yaml" \
+--noise_dir "/data/tianxing/HuggingDatasets/nx_noise/data/noise" \
+--speech_dir "/data/tianxing/HuggingDatasets/nx_noise/data/speech"
+
+
+sh run.sh --stage 2 --stop_stage 2 --system_version centos --file_folder_name file_dir --final_model_name frcrn-10-nx2 \
+--config_file "yaml/config-10.yaml" \
+--noise_dir "/data/tianxing/HuggingDatasets/nx_noise/data/noise/nx-noise" \
+--speech_dir "/data/tianxing/HuggingDatasets/nx_noise/data/speech/nx-speech2"
+
+END
+
+
+# params
+system_version="windows";
+verbose=true;
+stage=0 # start from 0 if you need to start from data preparation
+stop_stage=9
+
+work_dir="$(pwd)"
+file_folder_name=file_folder_name
+final_model_name=final_model_name
+config_file="yaml/config.yaml"
+limit=10
+
+noise_dir=/data/tianxing/HuggingDatasets/nx_noise/data/noise
+speech_dir=/data/tianxing/HuggingDatasets/aishell/data_aishell/wav/train
+
+max_count=10000000
+
+nohup_name=nohup.out
+
+# model params
+batch_size=64
+max_epochs=200
+save_top_k=10
+patience=5
+
+
+# parse options
+while true; do
+  [ -z "${1:-}" ] && break;  # break if there are no arguments
+  case "$1" in
+    --*) name=$(echo "$1" | sed s/^--// | sed s/-/_/g);
+      eval '[ -z "${'"$name"'+xxx}" ]' && echo "$0: invalid option $1" 1>&2 && exit 1;
+      old_value="(eval echo \\$$name)";
+      if [ "${old_value}" == "true" ] || [ "${old_value}" == "false" ]; then
+        was_bool=true;
+      else
+        was_bool=false;
+      fi
+
+      # Set the variable to the right value-- the escaped quotes make it work if
+      # the option had spaces, like --cmd "queue.pl -sync y"
+      eval "${name}=\"$2\"";
+
+      # Check that Boolean-valued arguments are really Boolean.
+      if $was_bool && [[ "$2" != "true" && "$2" != "false" ]]; then
+        echo "$0: expected \"true\" or \"false\": $1 $2" 1>&2
+        exit 1;
+      fi
+      shift 2;
+      ;;
+
+    *) break;
+  esac
+done
+
+file_dir="${work_dir}/${file_folder_name}"
+final_model_dir="${work_dir}/../../trained_models/${final_model_name}";
+evaluation_audio_dir="${file_dir}/evaluation_audio"
+
+train_dataset="${file_dir}/train.jsonl"
+valid_dataset="${file_dir}/valid.jsonl"
+
+$verbose && echo "system_version: ${system_version}"
+$verbose && echo "file_folder_name: ${file_folder_name}"
+
+if [ $system_version == "windows" ]; then
+  alias python3='D:/Users/tianx/PycharmProjects/virtualenv/nx_denoise/Scripts/python.exe'
+elif [ $system_version == "centos" ] || [ $system_version == "ubuntu" ]; then
+  #source /data/local/bin/nx_denoise/bin/activate
+  alias python3='/data/local/bin/nx_denoise/bin/python3'
+fi
+
+
+if [ ${stage} -le 1 ] && [ ${stop_stage} -ge 1 ]; then
+  $verbose && echo "stage 1: prepare data"
+  cd "${work_dir}" || exit 1
+  python3 step_1_prepare_data.py \
+  --file_dir "${file_dir}" \
+  --noise_dir "${noise_dir}" \
+  --speech_dir "${speech_dir}" \
+  --train_dataset "${train_dataset}" \
+  --valid_dataset "${valid_dataset}" \
+  --max_count "${max_count}" \
+
+fi
+
+
+if [ ${stage} -le 2 ] && [ ${stop_stage} -ge 2 ]; then
+  $verbose && echo "stage 2: train model"
+  cd "${work_dir}" || exit 1
+  python3 step_2_train_model.py \
+  --train_dataset "${train_dataset}" \
+  --valid_dataset "${valid_dataset}" \
+  --serialization_dir "${file_dir}" \
+  --config_file "${config_file}" \
+
+fi
+
+
+if [ ${stage} -le 3 ] && [ ${stop_stage} -ge 3 ]; then
+  $verbose && echo "stage 3: test model"
+  cd "${work_dir}" || exit 1
+  python3 step_3_evaluation.py \
+  --valid_dataset "${valid_dataset}" \
+  --model_dir "${file_dir}/best" \
+  --evaluation_audio_dir "${evaluation_audio_dir}" \
+  --limit "${limit}" \
+
+fi
+
+
+if [ ${stage} -le 4 ] && [ ${stop_stage} -ge 4 ]; then
+  $verbose && echo "stage 4: collect files"
+  cd "${work_dir}" || exit 1
+
+  mkdir -p ${final_model_dir}
+
+  cp "${file_dir}/best"/* "${final_model_dir}"
+  cp -r "${file_dir}/evaluation_audio" "${final_model_dir}"
+
+  cd "${final_model_dir}/.." || exit 1;
+
+  if [ -e "${final_model_name}.zip" ]; then
+    rm -rf "${final_model_name}_backup.zip"
+    mv "${final_model_name}.zip" "${final_model_name}_backup.zip"
+  fi
+
+  zip -r "${final_model_name}.zip" "${final_model_name}"
+  rm -rf "${final_model_name}"
+
+fi
+
+
+if [ ${stage} -le 5 ] && [ ${stop_stage} -ge 5 ]; then
+  $verbose && echo "stage 5: clear file_dir"
+  cd "${work_dir}" || exit 1
+
+  rm -rf "${file_dir}";
+
+fi

examples/frcrn/step_1_prepare_data.py

@@ -0,0 +1,164 @@
+#!/usr/bin/python3
+# -*- coding: utf-8 -*-
+import argparse
+import json
+import os
+from pathlib import Path
+import random
+import sys
+
+pwd = os.path.abspath(os.path.dirname(__file__))
+sys.path.append(os.path.join(pwd, "../../"))
+
+import librosa
+import numpy as np
+from tqdm import tqdm
+
+
+def get_args():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--file_dir", default="./", type=str)
+
+    parser.add_argument(
+        "--noise_dir",
+        default=r"E:\Users\tianx\HuggingDatasets\nx_noise\data\noise",
+        type=str
+    )
+    parser.add_argument(
+        "--speech_dir",
+        default=r"E:\programmer\asr_datasets\aishell\data_aishell\wav\train",
+        type=str
+    )
+
+    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=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=-1, type=int)
+
+    args = parser.parse_args()
+    return args
+
+
+def filename_generator(data_dir: str):
+    data_dir = Path(data_dir)
+    for filename in data_dir.glob("**/*.wav"):
+        yield filename.as_posix()
+
+
+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 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):
+                if np.sum(signal[begin: begin+win_size]) == 0:
+                    continue
+                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 main():
+    args = get_args()
+
+    file_dir = Path(args.file_dir)
+    file_dir.mkdir(exist_ok=True)
+
+    noise_dir = Path(args.noise_dir)
+    speech_dir = Path(args.speech_dir)
+
+    noise_generator = target_second_signal_generator(
+        noise_dir.as_posix(),
+        duration=args.duration,
+        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,
+        max_epoch=1,
+    )
+
+    dataset = list()
+
+    count = 0
+    process_bar = tqdm(desc="build dataset jsonl")
+    with open(args.train_dataset, "w", encoding="utf-8") as ftrain, open(args.valid_dataset, "w", encoding="utf-8") as fvalid:
+        for noise, speech in zip(noise_generator, speech_generator):
+            if count >= args.max_count > 0:
+                break
+
+            noise_filename = noise["filename"]
+            noise_raw_duration = noise["raw_duration"]
+            noise_offset = noise["offset"]
+            noise_duration = noise["duration"]
+
+            speech_filename = speech["filename"]
+            speech_raw_duration = speech["raw_duration"]
+            speech_offset = speech["offset"]
+            speech_duration = speech["duration"]
+
+            random1 = random.random()
+            random2 = random.random()
+
+            row = {
+                "count": count,
+
+                "noise_filename": noise_filename,
+                "noise_raw_duration": noise_raw_duration,
+                "noise_offset": noise_offset,
+                "noise_duration": noise_duration,
+
+                "speech_filename": speech_filename,
+                "speech_raw_duration": speech_raw_duration,
+                "speech_offset": speech_offset,
+                "speech_duration": speech_duration,
+
+                "snr_db": random.uniform(args.min_snr_db, args.max_snr_db),
+
+                "random1": random1,
+            }
+            row = json.dumps(row, ensure_ascii=False)
+            if random2 < (1 / 300 / 1):
+                fvalid.write(f"{row}\n")
+            else:
+                ftrain.write(f"{row}\n")
+
+            count += 1
+            duration_seconds = count * args.duration
+            duration_hours = duration_seconds / 3600
+
+            process_bar.update(n=1)
+            process_bar.set_postfix({
+                # "duration_seconds": round(duration_seconds, 4),
+                "duration_hours": round(duration_hours, 4),
+
+            })
+
+    return
+
+
+if __name__ == "__main__":
+    main()

examples/frcrn/step_2_train_model.py

@@ -0,0 +1,457 @@
+#!/usr/bin/python3
+# -*- coding: utf-8 -*-
+"""
+https://arxiv.org/abs/2206.07293
+
+FRCRN 论文中:
+在 WSJ0 数据集上训练了 120 个 epoch 得到 pesq 3.62, stoi 98.24, si-snr 21.33
+
+WSJ0 包含约 80小时的纯净英语语音录音.
+
+我的音频大约是 1300 小时, 则预期大约需要 10个 epoch
+"""
+import argparse
+import json
+import logging
+from logging.handlers import TimedRotatingFileHandler
+import os
+import platform
+from pathlib import Path
+import random
+import sys
+import shutil
+from typing import List
+
+pwd = os.path.abspath(os.path.dirname(__file__))
+sys.path.append(os.path.join(pwd, "../../"))
+
+import numpy as np
+import torch
+import torch.nn as nn
+from torch.nn import functional as F
+from torch.utils.data.dataloader import DataLoader
+from tqdm import tqdm
+
+from toolbox.torch.utils.data.dataset.denoise_jsonl_dataset import DenoiseJsonlDataset
+from toolbox.torchaudio.losses.snr import NegativeSISNRLoss
+from toolbox.torchaudio.losses.spectral import LSDLoss, MultiResolutionSTFTLoss
+from toolbox.torchaudio.metrics.pesq import run_pesq_score
+from toolbox.torchaudio.models.frcrn.configuration_frcrn import FRCRNConfig
+from toolbox.torchaudio.models.frcrn.modeling_frcrn import FRCRN, FRCRNPretrainedModel
+
+
+def get_args():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--train_dataset", default="train.jsonl", type=str)
+    parser.add_argument("--valid_dataset", default="valid.jsonl", type=str)
+
+    parser.add_argument("--num_serialized_models_to_keep", default=15, type=int)
+    parser.add_argument("--patience", default=30, type=int)
+    parser.add_argument("--serialization_dir", default="serialization_dir", type=str)
+
+    parser.add_argument("--config_file", default="config.yaml", type=str)
+
+    args = parser.parse_args()
+    return args
+
+
+def logging_config(file_dir: str):
+    fmt = "%(asctime)s - %(name)s - %(levelname)s  %(filename)s:%(lineno)d >  %(message)s"
+
+    logging.basicConfig(format=fmt,
+                        datefmt="%m/%d/%Y %H:%M:%S",
+                        level=logging.INFO)
+    file_handler = TimedRotatingFileHandler(
+        filename=os.path.join(file_dir, "main.log"),
+        encoding="utf-8",
+        when="D",
+        interval=1,
+        backupCount=7
+    )
+    file_handler.setLevel(logging.INFO)
+    file_handler.setFormatter(logging.Formatter(fmt))
+    logger = logging.getLogger(__name__)
+    logger.addHandler(file_handler)
+
+    return logger
+
+
+class CollateFunction(object):
+    def __init__(self):
+        pass
+
+    def __call__(self, batch: List[dict]):
+        clean_audios = list()
+        noisy_audios = list()
+
+        for sample in batch:
+            # 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"]
+
+            clean_audios.append(clean_audio)
+            noisy_audios.append(noisy_audio)
+
+        clean_audios = torch.stack(clean_audios)
+        noisy_audios = torch.stack(noisy_audios)
+
+        # assert
+        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()
+
+
+def main():
+    args = get_args()
+
+    config = FRCRNConfig.from_pretrained(
+        pretrained_model_name_or_path=args.config_file,
+    )
+
+    serialization_dir = Path(args.serialization_dir)
+    serialization_dir.mkdir(parents=True, exist_ok=True)
+
+    logger = logging_config(serialization_dir)
+
+    random.seed(config.seed)
+    np.random.seed(config.seed)
+    torch.manual_seed(config.seed)
+    logger.info(f"set seed: {config.seed}")
+
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    n_gpu = torch.cuda.device_count()
+    logger.info(f"GPU available count: {n_gpu}; device: {device}")
+
+    # datasets
+    train_dataset = DenoiseJsonlDataset(
+        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=225000,
+    )
+    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,
+    )
+    train_data_loader = DataLoader(
+        dataset=train_dataset,
+        batch_size=config.batch_size,
+        # shuffle=True,
+        sampler=None,
+        # Linux 系统中可以使用多个子进程加载数据, 而在 Windows 系统中不能.
+        num_workers=0 if platform.system() == "Windows" else os.cpu_count() // 2,
+        collate_fn=collate_fn,
+        pin_memory=False,
+        prefetch_factor=2,
+    )
+    valid_data_loader = DataLoader(
+        dataset=valid_dataset,
+        batch_size=config.batch_size,
+        # shuffle=True,
+        sampler=None,
+        # Linux 系统中可以使用多个子进程加载数据, 而在 Windows 系统中不能.
+        num_workers=0 if platform.system() == "Windows" else os.cpu_count() // 2,
+        collate_fn=collate_fn,
+        pin_memory=False,
+        prefetch_factor=2,
+    )
+
+    # models
+    logger.info(f"prepare models. config_file: {args.config_file}")
+    model = FRCRNPretrainedModel(config).to(device)
+    model.to(device)
+    model.train()
+
+    # optimizer
+    logger.info("prepare optimizer, lr_scheduler, loss_fn, evaluation_metric")
+    optimizer = torch.optim.AdamW(model.get_params(weight_decay=config.weight_decay), config.lr)
+
+    # resume training
+    last_step_idx = -1
+    last_epoch = -1
+    for step_idx_str in serialization_dir.glob("steps-*"):
+        step_idx_str = Path(step_idx_str)
+        step_idx = step_idx_str.stem.split("-")[1]
+        step_idx = int(step_idx)
+        if step_idx > last_step_idx:
+            last_step_idx = step_idx
+    # last_epoch = 0
+
+    if last_step_idx != -1:
+        logger.info(f"resume from steps-{last_step_idx}.")
+        model_pt = serialization_dir / f"steps-{last_step_idx}/model.pt"
+        # optimizer_pth = serialization_dir / f"steps-{last_step_idx}/optimizer.pth"
+
+        logger.info(f"load state dict for model.")
+        with open(model_pt.as_posix(), "rb") as f:
+            state_dict = torch.load(f, map_location="cpu", weights_only=True)
+        model.load_state_dict(state_dict, strict=True)
+
+        # logger.info(f"load state dict for optimizer.")
+        # with open(optimizer_pth.as_posix(), "rb") as f:
+        #     state_dict = torch.load(f, map_location="cpu", weights_only=True)
+        # optimizer.load_state_dict(state_dict)
+
+    if config.lr_scheduler == "CosineAnnealingLR":
+        lr_scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(
+            optimizer,
+            last_epoch=last_epoch,
+            # T_max=10 * config.eval_steps,
+            # eta_min=0.01 * config.lr,
+            **config.lr_scheduler_kwargs,
+        )
+    elif config.lr_scheduler == "MultiStepLR":
+        lr_scheduler = torch.optim.lr_scheduler.MultiStepLR(
+            optimizer,
+            last_epoch=last_epoch,
+            milestones=[10000, 20000, 30000, 40000, 50000], gamma=0.5
+        )
+    else:
+        raise AssertionError(f"invalid lr_scheduler: {config.lr_scheduler}")
+
+    neg_si_snr_loss_fn = NegativeSISNRLoss(reduction="mean").to(device)
+    mr_stft_loss_fn = MultiResolutionSTFTLoss(
+        fft_size_list=[256, 512, 1024],
+        win_size_list=[256, 512, 1024],
+        hop_size_list=[128, 256, 512],
+        factor_sc=1.5,
+        factor_mag=1.0,
+        reduction="mean"
+    ).to(device)
+
+    # training loop
+
+    # state
+    average_pesq_score = 1000000000
+    average_loss = 1000000000
+    average_neg_si_snr_loss = 1000000000
+    average_mask_loss = 1000000000
+
+    model_list = list()
+    best_epoch_idx = None
+    best_step_idx = None
+    best_metric = None
+    patience_count = 0
+
+    step_idx = 0 if last_step_idx == -1 else last_step_idx
+
+    logger.info("training")
+    early_stop_flag = False
+    for epoch_idx in range(max(0, last_epoch+1), config.max_epochs):
+        if early_stop_flag:
+            break
+
+        # train
+        model.train()
+
+        total_pesq_score = 0.
+        total_loss = 0.
+        total_mr_stft_loss = 0.
+        total_neg_si_snr_loss = 0.
+        total_mask_loss = 0.
+        total_batches = 0.
+
+        progress_bar_train = tqdm(
+            initial=step_idx,
+            desc="Training; epoch-{}".format(epoch_idx),
+        )
+        for train_batch in train_data_loader:
+            clean_audios, noisy_audios = train_batch
+            clean_audios: torch.Tensor = clean_audios.to(device)
+            noisy_audios: torch.Tensor = noisy_audios.to(device)
+
+            est_spec, est_wav, est_mask = model.forward(noisy_audios)
+            denoise_audios = est_wav
+
+            mr_stft_loss = mr_stft_loss_fn.forward(denoise_audios, clean_audios)
+            neg_si_snr_loss = neg_si_snr_loss_fn.forward(denoise_audios, clean_audios)
+            mask_loss = model.mask_loss_fn(est_mask, clean_audios, noisy_audios)
+
+            loss = 1.0 * mr_stft_loss + 1.0 * neg_si_snr_loss + 1.0 * mask_loss
+            if torch.any(torch.isnan(loss)) or torch.any(torch.isinf(loss)):
+                logger.info(f"find nan or inf in loss.")
+                continue
+
+            denoise_audios_list_r = list(denoise_audios.detach().cpu().numpy())
+            clean_audios_list_r = list(clean_audios.detach().cpu().numpy())
+            pesq_score = run_pesq_score(clean_audios_list_r, denoise_audios_list_r, sample_rate=config.sample_rate, mode="nb")
+
+            optimizer.zero_grad()
+            loss.backward()
+            torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=config.clip_grad_norm)
+            optimizer.step()
+            lr_scheduler.step()
+
+            total_pesq_score += pesq_score
+            total_loss += loss.item()
+            total_mr_stft_loss += mr_stft_loss.item()
+            total_neg_si_snr_loss += neg_si_snr_loss.item()
+            total_mask_loss += mask_loss.item()
+            total_batches += 1
+
+            average_pesq_score = round(total_pesq_score / total_batches, 4)
+            average_loss = round(total_loss / total_batches, 4)
+            average_mr_stft_loss = round(total_mr_stft_loss / total_batches, 4)
+            average_neg_si_snr_loss = round(total_neg_si_snr_loss / total_batches, 4)
+            average_mask_loss = round(total_mask_loss / total_batches, 4)
+
+            progress_bar_train.update(1)
+            progress_bar_train.set_postfix({
+                "lr": lr_scheduler.get_last_lr()[0],
+                "pesq_score": average_pesq_score,
+                "loss": average_loss,
+                "mr_stft_loss": average_mr_stft_loss,
+                "neg_si_snr_loss": average_neg_si_snr_loss,
+                "mask_loss": average_mask_loss,
+            })
+
+            # evaluation
+            step_idx += 1
+            if step_idx % config.eval_steps == 0:
+                model.eval()
+                with torch.no_grad():
+                    torch.cuda.empty_cache()
+
+                    total_pesq_score = 0.
+                    total_loss = 0.
+                    total_mr_stft_loss = 0.
+                    total_neg_si_snr_loss = 0.
+                    total_mask_loss = 0.
+                    total_batches = 0.
+
+                    progress_bar_train.close()
+                    progress_bar_eval = tqdm(
+                        desc="Evaluation; steps-{}k".format(int(step_idx/1000)),
+                    )
+                    for eval_batch in valid_data_loader:
+                        clean_audios, noisy_audios = eval_batch
+                        clean_audios = clean_audios.to(device)
+                        noisy_audios = noisy_audios.to(device)
+
+                        est_spec, est_wav, est_mask = model.forward(noisy_audios)
+                        denoise_audios = est_wav
+
+                        mr_stft_loss = mr_stft_loss_fn.forward(denoise_audios, clean_audios)
+                        neg_si_snr_loss = neg_si_snr_loss_fn.forward(denoise_audios, clean_audios)
+                        mask_loss = model.mask_loss_fn(est_mask, clean_audios, noisy_audios)
+
+                        loss = 1.0 * mr_stft_loss + 1.0 * neg_si_snr_loss + 1.0 * mask_loss
+                        if torch.any(torch.isnan(loss)) or torch.any(torch.isinf(loss)):
+                            logger.info(f"find nan or inf in loss.")
+                            continue
+
+                        denoise_audios_list_r = list(denoise_audios.detach().cpu().numpy())
+                        clean_audios_list_r = list(clean_audios.detach().cpu().numpy())
+                        pesq_score = run_pesq_score(clean_audios_list_r, denoise_audios_list_r, sample_rate=config.sample_rate, mode="nb")
+
+                        total_pesq_score += pesq_score
+                        total_loss += loss.item()
+                        total_neg_si_snr_loss += neg_si_snr_loss.item()
+                        total_mask_loss += mask_loss.item()
+                        total_batches += 1
+
+                        average_pesq_score = round(total_pesq_score / total_batches, 4)
+                        average_loss = round(total_loss / total_batches, 4)
+                        average_mr_stft_loss = round(total_mr_stft_loss / total_batches, 4)
+                        average_neg_si_snr_loss = round(total_neg_si_snr_loss / total_batches, 4)
+                        average_mask_loss = round(total_mask_loss / total_batches, 4)
+
+                        progress_bar_eval.update(1)
+                        progress_bar_eval.set_postfix({
+                            "lr": lr_scheduler.get_last_lr()[0],
+                            "pesq_score": average_pesq_score,
+                            "loss": average_loss,
+                            "mr_stft_loss": average_mr_stft_loss,
+                            "neg_si_snr_loss": average_neg_si_snr_loss,
+                            "mask_loss": average_mask_loss,
+                        })
+
+                    total_pesq_score = 0.
+                    total_loss = 0.
+                    total_mr_stft_loss = 0.
+                    total_neg_si_snr_loss = 0.
+                    total_mask_loss = 0.
+                    total_batches = 0.
+
+                    progress_bar_eval.close()
+                    progress_bar_train = tqdm(
+                        initial=progress_bar_train.n,
+                        postfix=progress_bar_train.postfix,
+                        desc=progress_bar_train.desc,
+                    )
+
+                    # save path
+                    save_dir = serialization_dir / "steps-{}".format(step_idx)
+                    save_dir.mkdir(parents=True, exist_ok=False)
+
+                    # save models
+                    model.save_pretrained(save_dir.as_posix())
+
+                    model_list.append(save_dir)
+                    if len(model_list) >= args.num_serialized_models_to_keep:
+                        model_to_delete: Path = model_list.pop(0)
+                        shutil.rmtree(model_to_delete.as_posix())
+
+                    # save metric
+                    if best_metric is None:
+                        best_epoch_idx = epoch_idx
+                        best_step_idx = step_idx
+                        best_metric = average_pesq_score
+                    elif average_pesq_score >= best_metric:
+                        # great is better.
+                        best_epoch_idx = epoch_idx
+                        best_step_idx = step_idx
+                        best_metric = average_pesq_score
+                    else:
+                        pass
+
+                    metrics = {
+                        "epoch_idx": epoch_idx,
+                        "best_epoch_idx": best_epoch_idx,
+                        "best_step_idx": best_step_idx,
+                        "pesq_score": average_pesq_score,
+                        "loss": average_loss,
+                        "neg_si_snr_loss": average_neg_si_snr_loss,
+                        "mask_loss": average_mask_loss,
+                    }
+                    metrics_filename = save_dir / "metrics_epoch.json"
+                    with open(metrics_filename, "w", encoding="utf-8") as f:
+                        json.dump(metrics, f, indent=4, ensure_ascii=False)
+
+                    # save best
+                    best_dir = serialization_dir / "best"
+                    if best_epoch_idx == epoch_idx and best_step_idx == step_idx:
+                        if best_dir.exists():
+                            shutil.rmtree(best_dir)
+                        shutil.copytree(save_dir, best_dir)
+
+                    # early stop
+                    early_stop_flag = False
+                    if best_epoch_idx == epoch_idx and best_step_idx == step_idx:
+                        patience_count = 0
+                    else:
+                        patience_count += 1
+                    if patience_count >= args.patience:
+                        early_stop_flag = True
+
+                    # early stop
+                    if early_stop_flag:
+                        break
+                model.train()
+
+    return
+
+
+if __name__ == "__main__":
+    main()

examples/frcrn/yaml/config-10.yaml

@@ -0,0 +1,31 @@
+model_name: "frcrn"
+
+sample_rate: 8000
+segment_size: 32000
+nfft: 128
+win_size: 128
+hop_size: 64
+win_type: hann
+
+use_complex_networks: true
+model_depth: 10
+model_complexity: -1
+
+min_snr_db: -10
+max_snr_db: 20
+
+num_workers: 8
+batch_size: 32
+eval_steps: 20000
+
+lr: 0.001
+lr_scheduler: "CosineAnnealingLR"
+lr_scheduler_kwargs:
+  T_max: 250000
+  eta_min: 0.0001
+
+max_epochs: 100
+weight_decay: 1.0e-05
+clip_grad_norm: 10.0
+seed: 1234
+num_gpus: -1

examples/frcrn/yaml/config-14.yaml

@@ -0,0 +1,31 @@
+model_name: "frcrn"
+
+sample_rate: 8000
+segment_size: 32000
+nfft: 640
+win_size: 640
+hop_size: 320
+win_type: hann
+
+use_complex_networks: true
+model_depth: 14
+model_complexity: -1
+
+min_snr_db: -10
+max_snr_db: 20
+
+num_workers: 8
+batch_size: 32
+eval_steps: 10000
+
+lr: 0.001
+lr_scheduler: "CosineAnnealingLR"
+lr_scheduler_kwargs:
+  T_max: 250000
+  eta_min: 0.0001
+
+max_epochs: 100
+weight_decay: 1.0e-05
+clip_grad_norm: 10.0
+seed: 1234
+num_gpus: -1

examples/frcrn/yaml/config-20.yaml

@@ -0,0 +1,31 @@
+model_name: "frcrn"
+
+sample_rate: 8000
+segment_size: 32000
+nfft: 512
+win_size: 512
+hop_size: 256
+win_type: hann
+
+use_complex_networks: true
+model_depth: 20
+model_complexity: 45
+
+min_snr_db: -10
+max_snr_db: 20
+
+num_workers: 8
+batch_size: 32
+eval_steps: 10000
+
+lr: 0.001
+lr_scheduler: "CosineAnnealingLR"
+lr_scheduler_kwargs:
+  T_max: 250000
+  eta_min: 0.0001
+
+max_epochs: 100
+weight_decay: 1.0e-05
+clip_grad_norm: 10.0
+seed: 1234
+num_gpus: -1

examples/frcrn_mp3_to_wav/run.sh

@@ -0,0 +1,156 @@
+#!/usr/bin/env bash
+
+: <<'END'
+
+
+sh run.sh --stage 1 --stop_stage 2 --system_version centos --file_folder_name file_dir --final_model_name frcrn-20-512-nx-dns3 \
+--config_file "yaml/config-10.yaml" \
+--noise_dir "/data/tianxing/HuggingDatasets/nx_noise/data/noise" \
+--speech_dir "/data/tianxing/HuggingDatasets/nx_noise/data/speech"
+
+
+sh run.sh --stage 2 --stop_stage 2 --system_version centos --file_folder_name file_dir --final_model_name frcrn-10-nx2 \
+--config_file "yaml/config-10.yaml" \
+--audio_dir "/data/tianxing/HuggingDatasets/nx_noise/data" \
+
+END
+
+
+# params
+system_version="windows";
+verbose=true;
+stage=0 # start from 0 if you need to start from data preparation
+stop_stage=9
+
+work_dir="$(pwd)"
+file_folder_name=file_folder_name
+final_model_name=final_model_name
+config_file="yaml/config.yaml"
+limit=10
+
+audio_dir=/data/tianxing/HuggingDatasets/nx_noise/data
+
+max_count=10000000
+
+nohup_name=nohup.out
+
+# model params
+batch_size=64
+max_epochs=200
+save_top_k=10
+patience=5
+
+
+# parse options
+while true; do
+  [ -z "${1:-}" ] && break;  # break if there are no arguments
+  case "$1" in
+    --*) name=$(echo "$1" | sed s/^--// | sed s/-/_/g);
+      eval '[ -z "${'"$name"'+xxx}" ]' && echo "$0: invalid option $1" 1>&2 && exit 1;
+      old_value="(eval echo \\$$name)";
+      if [ "${old_value}" == "true" ] || [ "${old_value}" == "false" ]; then
+        was_bool=true;
+      else
+        was_bool=false;
+      fi
+
+      # Set the variable to the right value-- the escaped quotes make it work if
+      # the option had spaces, like --cmd "queue.pl -sync y"
+      eval "${name}=\"$2\"";
+
+      # Check that Boolean-valued arguments are really Boolean.
+      if $was_bool && [[ "$2" != "true" && "$2" != "false" ]]; then
+        echo "$0: expected \"true\" or \"false\": $1 $2" 1>&2
+        exit 1;
+      fi
+      shift 2;
+      ;;
+
+    *) break;
+  esac
+done
+
+file_dir="${work_dir}/${file_folder_name}"
+final_model_dir="${work_dir}/../../trained_models/${final_model_name}";
+evaluation_audio_dir="${file_dir}/evaluation_audio"
+
+train_dataset="${file_dir}/train.jsonl"
+valid_dataset="${file_dir}/valid.jsonl"
+
+$verbose && echo "system_version: ${system_version}"
+$verbose && echo "file_folder_name: ${file_folder_name}"
+
+if [ $system_version == "windows" ]; then
+  alias python3='D:/Users/tianx/PycharmProjects/virtualenv/nx_denoise/Scripts/python.exe'
+elif [ $system_version == "centos" ] || [ $system_version == "ubuntu" ]; then
+  #source /data/local/bin/nx_denoise/bin/activate
+  alias python3='/data/local/bin/nx_denoise/bin/python3'
+fi
+
+
+if [ ${stage} -le 1 ] && [ ${stop_stage} -ge 1 ]; then
+  $verbose && echo "stage 1: prepare data"
+  cd "${work_dir}" || exit 1
+  python3 step_1_prepare_data.py \
+  --file_dir "${file_dir}" \
+  --audio_dir "${audio_dir}" \
+  --train_dataset "${train_dataset}" \
+  --valid_dataset "${valid_dataset}" \
+  --max_count "${max_count}" \
+
+fi
+
+
+if [ ${stage} -le 2 ] && [ ${stop_stage} -ge 2 ]; then
+  $verbose && echo "stage 2: train model"
+  cd "${work_dir}" || exit 1
+  python3 step_2_train_model.py \
+  --train_dataset "${train_dataset}" \
+  --valid_dataset "${valid_dataset}" \
+  --serialization_dir "${file_dir}" \
+  --config_file "${config_file}" \
+
+fi
+
+
+if [ ${stage} -le 3 ] && [ ${stop_stage} -ge 3 ]; then
+  $verbose && echo "stage 3: test model"
+  cd "${work_dir}" || exit 1
+  python3 step_3_evaluation.py \
+  --valid_dataset "${valid_dataset}" \
+  --model_dir "${file_dir}/best" \
+  --evaluation_audio_dir "${evaluation_audio_dir}" \
+  --limit "${limit}" \
+
+fi
+
+
+if [ ${stage} -le 4 ] && [ ${stop_stage} -ge 4 ]; then
+  $verbose && echo "stage 4: collect files"
+  cd "${work_dir}" || exit 1
+
+  mkdir -p ${final_model_dir}
+
+  cp "${file_dir}/best"/* "${final_model_dir}"
+  cp -r "${file_dir}/evaluation_audio" "${final_model_dir}"
+
+  cd "${final_model_dir}/.." || exit 1;
+
+  if [ -e "${final_model_name}.zip" ]; then
+    rm -rf "${final_model_name}_backup.zip"
+    mv "${final_model_name}.zip" "${final_model_name}_backup.zip"
+  fi
+
+  zip -r "${final_model_name}.zip" "${final_model_name}"
+  rm -rf "${final_model_name}"
+
+fi
+
+
+if [ ${stage} -le 5 ] && [ ${stop_stage} -ge 5 ]; then
+  $verbose && echo "stage 5: clear file_dir"
+  cd "${work_dir}" || exit 1
+
+  rm -rf "${file_dir}";
+
+fi

examples/frcrn_mp3_to_wav/step_1_prepare_data.py

@@ -0,0 +1,127 @@
+#!/usr/bin/python3
+# -*- coding: utf-8 -*-
+import argparse
+import json
+import os
+from pathlib import Path
+import random
+import sys
+
+pwd = os.path.abspath(os.path.dirname(__file__))
+sys.path.append(os.path.join(pwd, "../../"))
+
+import librosa
+import numpy as np
+from tqdm import tqdm
+
+
+def get_args():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--file_dir", default="./", type=str)
+
+    parser.add_argument(
+        "--audio_dir",
+        default="E:/Users/tianx/HuggingDatasets/nx_noise/data/speech",
+        type=str
+    )
+
+    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("--target_sample_rate", default=8000, type=int)
+
+    parser.add_argument("--max_count", default=-1, type=int)
+
+    args = parser.parse_args()
+    return args
+
+
+def target_second_signal_generator(data_dir: str, duration: int = 2, sample_rate: int = 8000, max_epoch: int = 1):
+    data_dir = Path(data_dir)
+    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):
+                if np.sum(signal[begin: begin+win_size]) == 0:
+                    continue
+                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 main():
+    args = get_args()
+
+    file_dir = Path(args.file_dir)
+    file_dir.mkdir(exist_ok=True)
+
+    audio_dir = Path(args.audio_dir)
+
+    audio_generator = target_second_signal_generator(
+        audio_dir.as_posix(),
+        duration=args.duration,
+        sample_rate=args.target_sample_rate,
+        max_epoch=1,
+    )
+    count = 0
+    process_bar = tqdm(desc="build dataset jsonl")
+    with open(args.train_dataset, "w", encoding="utf-8") as ftrain, open(args.valid_dataset, "w", encoding="utf-8") as fvalid:
+        for audio in audio_generator:
+            if count >= args.max_count > 0:
+                break
+
+            filename = audio["filename"]
+            raw_duration = audio["raw_duration"]
+            offset = audio["offset"]
+            duration = audio["duration"]
+
+            random1 = random.random()
+            random2 = random.random()
+
+            row = {
+                "count": count,
+
+                "filename": filename,
+                "raw_duration": raw_duration,
+                "offset": offset,
+                "duration": duration,
+
+                "random1": random1,
+            }
+            row = json.dumps(row, ensure_ascii=False)
+            if random2 < (1 / 10):
+                fvalid.write(f"{row}\n")
+            else:
+                ftrain.write(f"{row}\n")
+
+            count += 1
+            duration_seconds = count * args.duration
+            duration_hours = duration_seconds / 3600
+
+            process_bar.update(n=1)
+            process_bar.set_postfix({
+                "duration_hours": round(duration_hours, 4),
+            })
+
+    return
+
+
+if __name__ == "__main__":
+    main()

examples/frcrn_mp3_to_wav/step_2_train_model.py

@@ -0,0 +1,442 @@
+#!/usr/bin/python3
+# -*- coding: utf-8 -*-
+import argparse
+import json
+import logging
+from logging.handlers import TimedRotatingFileHandler
+import os
+import platform
+from pathlib import Path
+import random
+import sys
+import shutil
+from typing import List
+
+pwd = os.path.abspath(os.path.dirname(__file__))
+sys.path.append(os.path.join(pwd, "../../"))
+
+import numpy as np
+import torch
+import torch.nn as nn
+from torch.nn import functional as F
+from torch.utils.data.dataloader import DataLoader
+from tqdm import tqdm
+
+from toolbox.torch.utils.data.dataset.mp3_to_wav_jsonl_dataset import Mp3ToWavJsonlDataset
+from toolbox.torchaudio.losses.snr import NegativeSISNRLoss
+from toolbox.torchaudio.losses.spectral import LSDLoss, MultiResolutionSTFTLoss
+from toolbox.torchaudio.metrics.pesq import run_pesq_score
+from toolbox.torchaudio.models.frcrn.configuration_frcrn import FRCRNConfig
+from toolbox.torchaudio.models.frcrn.modeling_frcrn import FRCRN, FRCRNPretrainedModel
+
+
+def get_args():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--train_dataset", default="train.jsonl", type=str)
+    parser.add_argument("--valid_dataset", default="valid.jsonl", type=str)
+
+    parser.add_argument("--num_serialized_models_to_keep", default=15, type=int)
+    parser.add_argument("--patience", default=30, type=int)
+    parser.add_argument("--serialization_dir", default="serialization_dir", type=str)
+
+    parser.add_argument("--config_file", default="config.yaml", type=str)
+
+    args = parser.parse_args()
+    return args
+
+
+def logging_config(file_dir: str):
+    fmt = "%(asctime)s - %(name)s - %(levelname)s  %(filename)s:%(lineno)d >  %(message)s"
+
+    logging.basicConfig(format=fmt,
+                        datefmt="%m/%d/%Y %H:%M:%S",
+                        level=logging.INFO)
+    file_handler = TimedRotatingFileHandler(
+        filename=os.path.join(file_dir, "main.log"),
+        encoding="utf-8",
+        when="D",
+        interval=1,
+        backupCount=7
+    )
+    file_handler.setLevel(logging.INFO)
+    file_handler.setFormatter(logging.Formatter(fmt))
+    logger = logging.getLogger(__name__)
+    logger.addHandler(file_handler)
+
+    return logger
+
+
+class CollateFunction(object):
+    def __init__(self):
+        pass
+
+    def __call__(self, batch: List[dict]):
+        mp3_waveform_list = list()
+        wav_waveform_list = list()
+
+        for sample in batch:
+            mp3_waveform: torch.Tensor = sample["mp3_waveform"]
+            wav_waveform: torch.Tensor = sample["wav_waveform"]
+
+            mp3_waveform_list.append(mp3_waveform)
+            wav_waveform_list.append(wav_waveform)
+
+        mp3_waveform_list = torch.stack(mp3_waveform_list)
+        wav_waveform_list = torch.stack(wav_waveform_list)
+
+        # assert
+        if torch.any(torch.isnan(mp3_waveform_list)) or torch.any(torch.isinf(mp3_waveform_list)):
+            raise AssertionError("nan or inf in mp3_waveform_list")
+        if torch.any(torch.isnan(wav_waveform_list)) or torch.any(torch.isinf(wav_waveform_list)):
+            raise AssertionError("nan or inf in wav_waveform_list")
+
+        return mp3_waveform_list, wav_waveform_list
+
+
+collate_fn = CollateFunction()
+
+
+def main():
+    args = get_args()
+
+    config = FRCRNConfig.from_pretrained(
+        pretrained_model_name_or_path=args.config_file,
+    )
+
+    serialization_dir = Path(args.serialization_dir)
+    serialization_dir.mkdir(parents=True, exist_ok=True)
+
+    logger = logging_config(serialization_dir)
+
+    random.seed(config.seed)
+    np.random.seed(config.seed)
+    torch.manual_seed(config.seed)
+    logger.info(f"set seed: {config.seed}")
+
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    n_gpu = torch.cuda.device_count()
+    logger.info(f"GPU available count: {n_gpu}; device: {device}")
+
+    # datasets
+    train_dataset = Mp3ToWavJsonlDataset(
+        jsonl_file=args.train_dataset,
+        expected_sample_rate=config.sample_rate,
+        max_wave_value=32768.0,
+        # skip=225000,
+    )
+    valid_dataset = Mp3ToWavJsonlDataset(
+        jsonl_file=args.valid_dataset,
+        expected_sample_rate=config.sample_rate,
+        max_wave_value=32768.0,
+    )
+    train_data_loader = DataLoader(
+        dataset=train_dataset,
+        batch_size=config.batch_size,
+        # shuffle=True,
+        sampler=None,
+        # Linux 系统中可以使用多个子进程加载数据, 而在 Windows 系统中不能.
+        num_workers=0 if platform.system() == "Windows" else os.cpu_count() // 2,
+        collate_fn=collate_fn,
+        pin_memory=False,
+        prefetch_factor=2,
+    )
+    valid_data_loader = DataLoader(
+        dataset=valid_dataset,
+        batch_size=config.batch_size,
+        # shuffle=True,
+        sampler=None,
+        # Linux 系统中可以使用多个子进程加载数据, 而在 Windows 系统中不能.
+        num_workers=0 if platform.system() == "Windows" else os.cpu_count() // 2,
+        collate_fn=collate_fn,
+        pin_memory=False,
+        prefetch_factor=2,
+    )
+
+    # models
+    logger.info(f"prepare models. config_file: {args.config_file}")
+    model = FRCRNPretrainedModel(config).to(device)
+    model.to(device)
+    model.train()
+
+    # optimizer
+    logger.info("prepare optimizer, lr_scheduler, loss_fn, evaluation_metric")
+    optimizer = torch.optim.AdamW(model.get_params(weight_decay=config.weight_decay), config.lr)
+
+    # resume training
+    last_step_idx = -1
+    last_epoch = -1
+    for step_idx_str in serialization_dir.glob("steps-*"):
+        step_idx_str = Path(step_idx_str)
+        step_idx = step_idx_str.stem.split("-")[1]
+        step_idx = int(step_idx)
+        if step_idx > last_step_idx:
+            last_step_idx = step_idx
+    # last_epoch = 0
+
+    if last_step_idx != -1:
+        logger.info(f"resume from steps-{last_step_idx}.")
+        model_pt = serialization_dir / f"steps-{last_step_idx}/model.pt"
+        # optimizer_pth = serialization_dir / f"steps-{last_step_idx}/optimizer.pth"
+
+        logger.info(f"load state dict for model.")
+        with open(model_pt.as_posix(), "rb") as f:
+            state_dict = torch.load(f, map_location="cpu", weights_only=True)
+        model.load_state_dict(state_dict, strict=True)
+
+        # logger.info(f"load state dict for optimizer.")
+        # with open(optimizer_pth.as_posix(), "rb") as f:
+        #     state_dict = torch.load(f, map_location="cpu", weights_only=True)
+        # optimizer.load_state_dict(state_dict)
+
+    if config.lr_scheduler == "CosineAnnealingLR":
+        lr_scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(
+            optimizer,
+            last_epoch=last_epoch,
+            # T_max=10 * config.eval_steps,
+            # eta_min=0.01 * config.lr,
+            **config.lr_scheduler_kwargs,
+        )
+    elif config.lr_scheduler == "MultiStepLR":
+        lr_scheduler = torch.optim.lr_scheduler.MultiStepLR(
+            optimizer,
+            last_epoch=last_epoch,
+            milestones=[10000, 20000, 30000, 40000, 50000], gamma=0.5
+        )
+    else:
+        raise AssertionError(f"invalid lr_scheduler: {config.lr_scheduler}")
+
+    neg_si_snr_loss_fn = NegativeSISNRLoss(reduction="mean").to(device)
+    mr_stft_loss_fn = MultiResolutionSTFTLoss(
+        fft_size_list=[256, 512, 1024],
+        win_size_list=[256, 512, 1024],
+        hop_size_list=[128, 256, 512],
+        factor_sc=1.5,
+        factor_mag=1.0,
+        reduction="mean"
+    ).to(device)
+
+    # training loop
+
+    # state
+    average_pesq_score = 1000000000
+    average_loss = 1000000000
+    average_neg_si_snr_loss = 1000000000
+    average_mask_loss = 1000000000
+
+    model_list = list()
+    best_epoch_idx = None
+    best_step_idx = None
+    best_metric = None
+    patience_count = 0
+
+    step_idx = 0 if last_step_idx == -1 else last_step_idx
+
+    logger.info("training")
+    early_stop_flag = False
+    for epoch_idx in range(max(0, last_epoch+1), config.max_epochs):
+        if early_stop_flag:
+            break
+
+        # train
+        model.train()
+
+        total_pesq_score = 0.
+        total_loss = 0.
+        total_mr_stft_loss = 0.
+        total_neg_si_snr_loss = 0.
+        total_mask_loss = 0.
+        total_batches = 0.
+
+        progress_bar_train = tqdm(
+            initial=step_idx,
+            desc="Training; epoch-{}".format(epoch_idx),
+        )
+        for train_batch in train_data_loader:
+            mp3_audios, wav_audios = train_batch
+            noisy_audios: torch.Tensor = mp3_audios.to(device)
+            clean_audios: torch.Tensor = wav_audios.to(device)
+
+            est_spec, est_wav, est_mask = model.forward(noisy_audios)
+            denoise_audios = est_wav
+
+            mr_stft_loss = mr_stft_loss_fn.forward(denoise_audios, clean_audios)
+            neg_si_snr_loss = neg_si_snr_loss_fn.forward(denoise_audios, clean_audios)
+            mask_loss = model.mask_loss_fn(est_mask, clean_audios, noisy_audios)
+
+            loss = 1.0 * mr_stft_loss + 1.0 * neg_si_snr_loss + 1.0 * mask_loss
+            if torch.any(torch.isnan(loss)) or torch.any(torch.isinf(loss)):
+                logger.info(f"find nan or inf in loss.")
+                continue
+
+            denoise_audios_list_r = list(denoise_audios.detach().cpu().numpy())
+            clean_audios_list_r = list(clean_audios.detach().cpu().numpy())
+            pesq_score = run_pesq_score(clean_audios_list_r, denoise_audios_list_r, sample_rate=config.sample_rate, mode="nb")
+
+            optimizer.zero_grad()
+            loss.backward()
+            torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=config.clip_grad_norm)
+            optimizer.step()
+            lr_scheduler.step()
+
+            total_pesq_score += pesq_score
+            total_loss += loss.item()
+            total_mr_stft_loss += mr_stft_loss.item()
+            total_neg_si_snr_loss += neg_si_snr_loss.item()
+            total_mask_loss += mask_loss.item()
+            total_batches += 1
+
+            average_pesq_score = round(total_pesq_score / total_batches, 4)
+            average_loss = round(total_loss / total_batches, 4)
+            average_mr_stft_loss = round(total_mr_stft_loss / total_batches, 4)
+            average_neg_si_snr_loss = round(total_neg_si_snr_loss / total_batches, 4)
+            average_mask_loss = round(total_mask_loss / total_batches, 4)
+
+            progress_bar_train.update(1)
+            progress_bar_train.set_postfix({
+                "lr": lr_scheduler.get_last_lr()[0],
+                "pesq_score": average_pesq_score,
+                "loss": average_loss,
+                "mr_stft_loss": average_mr_stft_loss,
+                "neg_si_snr_loss": average_neg_si_snr_loss,
+                "mask_loss": average_mask_loss,
+            })
+
+            # evaluation
+            step_idx += 1
+            if step_idx % config.eval_steps == 0:
+                model.eval()
+                with torch.no_grad():
+                    torch.cuda.empty_cache()
+
+                    total_pesq_score = 0.
+                    total_loss = 0.
+                    total_mr_stft_loss = 0.
+                    total_neg_si_snr_loss = 0.
+                    total_mask_loss = 0.
+                    total_batches = 0.
+
+                    progress_bar_train.close()
+                    progress_bar_eval = tqdm(
+                        desc="Evaluation; steps-{}k".format(int(step_idx/1000)),
+                    )
+                    for eval_batch in valid_data_loader:
+                        mp3_audios, wav_audios = eval_batch
+                        noisy_audios: torch.Tensor = mp3_audios.to(device)
+                        clean_audios: torch.Tensor = wav_audios.to(device)
+
+                        est_spec, est_wav, est_mask = model.forward(noisy_audios)
+                        denoise_audios = est_wav
+
+                        mr_stft_loss = mr_stft_loss_fn.forward(denoise_audios, clean_audios)
+                        neg_si_snr_loss = neg_si_snr_loss_fn.forward(denoise_audios, clean_audios)
+                        mask_loss = model.mask_loss_fn(est_mask, clean_audios, noisy_audios)
+
+                        loss = 1.0 * mr_stft_loss + 1.0 * neg_si_snr_loss + 1.0 * mask_loss
+                        if torch.any(torch.isnan(loss)) or torch.any(torch.isinf(loss)):
+                            logger.info(f"find nan or inf in loss.")
+                            continue
+
+                        denoise_audios_list_r = list(denoise_audios.detach().cpu().numpy())
+                        clean_audios_list_r = list(clean_audios.detach().cpu().numpy())
+                        pesq_score = run_pesq_score(clean_audios_list_r, denoise_audios_list_r, sample_rate=config.sample_rate, mode="nb")
+
+                        total_pesq_score += pesq_score
+                        total_loss += loss.item()
+                        total_neg_si_snr_loss += neg_si_snr_loss.item()
+                        total_mask_loss += mask_loss.item()
+                        total_batches += 1
+
+                        average_pesq_score = round(total_pesq_score / total_batches, 4)
+                        average_loss = round(total_loss / total_batches, 4)
+                        average_mr_stft_loss = round(total_mr_stft_loss / total_batches, 4)
+                        average_neg_si_snr_loss = round(total_neg_si_snr_loss / total_batches, 4)
+                        average_mask_loss = round(total_mask_loss / total_batches, 4)
+
+                        progress_bar_eval.update(1)
+                        progress_bar_eval.set_postfix({
+                            "lr": lr_scheduler.get_last_lr()[0],
+                            "pesq_score": average_pesq_score,
+                            "loss": average_loss,
+                            "mr_stft_loss": average_mr_stft_loss,
+                            "neg_si_snr_loss": average_neg_si_snr_loss,
+                            "mask_loss": average_mask_loss,
+                        })
+
+                    total_pesq_score = 0.
+                    total_loss = 0.
+                    total_mr_stft_loss = 0.
+                    total_neg_si_snr_loss = 0.
+                    total_mask_loss = 0.
+                    total_batches = 0.
+
+                    progress_bar_eval.close()
+                    progress_bar_train = tqdm(
+                        initial=progress_bar_train.n,
+                        postfix=progress_bar_train.postfix,
+                        desc=progress_bar_train.desc,
+                    )
+
+                    # save path
+                    save_dir = serialization_dir / "steps-{}".format(step_idx)
+                    save_dir.mkdir(parents=True, exist_ok=False)
+
+                    # save models
+                    model.save_pretrained(save_dir.as_posix())
+
+                    model_list.append(save_dir)
+                    if len(model_list) >= args.num_serialized_models_to_keep:
+                        model_to_delete: Path = model_list.pop(0)
+                        shutil.rmtree(model_to_delete.as_posix())
+
+                    # save metric
+                    if best_metric is None:
+                        best_epoch_idx = epoch_idx
+                        best_step_idx = step_idx
+                        best_metric = average_pesq_score
+                    elif average_pesq_score >= best_metric:
+                        # great is better.
+                        best_epoch_idx = epoch_idx
+                        best_step_idx = step_idx
+                        best_metric = average_pesq_score
+                    else:
+                        pass
+
+                    metrics = {
+                        "epoch_idx": epoch_idx,
+                        "best_epoch_idx": best_epoch_idx,
+                        "best_step_idx": best_step_idx,
+                        "pesq_score": average_pesq_score,
+                        "loss": average_loss,
+                        "neg_si_snr_loss": average_neg_si_snr_loss,
+                        "mask_loss": average_mask_loss,
+                    }
+                    metrics_filename = save_dir / "metrics_epoch.json"
+                    with open(metrics_filename, "w", encoding="utf-8") as f:
+                        json.dump(metrics, f, indent=4, ensure_ascii=False)
+
+                    # save best
+                    best_dir = serialization_dir / "best"
+                    if best_epoch_idx == epoch_idx and best_step_idx == step_idx:
+                        if best_dir.exists():
+                            shutil.rmtree(best_dir)
+                        shutil.copytree(save_dir, best_dir)
+
+                    # early stop
+                    early_stop_flag = False
+                    if best_epoch_idx == epoch_idx and best_step_idx == step_idx:
+                        patience_count = 0
+                    else:
+                        patience_count += 1
+                    if patience_count >= args.patience:
+                        early_stop_flag = True
+
+                    # early stop
+                    if early_stop_flag:
+                        break
+                model.train()
+
+    return
+
+
+if __name__ == "__main__":
+    main()