update

examples/clean_unet_aishell/run.sh

@@ -0,0 +1,178 @@
+#!/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 1 --stop_stage 2 --system_version centos --file_folder_name file_dir \
+--noise_dir "/data/tianxing/HuggingDatasets/nx_noise/data/noise" \
+--speech_dir "/data/tianxing/HuggingDatasets/aishell/data_aishell/wav/train"
+
+
+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
+
+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}" \
+
+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_aishell/step_1_prepare_data.py

@@ -0,0 +1,202 @@
+#!/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("--scale", default=1, type=float)
+
+    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):
+        flag = random.random()
+        if flag > args.scale:
+            continue
+
+        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_aishell/step_2_train_model.py

@@ -0,0 +1,429 @@
+#!/usr/bin/python3
+# -*- coding: utf-8 -*-
+"""
+https://github.com/NVIDIA/CleanUNet/blob/main/train.py
+"""
+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
+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
+
+
+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("--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("--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(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 = 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=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=64,
+    )
+    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=64,
+    )
+
+    # models
+    logger.info(f"prepare models. config_file: {args.config_file}")
+    model = CleanUNetPretrainedModel(config).to(device)
+
+    # optimizer
+    logger.info("prepare optimizer, lr_scheduler")
+    optim_g = torch.optim.AdamW(model.parameters(), config.learning_rate, betas=[config.adam_b1, config.adam_b2])
+
+    # 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}.")
+        generator_pt = serialization_dir / f"epoch-{last_epoch}/generator.pt"
+        discriminator_pt = serialization_dir / f"epoch-{last_epoch}/discriminator.pt"
+        optim_g_pth = serialization_dir / f"epoch-{last_epoch}/optim_g.pth"
+        optim_d_pth = serialization_dir / f"epoch-{last_epoch}/optim_d.pth"
+
+        logger.info(f"load state dict for generator.")
+        with open(generator_pt.as_posix(), "rb") as f:
+            state_dict = torch.load(f, map_location="cpu", weights_only=True)
+        generator.load_state_dict(state_dict, strict=True)
+        logger.info(f"load state dict for discriminator.")
+        with open(discriminator_pt.as_posix(), "rb") as f:
+            state_dict = torch.load(f, map_location="cpu", weights_only=True)
+        discriminator.load_state_dict(state_dict, strict=True)
+
+        logger.info(f"load state dict for optim_g.")
+        with open(optim_g_pth.as_posix(), "rb") as f:
+            state_dict = torch.load(f, map_location="cpu", weights_only=True)
+        optim_g.load_state_dict(state_dict)
+        logger.info(f"load state dict for optim_d.")
+        with open(optim_d_pth.as_posix(), "rb") as f:
+            state_dict = torch.load(f, map_location="cpu", weights_only=True)
+        optim_d.load_state_dict(state_dict)
+
+    scheduler_g = torch.optim.lr_scheduler.ExponentialLR(optim_g, gamma=config.lr_decay, last_epoch=last_epoch)
+    scheduler_d = torch.optim.lr_scheduler.ExponentialLR(optim_d, gamma=config.lr_decay, last_epoch=last_epoch)
+
+    # training loop
+
+    # state
+    loss_d = 10000000000
+    loss_g = 10000000000
+    pesq_metric = 10000000000
+    mag_err = 10000000000
+    pha_err = 10000000000
+    com_err = 10000000000
+    stft_err = 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
+        generator.train()
+        discriminator.train()
+
+        total_loss_d = 0.
+        total_loss_g = 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_audio, noisy_audio = batch
+            clean_audio = clean_audio.to(device)
+            noisy_audio = noisy_audio.to(device)
+            one_labels = torch.ones(clean_audio.shape[0]).to(device)
+
+            clean_mag, clean_pha, clean_com = mag_pha_stft(clean_audio, config.n_fft, config.hop_size, config.win_size, config.compress_factor)
+            noisy_mag, noisy_pha, noisy_com = mag_pha_stft(noisy_audio, config.n_fft, config.hop_size, config.win_size, config.compress_factor)
+
+            mag_g, pha_g, com_g = generator.forward(noisy_mag, noisy_pha)
+
+            audio_g = mag_pha_istft(mag_g, pha_g, config.n_fft, config.hop_size, config.win_size, config.compress_factor)
+            mag_g_hat, pha_g_hat, com_g_hat = mag_pha_stft(audio_g, config.n_fft, config.hop_size, config.win_size, config.compress_factor)
+
+            audio_list_r, audio_list_g = list(clean_audio.cpu().numpy()), list(audio_g.detach().cpu().numpy())
+            batch_pesq_score = batch_pesq(audio_list_r, audio_list_g)
+
+            # Discriminator
+            optim_d.zero_grad()
+            metric_r = discriminator.forward(clean_mag, clean_mag)
+            metric_g = discriminator.forward(clean_mag, mag_g_hat.detach())
+            loss_disc_r = F.mse_loss(one_labels, metric_r.flatten())
+
+            if batch_pesq_score is not None:
+                loss_disc_g = F.mse_loss(batch_pesq_score.to(device), metric_g.flatten())
+            else:
+                # print("pesq is None!")
+                loss_disc_g = 0
+
+            loss_disc_all = loss_disc_r + loss_disc_g
+            loss_disc_all.backward()
+            optim_d.step()
+
+            # Generator
+            optim_g.zero_grad()
+            # L2 Magnitude Loss
+            loss_mag = F.mse_loss(clean_mag, mag_g)
+            # Anti-wrapping Phase Loss
+            loss_ip, loss_gd, loss_iaf = phase_losses(clean_pha, pha_g)
+            loss_pha = loss_ip + loss_gd + loss_iaf
+            # L2 Complex Loss
+            loss_com = F.mse_loss(clean_com, com_g) * 2
+            # L2 Consistency Loss
+            loss_stft = F.mse_loss(com_g, com_g_hat) * 2
+            # Time Loss
+            loss_time = F.l1_loss(clean_audio, audio_g)
+            # Metric Loss
+            metric_g = discriminator.forward(clean_mag, mag_g_hat)
+            loss_metric = F.mse_loss(metric_g.flatten(), one_labels)
+
+            loss_gen_all = loss_mag * 0.9 + loss_pha * 0.3  + loss_com * 0.1 + loss_stft * 0.1 + loss_metric * 0.05 + loss_time * 0.2
+
+            loss_gen_all.backward()
+            optim_g.step()
+
+            total_loss_d += loss_disc_all.item()
+            total_loss_g += loss_gen_all.item()
+            total_batches += 1
+
+            loss_d = round(total_loss_d / total_batches, 4)
+            loss_g = round(total_loss_g / total_batches, 4)
+
+            progress_bar.update(1)
+            progress_bar.set_postfix({
+                "loss_d": loss_d,
+                "loss_g": loss_g,
+            })
+
+        # evaluation
+        generator.eval()
+        discriminator.eval()
+
+        torch.cuda.empty_cache()
+        total_pesq_score = 0.
+        total_mag_err = 0.
+        total_pha_err = 0.
+        total_com_err = 0.
+        total_stft_err = 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_audio, noisy_audio = batch
+                clean_audio = clean_audio.to(device)
+                noisy_audio = noisy_audio.to(device)
+
+                clean_mag, clean_pha, clean_com = mag_pha_stft(clean_audio, config.n_fft, config.hop_size, config.win_size, config.compress_factor)
+                noisy_mag, noisy_pha, noisy_com = mag_pha_stft(noisy_audio, config.n_fft, config.hop_size, config.win_size, config.compress_factor)
+
+                mag_g, pha_g, com_g = generator.forward(noisy_mag, noisy_pha)
+
+                audio_g = mag_pha_istft(mag_g, pha_g, config.n_fft, config.hop_size, config.win_size, config.compress_factor)
+                mag_g_hat, pha_g_hat, com_g_hat = mag_pha_stft(audio_g, config.n_fft, config.hop_size, config.win_size, config.compress_factor)
+
+                total_pesq_score += pesq_score(
+                    torch.split(clean_audio, 1, dim=0),
+                    torch.split(audio_g, 1, dim=0),
+                    config
+                ).item()
+                total_mag_err += F.mse_loss(clean_mag, mag_g).item()
+                val_ip_err, val_gd_err, val_iaf_err = phase_losses(clean_pha, pha_g)
+                total_pha_err += (val_ip_err + val_gd_err + val_iaf_err).item()
+                total_com_err += F.mse_loss(clean_com, com_g).item()
+                total_stft_err += F.mse_loss(com_g, com_g_hat).item()
+
+                total_batches += 1
+
+                pesq_metric = round(total_pesq_score / total_batches, 4)
+                mag_err = round(total_mag_err / total_batches, 4)
+                pha_err = round(total_pha_err / total_batches, 4)
+                com_err = round(total_com_err / total_batches, 4)
+                stft_err = round(total_stft_err / total_batches, 4)
+
+                progress_bar.update(1)
+                progress_bar.set_postfix({
+                    "pesq_metric": pesq_metric,
+                    "mag_err": mag_err,
+                    "pha_err": pha_err,
+                    "com_err": com_err,
+                    "stft_err": stft_err,
+                })
+
+        # scheduler
+        scheduler_g.step()
+        scheduler_d.step()
+
+        # save path
+        epoch_dir = serialization_dir / "epoch-{}".format(idx_epoch)
+        epoch_dir.mkdir(parents=True, exist_ok=False)
+
+        # save models
+        generator.save_pretrained(epoch_dir.as_posix())
+        discriminator.save_pretrained(epoch_dir.as_posix())
+
+        # save optim
+        torch.save(optim_d.state_dict(), (epoch_dir / "optim_d.pth").as_posix())
+        torch.save(optim_g.state_dict(), (epoch_dir / "optim_g.pth").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 metric
+        if best_metric is None:
+            best_idx_epoch = idx_epoch
+            best_metric = pesq_metric
+        elif pesq_metric > best_metric:
+            # great is better.
+            best_idx_epoch = idx_epoch
+            best_metric = pesq_metric
+        else:
+            pass
+
+        metrics = {
+            "idx_epoch": idx_epoch,
+            "best_idx_epoch": best_idx_epoch,
+            "loss_d": loss_d,
+            "loss_g": loss_g,
+
+            "pesq_metric": pesq_metric,
+            "mag_err": mag_err,
+            "pha_err": pha_err,
+            "com_err": com_err,
+            "stft_err": stft_err,
+
+        }
+        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_aishell/step_3_evaluation.py

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

examples/clean_unet_aishell/yaml/config.yaml

@@ -0,0 +1,13 @@
+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/mpnet_aishell/run.sh

@@ -12,7 +12,7 @@ sh run.sh --stage 2 --stop_stage 2 --system_version centos --file_folder_name fi
 --noise_dir "/data/tianxing/HuggingDatasets/nx_noise/data/noise" \
 --speech_dir "/data/tianxing/HuggingDatasets/aishell/data_aishell/wav/train"
 
-sh run.sh --stage 2 --stop_stage 2 --system_version centos --file_folder_name file_dir \
+sh run.sh --stage 1 --stop_stage 2 --system_version centos --file_folder_name file_dir \
 --noise_dir "/data/tianxing/HuggingDatasets/nx_noise/data/noise" \
 --speech_dir "/data/tianxing/HuggingDatasets/aishell/data_aishell/wav/train"
 

examples/mpnet_aishell/step_2_train_model.py

@@ -20,18 +20,12 @@ sys.path.append(os.path.join(pwd, "../../"))
 
 import numpy as np
 import torch
-from torch.distributed import init_process_group
-import torch.multiprocessing as mp
-from torch.nn.parallel import DistributedDataParallel
-import torch.nn as nn
 from torch.nn import functional as F
-from torch.utils.data import DistributedSampler
 from torch.utils.data.dataloader import DataLoader
-import torchaudio
 from tqdm import tqdm
 
 from toolbox.torch.utils.data.dataset.denoise_excel_dataset import DenoiseExcelDataset
-from toolbox.torchaudio.models.mpnet.configuation_mpnet import MPNetConfig
+from toolbox.torchaudio.models.mpnet.configuration_mpnet import MPNetConfig
 from toolbox.torchaudio.models.mpnet.discriminator import MetricDiscriminatorPretrainedModel, batch_pesq
 from toolbox.torchaudio.models.mpnet.modeling_mpnet import MPNet, MPNetPretrainedModel, phase_losses, pesq_score
 from toolbox.torchaudio.models.mpnet.utils import mag_pha_stft, mag_pha_istft

examples/mpnet_aishell/step_3_evaluation.py

@@ -1,6 +1,184 @@
 #!/usr/bin/python3
 # -*- coding: utf-8 -*-
+"""
+https://github.com/yxlu-0102/MP-SENet/blob/main/inference.py
+"""
+import argparse
+import logging
+import os
+from pathlib import Path
+import sys
+import uuid
+
+pwd = os.path.abspath(os.path.dirname(__file__))
+sys.path.append(os.path.join(pwd, "../../"))
+
+import librosa
+import numpy as np
+import pandas as pd
+from scipy.io import wavfile
+import torch
+import torch.nn as nn
+import torchaudio
+from tqdm import tqdm
+
+from toolbox.torchaudio.models.mpnet.configuration_mpnet import MPNetConfig
+from toolbox.torchaudio.models.mpnet.modeling_mpnet import MPNetPretrainedModel
+from toolbox.torchaudio.models.mpnet.utils import mag_pha_stft, mag_pha_istft
+
+
+def get_args():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--valid_dataset", default="valid.xlsx", type=str)
+    parser.add_argument("--model_dir", default="serialization_dir/best", type=str)
+    parser.add_argument("--evaluation_audio_dir", default="evaluation_audio_dir", type=str)
+
+    parser.add_argument("--limit", default=10, type=int)
+
+    args = parser.parse_args()
+    return args
+
+
+def logging_config():
+    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)
+    stream_handler = logging.StreamHandler()
+    stream_handler.setLevel(logging.INFO)
+    stream_handler.setFormatter(logging.Formatter(fmt))
+
+    logger = logging.getLogger(__name__)
+
+    return logger
+
+
+def mix_speech_and_noise(speech: np.ndarray, noise: np.ndarray, snr_db: float):
+    l1 = len(speech)
+    l2 = len(noise)
+    l = min(l1, l2)
+    speech = speech[:l]
+    noise = noise[:l]
+
+    # np.float32, value between (-1, 1).
+
+    speech_power = np.mean(np.square(speech))
+    noise_power = speech_power / (10 ** (snr_db / 10))
+
+    noise_adjusted = np.sqrt(noise_power) * noise / np.sqrt(np.mean(noise ** 2))
+
+    noisy_signal = speech + noise_adjusted
+
+    return noisy_signal
+
+
+def save_audios(noise_audio: torch.Tensor,
+                clean_audio: torch.Tensor,
+                noisy_audio: torch.Tensor,
+                enhanced_audio: torch.Tensor,
+                output_dir: str,
+                sample_rate: int = 8000,
+                ):
+    basename = uuid.uuid4().__str__()
+    output_dir = Path(output_dir) / basename
+    output_dir.mkdir(parents=True, exist_ok=True)
+
+    filename = output_dir / "noise_audio.wav"
+    torchaudio.save(filename, noise_audio, sample_rate)
+    filename = output_dir / "clean_audio.wav"
+    torchaudio.save(filename, clean_audio, sample_rate)
+    filename = output_dir / "noisy_audio.wav"
+    torchaudio.save(filename, noisy_audio, sample_rate)
+
+    filename = output_dir / "enhanced_audio.wav"
+    torchaudio.save(filename, enhanced_audio, sample_rate)
+
+    return output_dir.as_posix()
+
+
+def main():
+    args = get_args()
+
+    logger = logging_config()
+
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    n_gpu = torch.cuda.device_count()
+    logger.info("GPU available count: {}; device: {}".format(n_gpu, device))
+
+    logger.info("prepare model")
+    config = MPNetConfig.from_pretrained(
+        pretrained_model_name_or_path=args.model_dir,
+    )
+    generator = MPNetPretrainedModel.from_pretrained(
+        pretrained_model_name_or_path=args.model_dir,
+    )
+    generator.to(device)
+    generator.eval()
+
+    logger.info("read excel")
+    df = pd.read_excel(args.valid_dataset)
+
+    progress_bar = tqdm(total=len(df), desc="Evaluation")
+    for idx, row in df.iterrows():
+        noise_filename = row["noise_filename"]
+        noise_offset = row["noise_offset"]
+        noise_duration = row["noise_duration"]
+
+        speech_filename = row["speech_filename"]
+        speech_offset = row["speech_offset"]
+        speech_duration = row["speech_duration"]
+
+        snr_db = row["snr_db"]
+
+        noise_audio, _ = librosa.load(
+            noise_filename,
+            sr=8000,
+            offset=noise_offset,
+            duration=noise_duration,
+        )
+        clean_audio, _ = librosa.load(
+            speech_filename,
+            sr=8000,
+            offset=speech_offset,
+            duration=speech_duration,
+        )
+        noisy_audio: np.ndarray = mix_speech_and_noise(
+            speech=clean_audio,
+            noise=noise_audio,
+            snr_db=snr_db,
+        )
+        noise_audio = torch.tensor(noise_audio, dtype=torch.float32)
+        clean_audio = torch.tensor(clean_audio, dtype=torch.float32)
+        noisy_audio: torch.Tensor = torch.tensor(noisy_audio, dtype=torch.float32)
+
+        noise_audio = noise_audio.unsqueeze(dim=0)
+        clean_audio = clean_audio.unsqueeze(dim=0)
+        noisy_audio = noisy_audio.unsqueeze(dim=0)
+
+        # inference
+        noisy_mag, noisy_pha, noisy_com = mag_pha_stft(
+            noisy_audio, config.n_fft, config.hop_size, config.win_size, config.compress_factor
+        )
+        mag_g, pha_g, com_g = generator.forward(noisy_mag, noisy_pha)
+        audio_g = mag_pha_istft(
+            mag_g, pha_g, config.n_fft, config.hop_size, config.win_size, config.compress_factor
+        )
+        enhanced_audio = audio_g.detach()
+
+        save_audios(
+            noise_audio, clean_audio, noisy_audio,
+            enhanced_audio,
+            args.evaluation_audio_dir
+        )
+
+        progress_bar.update(1)
+
+        if idx > args.limit:
+            break
+
+    return
 
 
 if __name__ == '__main__':
-    pass
+    main()

toolbox/torchaudio/models/clean_unet/configuration_clean_unet.py

@@ -0,0 +1,41 @@
+#!/usr/bin/python3
+# -*- coding: utf-8 -*-
+from toolbox.torchaudio.configuration_utils import PretrainedConfig
+
+
+class CleanUnetConfig(PretrainedConfig):
+    def __init__(self,
+                 channels_input: int = 1,
+                 channels_output: int = 1,
+
+                 channels_h: int = 64,
+                 max_h: int = 768,
+
+                 encoder_n_layers: int = 8,
+                 kernel_size: int = 4,
+                 stride: int = 2,
+                 tsfm_n_layers: int = 5,
+                 tsfm_n_head: int = 8,
+                 tsfm_d_model: int = 512,
+                 tsfm_d_inner: int = 2048,
+
+                 **kwargs
+                 ):
+        super(CleanUnetConfig, self).__init__(**kwargs)
+        self.channels_input = channels_input
+        self.channels_output = channels_output
+
+        self.channels_h = channels_h
+        self.max_h = max_h
+
+        self.encoder_n_layers = encoder_n_layers
+        self.kernel_size = kernel_size
+        self.stride = stride
+        self.tsfm_n_layers = tsfm_n_layers
+        self.tsfm_n_head = tsfm_n_head
+        self.tsfm_d_model = tsfm_d_model
+        self.tsfm_d_inner = tsfm_d_inner
+
+
+if __name__ == "__main__":
+    pass

toolbox/torchaudio/models/clean_unet/modeling_clean_unet.py

@@ -2,8 +2,287 @@
 # -*- coding: utf-8 -*-
 """
 https://arxiv.org/abs/2202.07790
+
+https://github.com/nvidia/cleanunet
+
+https://huggingface.co/spaces/fsoft-ai-center/Speech-Enhancement/blob/main/src/model.py
+
 """
+import os
+from typing import Optional, Union
+
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from toolbox.torchaudio.configuration_utils import CONFIG_FILE
+from toolbox.torchaudio.models.clean_unet.transformer import TransformerEncoder
+from toolbox.torchaudio.models.clean_unet.configuration_clean_unet import CleanUnetConfig
+
+
+def weight_scaling_init(layer):
+    """
+    weight rescaling initialization from https://arxiv.org/abs/1911.13254
+    """
+    w = layer.weight.detach()
+    alpha = 10.0 * w.std()
+    layer.weight.data /= torch.sqrt(alpha)
+    layer.bias.data /= torch.sqrt(alpha)
+
+
+def print_size(net, keyword=None):
+    """
+    Print the number of parameters of a network
+    """
+
+    if net is not None and isinstance(net, torch.nn.Module):
+        module_parameters = filter(lambda p: p.requires_grad, net.parameters())
+        params = sum([np.prod(p.size()) for p in module_parameters])
+
+        print("{} Parameters: {:.6f}M".format(
+            net.__class__.__name__, params / 1e6), flush=True, end="; ")
+
+        if keyword is not None:
+            keyword_parameters = [p for name, p in net.named_parameters() if p.requires_grad and keyword in name]
+            params = sum([np.prod(p.size()) for p in keyword_parameters])
+            print("{} Parameters: {:.6f}M".format(
+                keyword, params / 1e6), flush=True, end="; ")
+
+        print(" ")
+
+
+# CleanUNet architecture
+
+def padding(x, D, K, S):
+    """padding zeroes to x so that denoised audio has the same length"""
+
+    L = x.shape[-1]
+    for _ in range(D):
+        if L < K:
+            L = 1
+        else:
+            L = 1 + np.ceil((L - K) / S)
+
+    for _ in range(D):
+        L = (L - 1) * S + K
+
+    L = int(L)
+    x = F.pad(x, (0, L - x.shape[-1]))
+    return x
+
+
+class CleanUNet(nn.Module):
+    """
+    CleanUNet architecture.
+    """
+
+    def __init__(self,
+                 channels_input=1, channels_output=1,
+                 channels_h=64, max_h=768,
+                 encoder_n_layers=8, kernel_size=4, stride=2,
+                 tsfm_n_layers=3,
+                 tsfm_n_head=8,
+                 tsfm_d_model=512,
+                 tsfm_d_inner=2048):
+        """
+        Parameters:
+        channels_input (int):   input channels
+        channels_output (int):  output channels
+        channels_H (int):       middle channels H that controls capacity
+        max_H (int):            maximum H
+        encoder_n_layers (int): number of encoder/decoder layers D
+        kernel_size (int):      kernel size K
+        stride (int):           stride S
+        tsfm_n_layers (int):    number of self attention blocks N
+        tsfm_n_head (int):      number of heads in each self attention block
+        tsfm_d_model (int):     d_model of self attention
+        tsfm_d_inner (int):     d_inner of self attention
+        """
+
+        super(CleanUNet, self).__init__()
+
+        self.channels_input = channels_input
+        self.channels_output = channels_output
+        self.channels_h = channels_h
+        self.max_h = max_h
+        self.encoder_n_layers = encoder_n_layers
+        self.kernel_size = kernel_size
+        self.stride = stride
+
+        self.tsfm_n_layers = tsfm_n_layers
+        self.tsfm_n_head = tsfm_n_head
+        self.tsfm_d_model = tsfm_d_model
+        self.tsfm_d_inner = tsfm_d_inner
+
+        # encoder and decoder
+        self.encoder = nn.ModuleList()
+        self.decoder = nn.ModuleList()
+
+        for i in range(encoder_n_layers):
+            self.encoder.append(nn.Sequential(
+                nn.Conv1d(channels_input, channels_h, kernel_size, stride),
+                nn.ReLU(),
+                nn.Conv1d(channels_h, channels_h * 2, 1),
+                nn.GLU(dim=1)
+            ))
+            channels_input = channels_h
+
+            if i == 0:
+                # no relu at end
+                self.decoder.append(nn.Sequential(
+                    nn.Conv1d(channels_h, channels_h * 2, 1),
+                    nn.GLU(dim=1),
+                    nn.ConvTranspose1d(channels_h, channels_output, kernel_size, stride)
+                ))
+            else:
+                self.decoder.insert(0, nn.Sequential(
+                    nn.Conv1d(channels_h, channels_h * 2, 1),
+                    nn.GLU(dim=1),
+                    nn.ConvTranspose1d(channels_h, channels_output, kernel_size, stride),
+                    nn.ReLU()
+                ))
+            channels_output = channels_h
+
+            # double H but keep below max_H
+            channels_h *= 2
+            channels_h = min(channels_h, max_h)
+
+        # self attention block
+        self.tsfm_conv1 = nn.Conv1d(channels_output, tsfm_d_model, kernel_size=1)
+        self.tsfm_encoder = TransformerEncoder(d_word_vec=tsfm_d_model,
+                                               n_layers=tsfm_n_layers,
+                                               n_head=tsfm_n_head,
+                                               d_k=tsfm_d_model // tsfm_n_head,
+                                               d_v=tsfm_d_model // tsfm_n_head,
+                                               d_model=tsfm_d_model,
+                                               d_inner=tsfm_d_inner,
+                                               dropout=0.0,
+                                               n_position=0,
+                                               scale_emb=False)
+        self.tsfm_conv2 = nn.Conv1d(tsfm_d_model, channels_output, kernel_size=1)
+
+        # weight scaling initialization
+        for layer in self.modules():
+            if isinstance(layer, (nn.Conv1d, nn.ConvTranspose1d)):
+                weight_scaling_init(layer)
+
+    def forward(self, noisy_audio):
+        # (B, L) -> (B, C, L)
+        if len(noisy_audio.shape) == 2:
+            noisy_audio = noisy_audio.unsqueeze(1)
+        B, C, L = noisy_audio.shape
+        assert C == 1
+
+        # normalization and padding
+        std = noisy_audio.std(dim=2, keepdim=True) + 1e-3
+        noisy_audio /= std
+        x = padding(noisy_audio, self.encoder_n_layers, self.kernel_size, self.stride)
+
+        # encoder
+        skip_connections = []
+        for downsampling_block in self.encoder:
+            x = downsampling_block(x)
+            skip_connections.append(x)
+        skip_connections = skip_connections[::-1]
+
+        # attention mask for causal inference; for non-causal, set attn_mask to None
+        len_s = x.shape[-1]  # length at bottleneck
+        attn_mask = (1 - torch.triu(torch.ones((1, len_s, len_s), device=x.device), diagonal=1)).bool()
+
+        x = self.tsfm_conv1(x)  # C 1024 -> 512
+        x = x.permute(0, 2, 1)
+        x = self.tsfm_encoder(x, src_mask=attn_mask)
+        x = x.permute(0, 2, 1)
+        x = self.tsfm_conv2(x)  # C 512 -> 1024
+
+        # decoder
+        for i, upsampling_block in enumerate(self.decoder):
+            skip_i = skip_connections[i]
+            x += skip_i[:, :, :x.shape[-1]]
+            x = upsampling_block(x)
+
+        x = x[:, :, :L] * std
+        return x
+
+
+MODEL_FILE = "model.pt"
+
+
+class CleanUNetPretrainedModel(CleanUNet):
+    def __init__(self,
+                 config: CleanUnetConfig,
+                 ):
+        super(CleanUNetPretrainedModel, self).__init__(
+            channels_input=config.channels_input,
+            channels_output=config.channels_output,
+            channels_h=config.channels_h,
+            max_h=config.max_h,
+            encoder_n_layers=config.encoder_n_layers,
+            kernel_size=config.kernel_size,
+            stride=config.stride,
+            tsfm_n_layers=config.tsfm_n_layers,
+            tsfm_n_head=config.tsfm_n_head,
+            tsfm_d_model=config.tsfm_d_model,
+            tsfm_d_inner=config.tsfm_d_inner,
+        )
+        self.config = config
+
+    @classmethod
+    def from_pretrained(cls, pretrained_model_name_or_path, **kwargs):
+        config = CleanUnetConfig.from_pretrained(pretrained_model_name_or_path, **kwargs)
+
+        model = cls(config)
+
+        if os.path.isdir(pretrained_model_name_or_path):
+            ckpt_file = os.path.join(pretrained_model_name_or_path, MODEL_FILE)
+        else:
+            ckpt_file = pretrained_model_name_or_path
+
+        with open(ckpt_file, "rb") as f:
+            state_dict = torch.load(f, map_location="cpu", weights_only=True)
+        model.load_state_dict(state_dict, strict=True)
+        return model
+
+    def save_pretrained(self,
+                        save_directory: Union[str, os.PathLike],
+                        state_dict: Optional[dict] = None,
+                        ):
+
+        model = self
+
+        if state_dict is None:
+            state_dict = model.state_dict()
+
+        os.makedirs(save_directory, exist_ok=True)
+
+        # save state dict
+        model_file = os.path.join(save_directory, MODEL_FILE)
+        torch.save(state_dict, model_file)
+
+        # save config
+        config_file = os.path.join(save_directory, CONFIG_FILE)
+        self.config.to_yaml_file(config_file)
+        return save_directory
+
+
+def main():
+
+    config = CleanUnetConfig()
+    model = CleanUNetPretrainedModel(config)
+
+    print_size(model, keyword="tsfm")
+
+    input_data = torch.ones([4, 1, int(4.5 * 16000)])
+    output = model(input_data)
+    print(output.shape)
+
+    # y = torch.rand([4, 1, int(4.5 * 16000)])
+    # loss = torch.nn.MSELoss()(y, output)
+    # loss.backward()
+    # print(loss.item())
+    return
 
 
 if __name__ == '__main__':
-    pass
+    main()

toolbox/torchaudio/models/clean_unet/transformer.py

@@ -0,0 +1,216 @@
+#!/usr/bin/python3
+# -*- coding: utf-8 -*-
+import numpy as np
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+# Transformer (encoder) https://github.com/jadore801120/attention-is-all-you-need-pytorch
+# Original Copyright 2017 Victor Huang
+#  MIT License (https://opensource.org/licenses/MIT)
+
+class ScaledDotProductAttention(nn.Module):
+    """
+    Scaled Dot-Product Attention
+    """
+
+    def __init__(self, temperature, attn_dropout=0.1):
+        super().__init__()
+        self.temperature = temperature
+        self.dropout = nn.Dropout(attn_dropout)
+
+    def forward(self, q, k, v, mask=None):
+
+        attn = torch.matmul(q / self.temperature, k.transpose(2, 3))
+
+        if mask is not None:
+            attn = attn.masked_fill(mask == 0, -1e9)
+
+        attn = self.dropout(F.softmax(attn, dim=-1))
+        output = torch.matmul(attn, v)
+
+        return output, attn
+
+
+class MultiHeadAttention(nn.Module):
+    """
+    Multi-Head Attention module
+    """
+
+    def __init__(self, n_head, d_model, d_k, d_v, dropout=0.1):
+        super().__init__()
+
+        self.n_head = n_head
+        self.d_k = d_k
+        self.d_v = d_v
+
+        self.w_qs = nn.Linear(d_model, n_head * d_k, bias=False)
+        self.w_ks = nn.Linear(d_model, n_head * d_k, bias=False)
+        self.w_vs = nn.Linear(d_model, n_head * d_v, bias=False)
+        self.fc = nn.Linear(n_head * d_v, d_model, bias=False)
+
+        self.attention = ScaledDotProductAttention(temperature=d_k ** 0.5)
+
+        self.dropout = nn.Dropout(dropout)
+        self.layer_norm = nn.LayerNorm(d_model, eps=1e-6)
+
+    def forward(self, q, k, v, mask=None):
+
+        d_k, d_v, n_head = self.d_k, self.d_v, self.n_head
+        sz_b, len_q, len_k, len_v = q.size(0), q.size(1), k.size(1), v.size(1)
+
+        residual = q
+
+        # Pass through the pre-attention projection: b x lq x (n*dv)
+        # Separate different heads: b x lq x n x dv
+        q = self.w_qs(q).view(sz_b, len_q, n_head, d_k)
+        k = self.w_ks(k).view(sz_b, len_k, n_head, d_k)
+        v = self.w_vs(v).view(sz_b, len_v, n_head, d_v)
+
+        # Transpose for attention dot product: b x n x lq x dv
+        q, k, v = q.transpose(1, 2), k.transpose(1, 2), v.transpose(1, 2)
+
+        if mask is not None:
+            mask = mask.unsqueeze(1)   # For head axis broadcasting.
+
+        q, attn = self.attention(q, k, v, mask=mask)
+
+        # Transpose to move the head dimension back: b x lq x n x dv
+        # Combine the last two dimensions to concatenate all the heads together: b x lq x (n*dv)
+        q = q.transpose(1, 2).contiguous().view(sz_b, len_q, -1)
+        q = self.dropout(self.fc(q))
+        q += residual
+
+        q = self.layer_norm(q)
+
+        return q, attn
+
+
+class PositionwiseFeedForward(nn.Module):
+    """
+    A two-feed-forward-layer module
+    """
+
+    def __init__(self, d_in, d_hid, dropout=0.1):
+        super().__init__()
+        self.w_1 = nn.Linear(d_in, d_hid) # position-wise
+        self.w_2 = nn.Linear(d_hid, d_in) # position-wise
+        self.layer_norm = nn.LayerNorm(d_in, eps=1e-6)
+        self.dropout = nn.Dropout(dropout)
+
+    def forward(self, x):
+
+        residual = x
+
+        x = self.w_2(F.relu(self.w_1(x)))
+        x = self.dropout(x)
+        x += residual
+
+        x = self.layer_norm(x)
+
+        return x
+
+
+def get_subsequent_mask(seq):
+    """
+    For masking out the subsequent info.
+    """
+    sz_b, len_s = seq.size()
+    subsequent_mask = (1 - torch.triu(
+        torch.ones((1, len_s, len_s), device=seq.device), diagonal=1)).bool()
+    return subsequent_mask
+
+
+class PositionalEncoding(nn.Module):
+
+    def __init__(self, d_hid, n_position=200):
+        super(PositionalEncoding, self).__init__()
+
+        # Not a parameter
+        self.register_buffer('pos_table', self._get_sinusoid_encoding_table(n_position, d_hid))
+
+    def _get_sinusoid_encoding_table(self, n_position, d_hid):
+        """
+        Sinusoid position encoding table
+        """
+        # TODO: make it with torch instead of numpy
+
+        def get_position_angle_vec(position):
+            return [position / np.power(10000, 2 * (hid_j // 2) / d_hid) for hid_j in range(d_hid)]
+
+        sinusoid_table = np.array([get_position_angle_vec(pos_i) for pos_i in range(n_position)])
+        sinusoid_table[:, 0::2] = np.sin(sinusoid_table[:, 0::2])  # dim 2i
+        sinusoid_table[:, 1::2] = np.cos(sinusoid_table[:, 1::2])  # dim 2i+1
+
+        return torch.FloatTensor(sinusoid_table).unsqueeze(0)
+
+    def forward(self, x):
+        return x + self.pos_table[:, :x.size(1)].clone().detach()
+
+
+class EncoderLayer(nn.Module):
+    """
+    Compose with two layers
+    """
+
+    def __init__(self, d_model, d_inner, n_head, d_k, d_v, dropout=0.0):
+        super(EncoderLayer, self).__init__()
+        self.slf_attn = MultiHeadAttention(n_head, d_model, d_k, d_v, dropout=dropout)
+        self.pos_ffn = PositionwiseFeedForward(d_model, d_inner, dropout=dropout)
+
+    def forward(self, enc_input, slf_attn_mask=None):
+        enc_output, enc_slf_attn = self.slf_attn(
+            enc_input, enc_input, enc_input, mask=slf_attn_mask)
+        enc_output = self.pos_ffn(enc_output)
+        return enc_output, enc_slf_attn
+
+
+class TransformerEncoder(nn.Module):
+    """
+    A encoder model with self attention mechanism.
+    """
+
+    def __init__(
+            self, d_word_vec=512, n_layers=2, n_head=8, d_k=64, d_v=64,
+            d_model=512, d_inner=2048, dropout=0.1, n_position=624, scale_emb=False):
+
+        super().__init__()
+
+        # self.src_word_emb = nn.Embedding(n_src_vocab, d_word_vec, padding_idx=pad_idx)
+        if n_position > 0:
+            self.position_enc = PositionalEncoding(d_word_vec, n_position=n_position)
+        else:
+            self.position_enc = lambda x: x
+        self.dropout = nn.Dropout(p=dropout)
+        self.layer_stack = nn.ModuleList([
+            EncoderLayer(d_model, d_inner, n_head, d_k, d_v, dropout=dropout)
+            for _ in range(n_layers)])
+        self.layer_norm = nn.LayerNorm(d_model, eps=1e-6)
+        self.scale_emb = scale_emb
+        self.d_model = d_model
+
+    def forward(self, src_seq, src_mask, return_attns=False):
+
+        enc_slf_attn_list = []
+
+        # -- Forward
+        # enc_output = self.src_word_emb(src_seq)
+        enc_output = src_seq
+        if self.scale_emb:
+            enc_output *= self.d_model ** 0.5
+        enc_output = self.dropout(self.position_enc(enc_output))
+        enc_output = self.layer_norm(enc_output)
+
+        for enc_layer in self.layer_stack:
+            enc_output, enc_slf_attn = enc_layer(enc_output, slf_attn_mask=src_mask)
+            enc_slf_attn_list += [enc_slf_attn] if return_attns else []
+
+        if return_attns:
+            return enc_output, enc_slf_attn_list
+        return enc_output
+
+
+if __name__ == '__main__':
+    pass

toolbox/torchaudio/models/clean_unet/yaml/config.yaml

@@ -0,0 +1,13 @@
+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

toolbox/torchaudio/models/mpnet/modeling_mpnet.py

@@ -3,6 +3,8 @@
 """
 https://huggingface.co/spaces/LeeSangHoon/HierSpeech_TTS/blob/main/denoiser/generator.py
 
+https://huggingface.co/spaces/JacobLinCool/MP-SENet
+
 https://arxiv.org/abs/2305.13686
 https://github.com/yxlu-0102/MP-SENet
 
@@ -19,7 +21,7 @@ import torch.nn as nn
 from toolbox.torchaudio.configuration_utils import CONFIG_FILE
 from toolbox.torchaudio.models.mpnet.conformer import ConformerBlock
 from toolbox.torchaudio.models.mpnet.transformers import TransformerBlock
-from toolbox.torchaudio.models.mpnet.configuation_mpnet import MPNetConfig
+from toolbox.torchaudio.models.mpnet.configuration_mpnet import MPNetConfig
 from toolbox.torchaudio.models.mpnet.utils import LearnableSigmoid2d