update

examples/mpnet_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 3 --stop_stage 3 --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/mpnet_aishell/step_1_prepare_data.py

@@ -0,0 +1,197 @@
+#!/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)
+
+    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):
+
+        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/mpnet_aishell/step_2_train_model.py

@@ -0,0 +1,396 @@
+#!/usr/bin/python3
+# -*- coding: utf-8 -*-
+"""
+https://github.com/yxlu-0102/MP-SENet/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.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.discriminator import MetricDiscriminator, 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
+
+
+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=1e-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,
+                 n_fft: int = 512,
+                 win_length: int = 200,
+                 hop_length: int = 80,
+                 window_fn: str = "hamming",
+                 irm_beta: float = 1.0,
+                 epsilon: float = 1e-8,
+                 ):
+        self.n_fft = n_fft
+        self.win_length = win_length
+        self.hop_length = hop_length
+        self.window_fn = window_fn
+        self.irm_beta = irm_beta
+        self.epsilon = epsilon
+
+        self.transform = torchaudio.transforms.Spectrogram(
+            n_fft=self.n_fft,
+            win_length=self.win_length,
+            hop_length=self.hop_length,
+            power=2.0,
+            window_fn=torch.hamming_window if window_fn == "hamming" else torch.hann_window,
+        )
+
+    @staticmethod
+    def make_unfold_snr_db(x: torch.Tensor, n_time_steps: int = 3):
+        batch_size, channels, freq_dim, time_steps = x.shape
+
+        # kernel: [freq_dim, n_time_step]
+        kernel_size = (freq_dim, n_time_steps)
+
+        # pad
+        pad = n_time_steps // 2
+        x = torch.concat(tensors=[
+            x[:, :, :, :pad],
+            x,
+            x[:, :, :, -pad:],
+        ], dim=-1)
+
+        x = F.unfold(
+            input=x,
+            kernel_size=kernel_size,
+        )
+        # x shape: [batch_size, fold, time_steps]
+        return x
+
+    def __call__(self, batch: List[dict]):
+        mix_spec_list = list()
+        speech_irm_list = list()
+        snr_db_list = list()
+        for sample in batch:
+            noise_wave: torch.Tensor = sample["noise_wave"]
+            speech_wave: torch.Tensor = sample["speech_wave"]
+            mix_wave: torch.Tensor = sample["mix_wave"]
+            # snr_db: float = sample["snr_db"]
+
+            noise_spec = self.transform.forward(noise_wave)
+            speech_spec = self.transform.forward(speech_wave)
+            mix_spec = self.transform.forward(mix_wave)
+
+            # noise_irm = noise_spec / (noise_spec + speech_spec)
+            speech_irm = speech_spec / (noise_spec + speech_spec + self.epsilon)
+            speech_irm = torch.pow(speech_irm, self.irm_beta)
+
+            # noise_spec, speech_spec, mix_spec, speech_irm
+            # shape: [freq_dim, time_steps]
+
+            snr_db: torch.Tensor = 10 * torch.log10(
+                speech_spec / (noise_spec + self.epsilon)
+            )
+            snr_db = torch.clamp(snr_db, min=self.epsilon)
+
+            snr_db_ = torch.unsqueeze(snr_db, dim=0)
+            snr_db_ = torch.unsqueeze(snr_db_, dim=0)
+            snr_db_ = self.make_unfold_snr_db(snr_db_, n_time_steps=3)
+            snr_db_ = torch.squeeze(snr_db_, dim=0)
+            # snr_db_ shape: [fold, time_steps]
+
+            snr_db = torch.mean(snr_db_, dim=0, keepdim=True)
+            # snr_db shape: [1, time_steps]
+
+            mix_spec_list.append(mix_spec)
+            speech_irm_list.append(speech_irm)
+            snr_db_list.append(snr_db)
+
+        mix_spec_list = torch.stack(mix_spec_list)
+        speech_irm_list = torch.stack(speech_irm_list)
+        snr_db_list = torch.stack(snr_db_list)  # shape: (batch_size, time_steps, 1)
+
+        mix_spec_list = mix_spec_list[:, :-1, :]
+        speech_irm_list = speech_irm_list[:, :-1, :]
+
+        # mix_spec_list shape: [batch_size, freq_dim, time_steps]
+        # speech_irm_list shape: [batch_size, freq_dim, time_steps]
+        # snr_db shape: [batch_size, 1, time_steps]
+
+        # assert
+        if torch.any(torch.isnan(mix_spec_list)) or torch.any(torch.isinf(mix_spec_list)):
+            raise AssertionError("nan or inf in mix_spec_list")
+        if torch.any(torch.isnan(speech_irm_list)) or torch.any(torch.isinf(speech_irm_list)):
+            raise AssertionError("nan or inf in speech_irm_list")
+        if torch.any(torch.isnan(snr_db_list)) or torch.any(torch.isinf(snr_db_list)):
+            raise AssertionError("nan or inf in snr_db_list")
+
+        return mix_spec_list, speech_irm_list, snr_db_list
+
+
+collate_fn = CollateFunction()
+
+
+def main():
+    args = get_args()
+
+    config = MPNetConfig.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=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}")
+    generator = MPNetPretrainedModel(config).to(device)
+    discriminator = MetricDiscriminator().to(device)
+
+    # optimizer
+    logger.info("prepare optimizer, lr_scheduler, loss_fn, categorical_accuracy")
+    num_params = 0
+    for p in generator.parameters():
+        num_params += p.numel()
+    print("Total Parameters (generator): {:.3f}M".format(num_params/1e6))
+
+    optim_g = torch.optim.AdamW(generator.parameters(), config.learning_rate, betas=[config.adam_b1, config.adam_b2])
+    optim_d = torch.optim.AdamW(discriminator.parameters(), config.learning_rate, betas=[config.adam_b1, config.adam_b2])
+
+    scheduler_g = torch.optim.lr_scheduler.ExponentialLR(optim_g, gamma=config.lr_decay, last_epoch=-1)
+    scheduler_d = torch.optim.lr_scheduler.ExponentialLR(optim_d, gamma=config.lr_decay, last_epoch=-1)
+
+    # training loop
+    logger.info("training")
+    for idx_epoch in range(args.max_epochs):
+        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 = torch.autograd.Variable(clean_audio.to(device, non_blocking=True))
+            noisy_audio = torch.autograd.Variable(noisy_audio.to(device, non_blocking=True))
+            one_labels = torch.ones(config.batch_size).to(device, non_blocking=True)
+
+            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
+
+            progress_bar.update(1)
+            progress_bar.set_postfix({
+                "loss_d": round(total_loss_d / total_batches, 4),
+                "loss_g": round(total_loss_g / total_batches, 4),
+            })
+
+        generator.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 = torch.autograd.Variable(clean_audio.to(device, non_blocking=True))
+                noisy_audio = torch.autograd.Variable(noisy_audio.to(device, non_blocking=True))
+
+                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
+
+                progress_bar.update(1)
+                progress_bar.set_postfix({
+                    "pesq_score": 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),
+
+                })
+
+    return
+
+
+if __name__ == '__main__':
+    main()

examples/mpnet_aishell/step_3_evaluation.py

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

examples/mpnet_aishell/yaml/config.yaml

@@ -0,0 +1,27 @@
+model_name: "mpnet"
+
+num_gpus: 0
+batch_size: 4
+learning_rate: 0.0005
+adam_b1: 0.8
+adam_b2: 0.99
+lr_decay: 0.99
+seed: 1234
+
+dense_channel: 64
+compress_factor: 0.3
+num_tsconformers: 4
+beta: 2.0
+
+sample_rate: 16000
+segment_size: 32000
+n_fft: 400
+hop_size: 100
+win_size: 400
+
+num_workers: 4
+
+dist_config:
+  dist_backend: nccl
+  dist_url: tcp://localhost:54321
+  world_size: 1

examples/spectrum_dfnet_aishell/step_3_evaluation.py

@@ -255,7 +255,6 @@ def main():
         # speech_irm_prediction shape: [batch_size, freq_dim (256), time_steps]
         batch_size, _, time_steps = speech_irm_prediction.shape
 
-
         mix_spec_complex = torch.concat(
             [
                 mix_spec_complex,

requirements-python-3-9-9.txt

@@ -11,3 +11,4 @@ overrides==7.7.0
 torch-pesq
 torchmetrics
 torchmetrics[audio]
+einops

requirements.txt

@@ -10,4 +10,5 @@ torchaudio==2.5.1
 overrides==7.7.0
 torch-pesq==0.1.2
 torchmetrics==1.6.1
-torchmetrics[audio]
+torchmetrics[audio]==1.6.1
+einops==0.8.1

toolbox/torchaudio/models/mpnet/configuation_mpnet.py

@@ -0,0 +1,68 @@
+#!/usr/bin/python3
+# -*- coding: utf-8 -*-
+from typing import Tuple
+
+from toolbox.torchaudio.configuration_utils import PretrainedConfig
+
+
+class MPNetConfig(PretrainedConfig):
+    """
+    https://github.com/yxlu-0102/MP-SENet/blob/main/config.json
+    """
+    def __init__(self,
+                 num_gpus: int = 0,
+                 batch_size: int = 4,
+                 learning_rate: float = 0.0005,
+                 adam_b1: float = 0.8,
+                 adam_b2: float = 0.99,
+                 lr_decay: float = 0.99,
+                 seed: int = 1234,
+
+                 dense_channel: int = 64,
+                 compress_factor: float = 0.3,
+                 num_tsconformers: int = 4,
+                 beta: float = 2.0,
+
+                 sample_rate: int = 16000,
+                 segment_size: int = 32000,
+                 n_fft: int = 400,
+                 hop_size: int = 100,
+                 win_size: int = 400,
+
+                 num_workers: int = 4,
+
+                 dist_config: dict = None,
+
+                 **kwargs
+                 ):
+        super(MPNetConfig, self).__init__(**kwargs)
+        self.num_gpus = num_gpus
+        self.batch_size = batch_size
+        self.learning_rate = learning_rate
+        self.adam_b1 = adam_b1
+        self.adam_b2 = adam_b2
+        self.lr_decay = lr_decay
+        self.seed = seed
+
+        self.dense_channel = dense_channel
+        self.compress_factor = compress_factor
+        self.num_tsconformers = num_tsconformers
+        self.beta = beta
+
+        self.sample_rate = sample_rate
+        self.segment_size = segment_size
+        self.n_fft = n_fft
+        self.hop_size = hop_size
+        self.win_size = win_size
+
+        self.num_workers = num_workers
+
+        self.dist_config = dist_config or {
+            "dist_backend": "nccl",
+            "dist_url": "tcp://localhost:54321",
+            "world_size": 1
+        }
+
+
+if __name__ == "__main__":
+    pass

toolbox/torchaudio/models/mpnet/conformer.py

@@ -0,0 +1,83 @@
+#!/usr/bin/python3
+# -*- coding: utf-8 -*-
+from einops.layers.torch import Rearrange
+import torch.nn as nn
+
+
+def get_padding(kernel_size: int, dilation: int = 1):
+    return int((kernel_size * dilation - dilation) / 2)
+
+
+class FeedForwardModule(nn.Module):
+    def __init__(self, dim, mult=4, dropout=0):
+        super(FeedForwardModule, self).__init__()
+        self.ffm = nn.Sequential(
+            nn.LayerNorm(dim),
+            nn.Linear(dim, dim * mult),
+            nn.SiLU(),
+            nn.Dropout(dropout),
+            nn.Linear(dim * mult, dim),
+            nn.Dropout(dropout)
+        )
+
+    def forward(self, x):
+        return self.ffm(x)
+
+
+class ConformerConvModule(nn.Module):
+    def __init__(self, dim, expansion_factor=2, kernel_size=31, dropout=0.):
+        super(ConformerConvModule, self).__init__()
+        inner_dim = dim * expansion_factor
+        self.ccm = nn.Sequential(
+            nn.LayerNorm(dim),
+            Rearrange('b n c -> b c n'),
+            nn.Conv1d(dim, inner_dim*2, 1),
+            nn.GLU(dim=1),
+            nn.Conv1d(inner_dim, inner_dim, kernel_size=kernel_size,
+                      padding=get_padding(kernel_size), groups=inner_dim), # DepthWiseConv1d
+            nn.BatchNorm1d(inner_dim),
+            nn.SiLU(),
+            nn.Conv1d(inner_dim, dim, 1),
+            Rearrange('b c n -> b n c'),
+            nn.Dropout(dropout)
+        )
+
+    def forward(self, x):
+        return self.ccm(x)
+
+
+class AttentionModule(nn.Module):
+    def __init__(self, dim, n_head=8, dropout=0.):
+        super(AttentionModule, self).__init__()
+        self.attn = nn.MultiheadAttention(dim, n_head, dropout=dropout)
+        self.layernorm = nn.LayerNorm(dim)
+
+    def forward(self, x, attn_mask=None, key_padding_mask=None):
+        x = self.layernorm(x)
+        x, _ = self.attn(x, x, x,
+                         attn_mask=attn_mask,
+                         key_padding_mask=key_padding_mask)
+        return x
+
+
+class ConformerBlock(nn.Module):
+    def __init__(self, dim, n_head=8, ffm_mult=4, ccm_expansion_factor=2, ccm_kernel_size=31,
+                 ffm_dropout=0., attn_dropout=0., ccm_dropout=0.):
+        super(ConformerBlock, self).__init__()
+        self.ffm1 = FeedForwardModule(dim, ffm_mult, dropout=ffm_dropout)
+        self.attn = AttentionModule(dim, n_head, dropout=attn_dropout)
+        self.ccm = ConformerConvModule(dim, ccm_expansion_factor, ccm_kernel_size, dropout=ccm_dropout)
+        self.ffm2 = FeedForwardModule(dim, ffm_mult, dropout=ffm_dropout)
+        self.post_norm = nn.LayerNorm(dim)
+
+    def forward(self, x):
+        x = x + 0.5 * self.ffm1(x)
+        x = x + self.attn(x)
+        x = x + self.ccm(x)
+        x = x + 0.5 * self.ffm2(x)
+        x = self.post_norm(x)
+        return x
+
+
+if __name__ == '__main__':
+    pass

toolbox/torchaudio/models/mpnet/discriminator.py

@@ -0,0 +1,71 @@
+#!/usr/bin/python3
+# -*- coding: utf-8 -*-
+import torch
+import torch.nn as nn
+import numpy as np
+import torch.nn.functional as F
+from pesq import pesq
+from joblib import Parallel, delayed
+
+from toolbox.torchaudio.models.mpnet.utils import LearnableSigmoid1d
+
+
+def cal_pesq(clean, noisy, sr=16000):
+    try:
+        pesq_score = pesq(sr, clean, noisy, 'wb')
+    except:
+        # error can happen due to silent period
+        pesq_score = -1
+    return pesq_score
+
+
+def batch_pesq(clean, noisy):
+    pesq_score = Parallel(n_jobs=15)(delayed(cal_pesq)(c, n) for c, n in zip(clean, noisy))
+    pesq_score = np.array(pesq_score)
+    if -1 in pesq_score:
+        return None
+    pesq_score = (pesq_score - 1) / 3.5
+    return torch.FloatTensor(pesq_score)
+
+
+def metric_loss(metric_ref, metrics_gen):
+    loss = 0
+    for metric_gen in metrics_gen:
+        metric_loss = F.mse_loss(metric_ref, metric_gen.flatten())
+        loss += metric_loss
+
+    return loss
+
+
+class MetricDiscriminator(nn.Module):
+    def __init__(self, dim=16, in_channel=2):
+        super(MetricDiscriminator, self).__init__()
+        self.layers = nn.Sequential(
+            nn.utils.spectral_norm(nn.Conv2d(in_channel, dim, (4,4), (2,2), (1,1), bias=False)),
+            nn.InstanceNorm2d(dim, affine=True),
+            nn.PReLU(dim),
+            nn.utils.spectral_norm(nn.Conv2d(dim, dim*2, (4,4), (2,2), (1,1), bias=False)),
+            nn.InstanceNorm2d(dim*2, affine=True),
+            nn.PReLU(dim*2),
+            nn.utils.spectral_norm(nn.Conv2d(dim*2, dim*4, (4,4), (2,2), (1,1), bias=False)),
+            nn.InstanceNorm2d(dim*4, affine=True),
+            nn.PReLU(dim*4),
+            nn.utils.spectral_norm(nn.Conv2d(dim*4, dim*8, (4,4), (2,2), (1,1), bias=False)),
+            nn.InstanceNorm2d(dim*8, affine=True),
+            nn.PReLU(dim*8),
+            nn.AdaptiveMaxPool2d(1),
+            nn.Flatten(),
+            nn.utils.spectral_norm(nn.Linear(dim*8, dim*4)),
+            nn.Dropout(0.3),
+            nn.PReLU(dim*4),
+            nn.utils.spectral_norm(nn.Linear(dim*4, 1)),
+            LearnableSigmoid1d(1)
+        )
+
+    def forward(self, x, y):
+        xy = torch.stack((x, y), dim=1)
+        return self.layers(xy)
+
+
+if __name__ == '__main__':
+    pass

toolbox/torchaudio/models/mpnet/modeling_mpnet.py

@@ -2,8 +2,295 @@
 # -*- coding: utf-8 -*-
 """
 https://huggingface.co/spaces/LeeSangHoon/HierSpeech_TTS/blob/main/denoiser/generator.py
+
+https://arxiv.org/abs/2305.13686
+https://github.com/yxlu-0102/MP-SENet
+
 """
+import os
+from typing import Optional, Union
+
+from pesq import pesq
+from joblib import Parallel, delayed
+import numpy as np
+import torch
+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.utils import LearnableSigmoid2d
+
+
+class SPConvTranspose2d(nn.Module):
+    def __init__(self, in_channels, out_channels, kernel_size, r=1):
+        super(SPConvTranspose2d, self).__init__()
+        self.pad1 = nn.ConstantPad2d((1, 1, 0, 0), value=0.)
+        self.out_channels = out_channels
+        self.conv = nn.Conv2d(in_channels, out_channels * r, kernel_size=kernel_size, stride=(1, 1))
+        self.r = r
+
+    def forward(self, x):
+        x = self.pad1(x)
+        out = self.conv(x)
+        batch_size, nchannels, H, W = out.shape
+        out = out.view((batch_size, self.r, nchannels // self.r, H, W))
+        out = out.permute(0, 2, 3, 4, 1)
+        out = out.contiguous().view((batch_size, nchannels // self.r, H, -1))
+        return out
+
+
+class DenseBlock(nn.Module):
+    def __init__(self, h, kernel_size=(2, 3), depth=4):
+        super(DenseBlock, self).__init__()
+        self.h = h
+        self.depth = depth
+        self.dense_block = nn.ModuleList([])
+        for i in range(depth):
+            dilation = 2 ** i
+            pad_length = dilation
+            dense_conv = nn.Sequential(
+                nn.ConstantPad2d((1, 1, pad_length, 0), value=0.),
+                nn.Conv2d(h.dense_channel*(i+1), h.dense_channel, kernel_size, dilation=(dilation, 1)),
+                nn.InstanceNorm2d(h.dense_channel, affine=True),
+                nn.PReLU(h.dense_channel)
+            )
+            self.dense_block.append(dense_conv)
+
+    def forward(self, x):
+        skip = x
+        for i in range(self.depth):
+            x = self.dense_block[i](skip)
+            skip = torch.cat([x, skip], dim=1)
+        return x
+
+
+class DenseEncoder(nn.Module):
+    def __init__(self, h, in_channel):
+        super(DenseEncoder, self).__init__()
+        self.h = h
+        self.dense_conv_1 = nn.Sequential(
+            nn.Conv2d(in_channel, h.dense_channel, (1, 1)),
+            nn.InstanceNorm2d(h.dense_channel, affine=True),
+            nn.PReLU(h.dense_channel))
+
+        self.dense_block = DenseBlock(h, depth=4)
+
+        self.dense_conv_2 = nn.Sequential(
+            nn.Conv2d(h.dense_channel, h.dense_channel, (1, 3), (1, 2), padding=(0, 1)),
+            nn.InstanceNorm2d(h.dense_channel, affine=True),
+            nn.PReLU(h.dense_channel))
+
+    def forward(self, x):
+        x = self.dense_conv_1(x)  # [b, 64, T, F]
+        x = self.dense_block(x)   # [b, 64, T, F]
+        x = self.dense_conv_2(x)  # [b, 64, T, F//2]
+        return x
+
+
+class MaskDecoder(nn.Module):
+    def __init__(self, h, out_channel=1):
+        super(MaskDecoder, self).__init__()
+        self.dense_block = DenseBlock(h, depth=4)
+        self.mask_conv = nn.Sequential(
+            SPConvTranspose2d(h.dense_channel, h.dense_channel, (1, 3), 2),
+            nn.InstanceNorm2d(h.dense_channel, affine=True),
+            nn.PReLU(h.dense_channel),
+            nn.Conv2d(h.dense_channel, out_channel, (1, 2))
+        )
+        self.lsigmoid = LearnableSigmoid2d(h.n_fft//2+1, beta=h.beta)
+
+    def forward(self, x):
+        x = self.dense_block(x)
+        x = self.mask_conv(x)
+        x = x.permute(0, 3, 2, 1).squeeze(-1) # [B, F, T]
+        x = self.lsigmoid(x)
+        return x
+
+
+class PhaseDecoder(nn.Module):
+    def __init__(self, h, out_channel=1):
+        super(PhaseDecoder, self).__init__()
+        self.dense_block = DenseBlock(h, depth=4)
+        self.phase_conv = nn.Sequential(
+            SPConvTranspose2d(h.dense_channel, h.dense_channel, (1, 3), 2),
+            nn.InstanceNorm2d(h.dense_channel, affine=True),
+            nn.PReLU(h.dense_channel)
+        )
+        self.phase_conv_r = nn.Conv2d(h.dense_channel, out_channel, (1, 2))
+        self.phase_conv_i = nn.Conv2d(h.dense_channel, out_channel, (1, 2))
+
+    def forward(self, x):
+        x = self.dense_block(x)
+        x = self.phase_conv(x)
+        x_r = self.phase_conv_r(x)
+        x_i = self.phase_conv_i(x)
+        x = torch.atan2(x_i, x_r)
+        x = x.permute(0, 3, 2, 1).squeeze(-1)  # [B, F, T]
+        return x
+
+
+class TSTransformerBlock(nn.Module):
+    def __init__(self, h):
+        super(TSTransformerBlock, self).__init__()
+        self.h = h
+        self.time_transformer = TransformerBlock(d_model=h.dense_channel, n_heads=4)
+        self.freq_transformer = TransformerBlock(d_model=h.dense_channel, n_heads=4)
+
+    def forward(self, x):
+        b, c, t, f = x.size()
+        x = x.permute(0, 3, 2, 1).contiguous().view(b*f, t, c)
+        x = self.time_transformer(x) + x
+        x = x.view(b, f, t, c).permute(0, 2, 1, 3).contiguous().view(b*t, f, c)
+        x = self.freq_transformer(x) + x
+        x = x.view(b, t, f, c).permute(0, 3, 1, 2)
+        return x
+
+
+class MPNet(nn.Module):
+    def __init__(self, config: MPNetConfig, num_tsblocks=4):
+        super(MPNet, self).__init__()
+        self.config = config
+        self.num_tscblocks = num_tsblocks
+        self.dense_encoder = DenseEncoder(config, in_channel=2)
+
+        self.TSTransformer = nn.ModuleList([])
+        for i in range(num_tsblocks):
+            self.TSTransformer.append(TSTransformerBlock(config))
+
+        self.mask_decoder = MaskDecoder(config, out_channel=1)
+        self.phase_decoder = PhaseDecoder(config, out_channel=1)
+
+    def forward(self, noisy_amp, noisy_pha):  # [B, F, T]
+
+        x = torch.stack((noisy_amp, noisy_pha), dim=-1).permute(0, 3, 2, 1)  # [B, 2, T, F]
+        x = self.dense_encoder(x)
+
+        for i in range(self.num_tscblocks):
+            x = self.TSTransformer[i](x)
+
+        denoised_amp = noisy_amp * self.mask_decoder(x)
+        denoised_pha = self.phase_decoder(x)
+        denoised_com = torch.stack(
+            tensors=(
+                denoised_amp * torch.cos(denoised_pha),
+                denoised_amp * torch.sin(denoised_pha)
+            ),
+            dim=-1
+        )
+
+        return denoised_amp, denoised_pha, denoised_com
+
+
+MODEL_FILE = "model.pt"
+
+
+class MPNetPretrainedModel(MPNet):
+    def __init__(self,
+                 config: MPNetConfig,
+                 ):
+        super(MPNetPretrainedModel, self).__init__(
+            config=config,
+        )
+
+    @classmethod
+    def from_pretrained(cls, pretrained_model_name_or_path, **kwargs):
+        config = MPNetConfig.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 phase_losses(phase_r, phase_g):
+
+    ip_loss = torch.mean(anti_wrapping_function(phase_r - phase_g))
+    gd_loss = torch.mean(anti_wrapping_function(torch.diff(phase_r, dim=1) - torch.diff(phase_g, dim=1)))
+    iaf_loss = torch.mean(anti_wrapping_function(torch.diff(phase_r, dim=2) - torch.diff(phase_g, dim=2)))
+
+    return ip_loss, gd_loss, iaf_loss
+
+
+def anti_wrapping_function(x):
+
+    return torch.abs(x - torch.round(x / (2 * np.pi)) * 2 * np.pi)
+
+
+def pesq_score(utts_r, utts_g, h):
+
+    pesq_score = Parallel(n_jobs=30)(delayed(eval_pesq)(
+                            utts_r[i].squeeze().cpu().numpy(),
+                            utts_g[i].squeeze().cpu().numpy(),
+                            h.sampling_rate)
+                          for i in range(len(utts_r)))
+    pesq_score = np.mean(pesq_score)
+
+    return pesq_score
+
+
+def eval_pesq(clean_utt, esti_utt, sr):
+    try:
+        pesq_score = pesq(sr, clean_utt, esti_utt)
+    except:
+        pesq_score = -1
+
+    return pesq_score
+
+
+def main():
+    import torchaudio
+
+    config = MPNetConfig()
+    model = MPNet(config=config)
+
+    transformer = torchaudio.transforms.Spectrogram(
+        n_fft=config.n_fft,
+        win_length=config.win_size,
+        hop_length=config.hop_size,
+        window_fn=torch.hamming_window,
+    )
+
+    inputs = torch.randn(size=(1, 32000), dtype=torch.float32)
+    spec = transformer.forward(inputs)
+    print(spec.shape)
+
+    denoised_amp, denoised_pha, denoised_com = model.forward(spec, spec)
+    print(denoised_amp.shape)
+    print(denoised_pha.shape)
+    print(denoised_com.shape)
+
+    return
 
 
 if __name__ == '__main__':
-    pass
+    main()

toolbox/torchaudio/models/mpnet/transformers.py

@@ -0,0 +1,70 @@
+#!/usr/bin/python3
+# -*- coding: utf-8 -*-
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch.nn import MultiheadAttention, GRU, Linear, LayerNorm, Dropout
+
+
+class FFN(nn.Module):
+    def __init__(self, d_model, bidirectional=True, dropout=0):
+        super(FFN, self).__init__()
+        self.gru = GRU(d_model, d_model * 2, 1, bidirectional=bidirectional)
+        if bidirectional:
+            self.linear = Linear(d_model * 2 * 2, d_model)
+        else:
+            self.linear = Linear(d_model * 2, d_model)
+        self.dropout = Dropout(dropout)
+
+    def forward(self, x):
+        self.gru.flatten_parameters()
+        x, _ = self.gru(x)
+        x = F.leaky_relu(x)
+        x = self.dropout(x)
+        x = self.linear(x)
+
+        return x
+
+
+class TransformerBlock(nn.Module):
+    def __init__(self, d_model, n_heads, bidirectional=True, dropout=0):
+        super(TransformerBlock, self).__init__()
+
+        self.norm1 = LayerNorm(d_model)
+        self.attention = MultiheadAttention(d_model, n_heads, dropout=dropout)
+        self.dropout1 = Dropout(dropout)
+
+        self.norm2 = LayerNorm(d_model)
+        self.ffn = FFN(d_model, bidirectional=bidirectional)
+        self.dropout2 = Dropout(dropout)
+
+        self.norm3 = LayerNorm(d_model)
+
+    def forward(self, x, attn_mask=None, key_padding_mask=None):
+        xt = self.norm1(x)
+        xt, _ = self.attention(xt, xt, xt,
+                               attn_mask=attn_mask,
+                               key_padding_mask=key_padding_mask)
+        x = x + self.dropout1(xt)
+
+        xt = self.norm2(x)
+        xt = self.ffn(xt)
+        x = x + self.dropout2(xt)
+
+        x = self.norm3(x)
+
+        return x
+
+
+def main():
+    x = torch.randn(4, 64, 401, 201)
+    b, c, t, f = x.size()
+    x = x.permute(0, 3, 2, 1).contiguous().view(b, f * t, c)
+    transformer = TransformerBlock(d_model=64, n_heads=4)
+    x = transformer(x)
+    x = x.view(b, f, t, c).permute(0, 3, 2, 1)
+    print(x.size())
+
+
+if __name__ == '__main__':
+    main()

toolbox/torchaudio/models/mpnet/utils.py

@@ -0,0 +1,106 @@
+#!/usr/bin/python3
+# -*- coding: utf-8 -*-
+from einops.layers.torch import Rearrange
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import numpy as np
+from pesq import pesq
+from joblib import Parallel, delayed
+
+
+def phase_losses(phase_r, phase_g):
+
+    ip_loss = torch.mean(anti_wrapping_function(phase_r - phase_g))
+    gd_loss = torch.mean(anti_wrapping_function(torch.diff(phase_r, dim=1) - torch.diff(phase_g, dim=1)))
+    iaf_loss = torch.mean(anti_wrapping_function(torch.diff(phase_r, dim=2) - torch.diff(phase_g, dim=2)))
+
+    return ip_loss, gd_loss, iaf_loss
+
+
+def anti_wrapping_function(x):
+
+    return torch.abs(x - torch.round(x / (2 * np.pi)) * 2 * np.pi)
+
+
+def pesq_score(utts_r, utts_g, h):
+
+    pesq_score = Parallel(n_jobs=30)(delayed(eval_pesq)(
+                            utts_r[i].squeeze().cpu().numpy(),
+                            utts_g[i].squeeze().cpu().numpy(),
+                            h.sampling_rate)
+                          for i in range(len(utts_r)))
+    pesq_score = np.mean(pesq_score)
+
+    return pesq_score
+
+
+def eval_pesq(clean_utt, esti_utt, sr):
+    try:
+        pesq_score = pesq(sr, clean_utt, esti_utt)
+    except:
+        pesq_score = -1
+
+    return pesq_score
+
+
+def mag_pha_stft(y, n_fft, hop_size, win_size, compress_factor=1.0, center=True):
+
+    hann_window = torch.hann_window(win_size).to(y.device)
+    stft_spec = torch.stft(y, n_fft, hop_length=hop_size, win_length=win_size, window=hann_window,
+                           center=center, pad_mode='reflect', normalized=False, return_complex=True)
+    stft_spec = torch.view_as_real(stft_spec)
+    mag = torch.sqrt(stft_spec.pow(2).sum(-1) + 1e-9)
+    pha = torch.atan2(stft_spec[:, :, :, 1] + 1e-10, stft_spec[:, :, :, 0] + 1e-5)
+    # Magnitude Compression
+    mag = torch.pow(mag, compress_factor)
+    com = torch.stack((mag*torch.cos(pha), mag*torch.sin(pha)), dim=-1)
+
+    return mag, pha, com
+
+
+def mag_pha_istft(mag, pha, n_fft, hop_size, win_size, compress_factor=1.0, center=True):
+    # Magnitude Decompression
+    mag = torch.pow(mag, (1.0/compress_factor))
+    com = torch.complex(mag*torch.cos(pha), mag*torch.sin(pha))
+    hann_window = torch.hann_window(win_size).to(com.device)
+    wav = torch.istft(com, n_fft, hop_length=hop_size, win_length=win_size, window=hann_window, center=center)
+
+    return wav
+
+
+class LearnableSigmoid1d(nn.Module):
+    def __init__(self, in_features, beta=1):
+        super().__init__()
+        self.beta = beta
+        self.slope = nn.Parameter(torch.ones(in_features))
+        self.slope.requiresGrad = True
+
+    def forward(self, x):
+        # x shape: [batch_size, time_steps, spec_bins]
+        return self.beta * torch.sigmoid(self.slope * x)
+
+
+class LearnableSigmoid2d(nn.Module):
+    def __init__(self, in_features, beta=1):
+        super().__init__()
+        self.beta = beta
+        self.slope = nn.Parameter(torch.ones(in_features, 1))
+        self.slope.requiresGrad = True
+
+    def forward(self, x):
+        return self.beta * torch.sigmoid(self.slope * x)
+
+
+def main():
+    learnable_sigmoid = LearnableSigmoid1d(201)
+    a = torch.randn(4, 100, 201)
+
+    result = learnable_sigmoid.forward(a)
+    print(result.shape)
+
+    return
+
+
+if __name__ == '__main__':
+    main()

toolbox/torchaudio/models/mpnet/yaml/config.yaml

@@ -0,0 +1,27 @@
+model_name: "mpnet"
+
+num_gpus: 0
+batch_size: 4
+learning_rate: 0.0005
+adam_b1: 0.8
+adam_b2: 0.99
+lr_decay: 0.99
+seed: 1234
+
+dense_channel: 64
+compress_factor: 0.3
+num_tsconformers: 4
+beta: 2.0
+
+sample_rate: 16000
+segment_size: 32000
+n_fft: 400
+hop_size: 100
+win_size: 400
+
+num_workers: 4
+
+dist_config:
+  dist_backend: nccl
+  dist_url: tcp://localhost:54321
+  world_size: 1