update

examples/nx_denoise/run.sh

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

examples/nx_denoise/step_1_prepare_data.py

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

examples/nx_denoise/step_2_train_model.py

@@ -0,0 +1,439 @@
+#!/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.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.nx_denoise.configuration_nx_denoise import NXDenoiseConfig
+from toolbox.torchaudio.models.nx_denoise.discriminator import MetricDiscriminator, MetricDiscriminatorPretrainedModel
+from toolbox.torchaudio.models.nx_denoise.modeling_nx_denoise import NXDenoise, NXDenoisePretrainedModel
+from toolbox.torchaudio.models.nx_denoise.metrics import run_batch_pesq, run_pesq_score
+from toolbox.torchaudio.models.nx_denoise.utils import mag_pha_stft, mag_pha_istft
+from toolbox.torchaudio.models.nx_denoise.loss import phase_losses
+
+
+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 = NXDenoiseConfig.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=16,
+    )
+    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=16,
+    )
+
+    # models
+    logger.info(f"prepare models. config_file: {args.config_file}")
+    generator = NXDenoisePretrainedModel(config).to(device)
+    discriminator = MetricDiscriminatorPretrainedModel(config).to(device)
+
+    # optimizer
+    logger.info("prepare optimizer, lr_scheduler")
+    num_params = 0
+    for p in generator.parameters():
+        num_params += p.numel()
+    logger.info("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])
+
+    # 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
+
+    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_audios, noisy_audios = batch
+            clean_audios = clean_audios.to(device)
+            noisy_audios = noisy_audios.to(device)
+            one_labels = torch.ones(clean_audios.shape[0]).to(device)
+
+            audio_g = generator.forward(noisy_audios)
+
+            clean_mag, clean_pha, clean_com = mag_pha_stft(clean_audios, config.n_fft, config.hop_length, config.win_length, config.compress_factor)
+            mag_g, pha_g, com_g = mag_pha_stft(audio_g, config.n_fft, config.hop_length, config.win_length, config.compress_factor)
+
+            clean_audio_list = torch.split(clean_audios, 1, dim=0)
+            enhanced_audio_list = torch.split(audio_g, 1, dim=0)
+            clean_audio_list = [t.squeeze().detach().cpu().numpy() for t in clean_audio_list]
+            enhanced_audio_list = [t.squeeze().detach().cpu().numpy() for t in enhanced_audio_list]
+
+            pesq_score_list: List[float] = run_batch_pesq(clean_audio_list, enhanced_audio_list, sample_rate=config.sample_rate, mode="nb")
+
+            # Discriminator
+            optim_d.zero_grad()
+            metric_r = discriminator.forward(clean_audios, clean_audios)
+            metric_g = discriminator.forward(clean_audios, audio_g.detach())
+            loss_disc_r = F.mse_loss(one_labels, metric_r.flatten())
+
+            if -1 in pesq_score_list:
+                # print("-1 in batch_pesq_score!")
+                loss_disc_g = 0
+            else:
+                pesq_score_list: torch.FloatTensor = torch.tensor([(score - 1) / 3.5 for score in pesq_score_list], dtype=torch.float32)
+                loss_disc_g = F.mse_loss(pesq_score_list.to(device), metric_g.flatten())
+
+            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
+            # Time Loss
+            loss_time = F.l1_loss(clean_audios, audio_g)
+            # Metric Loss
+            metric_g = discriminator.forward(clean_audios, audio_g.detach())
+            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_metric * 0.05 + loss_time * 0.2
+            loss_gen_all = loss_mag * 0.1 + loss_pha * 0.1  + loss_com * 0.1 + loss_metric * 0.9 + loss_time * 0.9
+
+            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_batches = 0.
+
+        progress_bar = tqdm(
+            total=len(valid_data_loader),
+            desc="Evaluation; epoch: {}".format(idx_epoch),
+        )
+        with torch.no_grad():
+            for batch in valid_data_loader:
+                clean_audios, noisy_audios = batch
+                clean_audios = clean_audios.to(device)
+                noisy_audios = noisy_audios.to(device)
+
+                audio_g = generator.forward(noisy_audios)
+
+                clean_mag, clean_pha, clean_com = mag_pha_stft(clean_audios, config.n_fft, config.hop_length, config.win_length, config.compress_factor)
+                mag_g, pha_g, com_g = mag_pha_stft(audio_g, config.n_fft, config.hop_length, config.win_length, config.compress_factor)
+
+                clean_audio_list = torch.split(clean_audios, 1, dim=0)
+                enhanced_audio_list = torch.split(audio_g, 1, dim=0)
+                clean_audio_list = [t.squeeze().cpu().numpy() for t in clean_audio_list]
+                enhanced_audio_list = [t.squeeze().cpu().numpy() for t in enhanced_audio_list]
+                pesq_score = run_pesq_score(
+                    clean_audio_list,
+                    enhanced_audio_list,
+                    sample_rate = config.sample_rate,
+                    mode = "nb",
+                )
+                total_pesq_score += pesq_score
+                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_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)
+
+                progress_bar.update(1)
+                progress_bar.set_postfix({
+                    "pesq_metric": pesq_metric,
+                    "mag_err": mag_err,
+                    "pha_err": pha_err,
+                    "com_err": com_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,
+
+        }
+        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/nx_denoise/step_3_evaluation.py

@@ -0,0 +1,55 @@
+#!/usr/bin/python3
+# -*- coding: utf-8 -*-
+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
+
+
+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 main():
+    return
+
+
+if __name__ == '__main__':
+    main()

examples/nx_denoise/yaml/config.yaml

@@ -0,0 +1,42 @@
+model_name: "nx_clean_unet"
+
+sample_rate: 8000
+segment_size: 16000
+n_fft: 512
+win_size: 200
+hop_size: 80
+
+down_sampling_num_layers: 6
+down_sampling_in_channels: 1
+down_sampling_hidden_channels: 64
+down_sampling_kernel_size: 4
+down_sampling_stride: 2
+
+causal_in_channels: 1
+causal_out_channels: 1
+causal_kernel_size: 3
+causal_bias: false
+causal_separable: true
+causal_f_stride: 1
+causal_num_layers: 5
+
+tsfm_hidden_size: 256
+tsfm_attention_heads: 8
+tsfm_num_blocks: 6
+tsfm_dropout_rate: 0.1
+tsfm_max_length: 512
+tsfm_chunk_size: 1
+tsfm_num_left_chunks: 128
+tsfm_num_right_chunks: 4
+
+discriminator_dim: 32
+discriminator_in_channel: 2
+
+compress_factor: 0.3
+
+batch_size: 64
+learning_rate: 0.0005
+adam_b1: 0.8
+adam_b2: 0.99
+lr_decay: 0.99
+seed: 1234

toolbox/torchaudio/models/nx_clean_unet/modeling_nx_clean_unet.py

@@ -143,7 +143,7 @@ class UpSampling(nn.Module):
         for idx, up_sampling_block in enumerate(self.up_sampling_block_list):
             skip_x = skip_connection_list[idx]
             x = x + skip_x
-            # x = x + skip_x[:, :, :x.shape[-1]]
+            # x = x + skip_x[:, :, :x.size(2)]
             x = up_sampling_block.forward(x)
         return x
 

toolbox/torchaudio/models/nx_clean_unet/transformers/attention.py

@@ -207,6 +207,13 @@ class RelativeMultiHeadSelfAttention(nn.Module):
                 mask: torch.Tensor = None,
                 cache: torch.Tensor = None
                 ) -> Tuple[torch.Tensor, torch.Tensor]:
+        """
+
+        :param x:
+        :param mask:
+        :param cache: Tensor, shape: [1, n_heads, time_steps, dim]
+        :return:
+        """
         # attention! self attention.
 
         q, k, v = self.forward_qkv(x, x, x)

toolbox/torchaudio/models/nx_denoise/__init__.py

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

toolbox/torchaudio/models/nx_denoise/causal_convolution/__init__.py

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

toolbox/torchaudio/models/nx_denoise/causal_convolution/causal_conv2d.py

@@ -0,0 +1,281 @@
+#!/usr/bin/python3
+# -*- coding: utf-8 -*-
+import math
+import os
+from typing import List, Optional, Union, Iterable
+
+import numpy as np
+import torch
+import torch.nn as nn
+from torch.nn import functional as F
+
+
+norm_layer_dict = {
+    "batch_norm_2d": torch.nn.BatchNorm2d
+}
+
+
+activation_layer_dict = {
+    "relu": torch.nn.ReLU,
+    "identity": torch.nn.Identity,
+    "sigmoid": torch.nn.Sigmoid,
+}
+
+
+class CausalConv2d(nn.Module):
+    def __init__(self,
+                 in_channels: int,
+                 out_channels: int,
+                 kernel_size: Union[int, Iterable[int]],
+                 f_stride: int = 1,
+                 dilation: int = 1,
+                 do_f_pad: bool = True,
+                 bias: bool = True,
+                 separable: bool = False,
+                 norm_layer: str = "batch_norm_2d",
+                 activation_layer: str = "relu",
+                 lookahead: int = 0
+                 ):
+        super(CausalConv2d, self).__init__()
+        kernel_size = (kernel_size, kernel_size) if isinstance(kernel_size, int) else tuple(kernel_size)
+
+        if do_f_pad:
+            f_pad = kernel_size[1] // 2 + dilation - 1
+        else:
+            f_pad = 0
+
+        self.causal_left_pad = kernel_size[0] - 1 - lookahead
+        self.causal_right_pad = lookahead
+        self.constant_pad = nn.ConstantPad2d(
+            padding=(0, 0, self.causal_left_pad, self.causal_right_pad),
+            value=0.0
+        )
+
+        groups = math.gcd(in_channels, out_channels) if separable else 1
+        self.conv1 = nn.Conv2d(
+            in_channels,
+            out_channels,
+            kernel_size=kernel_size,
+            padding=(0, f_pad),
+            stride=(1, f_stride),
+            dilation=(1, dilation),
+            groups=groups,
+            bias=bias,
+        )
+
+        self.conv2 = None
+        if not any([groups == 1, max(kernel_size) == 1]):
+            self.conv2 = nn.Conv2d(
+                out_channels,
+                out_channels,
+                kernel_size=1,
+                bias=False,
+            )
+
+        self.norm = None
+        if norm_layer is not None:
+            norm_layer = norm_layer_dict[norm_layer]
+            self.norm = norm_layer(out_channels)
+
+        self.activation = None
+        if activation_layer is not None:
+            activation_layer = activation_layer_dict[activation_layer]
+            self.activation = activation_layer()
+
+    def forward(self,
+                inputs: torch.Tensor,
+                causal_cache: torch.Tensor = None,
+                ):
+
+        if causal_cache is None:
+            # inputs shape: [batch_size, 1, time_steps, hidden_size]
+            x = self.constant_pad.forward(inputs)
+        else:
+            # inputs shape: [batch_size, 1, time_steps + self.causal_right_pad, hidden_size]
+            # causal_cache shape: [batch_size, 1, self.causal_left_pad, hidden_size]
+            x = torch.concat(tensors=[causal_cache, inputs], dim=2)
+        # x shape: [batch_size, 1, time_steps2, hidden_size]
+        # time_steps2 = time_steps + self.causal_left_pad + self.causal_right_pad
+
+        x = self.conv1.forward(x)
+        # inputs shape: [batch_size, 1, time_steps, hidden_size]
+
+        if self.conv2:
+            x = self.conv2.forward(x)
+
+        if self.norm:
+            x = self.norm(x)
+        if self.activation:
+            x = self.activation(x)
+
+        causal_cache = x[:, :, -self.causal_left_pad:, :]
+
+        # inputs shape: [batch_size, 1, time_steps, hidden_size]
+        return x, causal_cache
+
+
+class CausalConv2dEncoder(nn.Module):
+    def __init__(self,
+                 in_channels: int,
+                 hidden_channels: int,
+                 out_channels: int,
+                 kernel_size: Union[int, Iterable[int]],
+                 f_stride: int = 1,
+                 dilation: int = 1,
+                 do_f_pad: bool = True,
+                 bias: bool = True,
+                 separable: bool = False,
+                 norm_layer: str = "batch_norm_2d",
+                 activation_layer: str = "relu",
+                 lookahead: int = 0,
+                 num_layers: int = 5,
+                 ):
+        super(CausalConv2dEncoder, self).__init__()
+        self.num_layers = num_layers
+
+        self.total_causal_left_pad = 0
+        self.total_causal_right_pad = 0
+
+        self.causal_conv_list: List[CausalConv2d] = nn.ModuleList(modules=[])
+        for i_layer in range(num_layers):
+            conv = CausalConv2d(
+                in_channels=in_channels,
+                out_channels=hidden_channels,
+                kernel_size=kernel_size,
+                f_stride=f_stride,
+                dilation=dilation,
+                do_f_pad=do_f_pad,
+                bias=bias,
+                separable=separable,
+                norm_layer=norm_layer,
+                activation_layer=activation_layer,
+                lookahead=lookahead,
+            )
+            self.causal_conv_list.append(conv)
+
+            self.total_causal_left_pad += conv.causal_left_pad
+            self.total_causal_right_pad += conv.causal_right_pad
+
+            in_channels = hidden_channels
+        else:
+            conv = CausalConv2d(
+                in_channels=hidden_channels,
+                out_channels=out_channels,
+                kernel_size=kernel_size,
+                f_stride=f_stride,
+                dilation=dilation,
+                do_f_pad=do_f_pad,
+                bias=bias,
+                separable=separable,
+                norm_layer=norm_layer,
+                activation_layer=activation_layer,
+                lookahead=lookahead,
+            )
+            self.causal_conv_list.append(conv)
+
+            self.total_causal_left_pad += conv.causal_left_pad
+            self.total_causal_right_pad += conv.causal_right_pad
+
+
+    def forward(self, inputs: torch.Tensor):
+        # inputs shape: [batch_size, 1, time_steps, hidden_size]
+
+        x = inputs
+        for layer in self.causal_conv_list:
+            x, _ = layer.forward(x)
+        return x
+
+    def forward_chunk(self,
+                      chunk: torch.Tensor,
+                      causal_cache: torch.Tensor = None,
+                      ):
+        # causal_cache shape: [self.num_layers, 1, causal_left_pad, hidden_size]
+
+        new_causal_cache_list = list()
+        for idx, causal_conv in enumerate(self.causal_conv_list):
+            chunk, new_causal_cache = causal_conv.forward(
+                inputs=chunk, causal_cache=causal_cache[idx: idx+1] if causal_cache is not None else None
+            )
+            new_causal_cache_list.append(new_causal_cache)
+
+        new_causal_cache = torch.cat(new_causal_cache_list, dim=0)
+        return chunk, new_causal_cache
+
+    def forward_chunk_by_chunk(self, inputs: torch.Tensor):
+        # inputs shape: [batch_size, 1, time_steps, hidden_size]
+        # batch_size = 1
+
+        batch_size, channels, time_steps, hidden_size = inputs.shape
+
+        causal_cache = None
+
+        outputs = []
+        for idx in range(0, time_steps, 1):
+            begin = idx
+            end = begin + self.total_causal_right_pad + 1
+            chunk_xs = inputs[:, :, begin:end, :]
+
+            ys, attention_cache = self.forward_chunk(
+                chunk=chunk_xs,
+                causal_cache=causal_cache,
+            )
+            # ys shape: [batch_size, channels, self.total_causal_right_pad + 1 , hidden_size]
+            ys = ys[:, :, :1, :]
+
+            # ys shape: [batch_size, chunk_size, hidden_size]
+            outputs.append(ys)
+
+        ys = torch.cat(outputs, 2)
+        return ys
+
+
+def main2():
+    conv = CausalConv2d(
+        in_channels=1,
+        out_channels=64,
+        kernel_size=3,
+        bias=False,
+        separable=True,
+        f_stride=1,
+        lookahead=0,
+    )
+
+    spec = torch.randn(size=(1, 1, 200, 64), dtype=torch.float32)
+    # spec shape: [batch_size, 1, time_steps, hidden_size]
+    cache = torch.randn(size=(1, 1, conv.causal_left_pad, 64), dtype=torch.float32)
+
+    output, _ = conv.forward(spec)
+    print(output.shape)
+
+    output, _ = conv.forward(spec, cache)
+    print(output.shape)
+
+    return
+
+
+def main():
+    causal = CausalConv2dEncoder(
+        in_channels=1,
+        out_channels=1,
+        kernel_size=3,
+        bias=False,
+        separable=True,
+        f_stride=1,
+        lookahead=0,
+        num_layers=3,
+    )
+
+    spec = torch.randn(size=(1, 1, 200, 64), dtype=torch.float32)
+    # spec shape: [batch_size, 1, time_steps, hidden_size]
+
+    output = causal.forward(spec)
+    print(output.shape)
+
+    output = causal.forward_chunk_by_chunk(spec)
+    print(output.shape)
+
+    return
+
+
+if __name__ == '__main__':
+    main()

toolbox/torchaudio/models/nx_denoise/configuration_nx_denoise.py

@@ -0,0 +1,102 @@
+#!/usr/bin/python3
+# -*- coding: utf-8 -*-
+from toolbox.torchaudio.configuration_utils import PretrainedConfig
+
+
+class NXDenoiseConfig(PretrainedConfig):
+    """
+    https://github.com/yxlu-0102/MP-SENet/blob/main/config.json
+    """
+    def __init__(self,
+                 sample_rate: int = 8000,
+                 segment_size: int = 16000,
+                 n_fft: int = 512,
+                 win_length: int = 200,
+                 hop_length: int = 80,
+
+                 down_sampling_num_layers: int = 5,
+                 down_sampling_in_channels: int = 1,
+                 down_sampling_hidden_channels: int = 64,
+                 down_sampling_kernel_size: int = 4,
+                 down_sampling_stride: int = 2,
+
+                 causal_in_channels: int = 1,
+                 causal_hidden_channels: int = 64,
+                 causal_kernel_size: int = 3,
+                 causal_bias: bool = False,
+                 causal_separable: bool = True,
+                 causal_f_stride: int = 1,
+                 # causal_lookahead: int = 0,
+                 causal_num_layers: int = 3,
+
+                 tsfm_hidden_size: int = 256,
+                 tsfm_attention_heads: int = 4,
+                 tsfm_num_blocks: int = 6,
+                 tsfm_dropout_rate: float = 0.1,
+                 tsfm_max_time_relative_position: int = 1024,
+                 tsfm_max_freq_relative_position: int = 128,
+                 tsfm_chunk_size: int = 4,
+                 tsfm_num_left_chunks: int = 128,
+                 tsfm_num_right_chunks: int = 2,
+
+                 discriminator_dim: int = 16,
+                 discriminator_in_channel: int = 2,
+
+                 compress_factor: float = 0.3,
+
+                 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,
+
+                 **kwargs
+                 ):
+        super(NXDenoiseConfig, self).__init__(**kwargs)
+        self.sample_rate = sample_rate
+        self.segment_size = segment_size
+        self.n_fft = n_fft
+        self.win_length = win_length
+        self.hop_length = hop_length
+
+        self.down_sampling_num_layers = down_sampling_num_layers
+        self.down_sampling_in_channels = down_sampling_in_channels
+        self.down_sampling_hidden_channels = down_sampling_hidden_channels
+        self.down_sampling_kernel_size = down_sampling_kernel_size
+        self.down_sampling_stride = down_sampling_stride
+
+        self.causal_in_channels = causal_in_channels
+        self.causal_hidden_channels = causal_hidden_channels
+        self.causal_kernel_size = causal_kernel_size
+        self.causal_bias = causal_bias
+        self.causal_separable = causal_separable
+        self.causal_f_stride = causal_f_stride
+        # self.causal_lookahead = causal_lookahead
+        self.causal_num_layers = causal_num_layers
+
+        self.tsfm_hidden_size = tsfm_hidden_size
+        self.tsfm_attention_heads = tsfm_attention_heads
+        self.tsfm_num_blocks = tsfm_num_blocks
+        self.tsfm_dropout_rate = tsfm_dropout_rate
+        self.tsfm_max_time_relative_position = tsfm_max_time_relative_position
+        self.tsfm_max_freq_relative_position = tsfm_max_freq_relative_position
+        self.tsfm_chunk_size = tsfm_chunk_size
+        self.tsfm_num_left_chunks = tsfm_num_left_chunks
+        self.tsfm_num_right_chunks = tsfm_num_right_chunks
+
+        self.discriminator_dim = discriminator_dim
+        self.discriminator_in_channel = discriminator_in_channel
+
+        self.compress_factor = compress_factor
+
+        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
+
+
+if __name__ == '__main__':
+    pass

toolbox/torchaudio/models/nx_denoise/discriminator.py

@@ -0,0 +1,132 @@
+#!/usr/bin/python3
+# -*- coding: utf-8 -*-
+import os
+from typing import Optional, Union
+
+import torch
+import torch.nn as nn
+import torchaudio
+
+from toolbox.torchaudio.configuration_utils import CONFIG_FILE
+from toolbox.torchaudio.models.nx_denoise.configuration_nx_denoise import NXDenoiseConfig
+from toolbox.torchaudio.models.nx_denoise.utils import LearnableSigmoid1d
+
+
+class MetricDiscriminator(nn.Module):
+    def __init__(self, config: NXDenoiseConfig):
+        super(MetricDiscriminator, self).__init__()
+        dim = config.discriminator_dim
+        self.in_channel = config.discriminator_in_channel
+
+        self.n_fft = config.n_fft
+        self.win_length = config.win_length
+        self.hop_length = config.hop_length
+
+        self.transform = torchaudio.transforms.Spectrogram(
+            n_fft=self.n_fft,
+            win_length=self.win_length,
+            hop_length=self.hop_length,
+            power=1.0,
+            window_fn=torch.hann_window,
+            # window_fn=torch.hamming_window if window_fn == "hamming" else torch.hann_window,
+        )
+
+        self.layers = nn.Sequential(
+            nn.utils.spectral_norm(nn.Conv2d(self.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):
+        x = self.transform.forward(x)
+        y = self.transform.forward(y)
+
+        xy = torch.stack((x, y), dim=1)
+        return self.layers(xy)
+
+
+MODEL_FILE = "discriminator.pt"
+
+
+class MetricDiscriminatorPretrainedModel(MetricDiscriminator):
+    def __init__(self,
+                 config: NXDenoiseConfig,
+                 ):
+        super(MetricDiscriminatorPretrainedModel, self).__init__(
+            config=config,
+        )
+        self.config = config
+
+    @classmethod
+    def from_pretrained(cls, pretrained_model_name_or_path, **kwargs):
+        config = NXDenoiseConfig.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 = NXDenoiseConfig()
+    discriminator = MetricDiscriminator(config=config)
+
+    # shape: [batch_size, num_samples]
+    # x = torch.ones([4, int(4.5 * 16000)])
+    # y = torch.ones([4, int(4.5 * 16000)])
+    x = torch.ones([4, 16000])
+    y = torch.ones([4, 16000])
+
+    output = discriminator.forward(x, y)
+    print(output.shape)
+    print(output)
+
+    return
+
+
+if __name__ == "__main__":
+    main()

toolbox/torchaudio/models/nx_denoise/loss.py

@@ -0,0 +1,22 @@
+#!/usr/bin/python3
+# -*- coding: utf-8 -*-
+import numpy as np
+import torch
+
+
+def anti_wrapping_function(x):
+
+    return torch.abs(x - torch.round(x / (2 * np.pi)) * 2 * np.pi)
+
+
+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
+
+
+if __name__ == '__main__':
+    pass

toolbox/torchaudio/models/nx_denoise/metrics.py

@@ -0,0 +1,80 @@
+#!/usr/bin/python3
+# -*- coding: utf-8 -*-
+from joblib import Parallel, delayed
+import numpy as np
+from pesq import pesq
+from typing import List
+
+from pesq import cypesq
+
+
+def run_pesq(clean_audio: np.ndarray,
+             noisy_audio: np.ndarray,
+             sample_rate: int = 16000,
+             mode: str = "wb",
+             ) -> float:
+    if sample_rate == 8000 and mode == "wb":
+        raise AssertionError(f"mode should be `nb` when sample_rate is 8000")
+    try:
+        pesq_score = pesq(sample_rate, clean_audio, noisy_audio, mode)
+    except cypesq.NoUtterancesError as e:
+        pesq_score = -1
+    except Exception as e:
+        print(f"pesq failed. error type: {type(e)}, error text: {str(e)}")
+        pesq_score = -1
+    return pesq_score
+
+
+def run_batch_pesq(clean_audio_list: List[np.ndarray],
+                   noisy_audio_list: List[np.ndarray],
+                   sample_rate: int = 16000,
+                   mode: str = "wb",
+                   n_jobs: int = 4,
+                   ) -> List[float]:
+    parallel = Parallel(n_jobs=n_jobs)
+
+    parallel_tasks = list()
+    for clean_audio, noisy_audio in zip(clean_audio_list, noisy_audio_list):
+        parallel_task = delayed(run_pesq)(clean_audio, noisy_audio, sample_rate, mode)
+        parallel_tasks.append(parallel_task)
+
+    pesq_score_list = parallel.__call__(parallel_tasks)
+    return pesq_score_list
+
+
+def run_pesq_score(clean_audio_list: List[np.ndarray],
+                   noisy_audio_list: List[np.ndarray],
+                   sample_rate: int = 16000,
+                   mode: str = "wb",
+                   n_jobs: int = 4,
+                   ) -> List[float]:
+
+    pesq_score_list = run_batch_pesq(clean_audio_list=clean_audio_list,
+                                     noisy_audio_list=noisy_audio_list,
+                                     sample_rate=sample_rate,
+                                     mode=mode,
+                                     n_jobs=n_jobs,
+                                     )
+
+    pesq_score = np.mean(pesq_score_list)
+    return pesq_score
+
+
+def main():
+    clean_audio = np.random.uniform(low=0, high=1, size=(2, 160000,))
+    noisy_audio = np.random.uniform(low=0, high=1, size=(2, 160000,))
+
+    clean_audio_list = list(clean_audio)
+    noisy_audio_list = list(noisy_audio)
+
+    pesq_score_list = run_batch_pesq(clean_audio_list, noisy_audio_list)
+    print(pesq_score_list)
+
+    pesq_score = run_pesq_score(clean_audio_list, noisy_audio_list)
+    print(pesq_score)
+
+    return
+
+
+if __name__ == "__main__":
+    main()

toolbox/torchaudio/models/nx_denoise/modeling_nx_denoise.py

@@ -0,0 +1,338 @@
+#!/usr/bin/python3
+# -*- coding: utf-8 -*-
+import os
+from typing import List, Optional, Union
+
+import numpy as np
+import torch
+import torch.nn as nn
+from torch.nn import functional as F
+
+from toolbox.torchaudio.configuration_utils import CONFIG_FILE
+from toolbox.torchaudio.models.nx_denoise.configuration_nx_denoise import NXDenoiseConfig
+from toolbox.torchaudio.models.nx_denoise.causal_convolution.causal_conv2d import CausalConv2dEncoder
+from toolbox.torchaudio.models.nx_denoise.transformers.transformers import TSTransformerEncoder
+
+
+class DownSamplingBlock(nn.Module):
+    def __init__(self,
+                 in_channels: int,
+                 hidden_channels: int,
+                 kernel_size: int,
+                 stride: int,
+                 ):
+        super(DownSamplingBlock, self).__init__()
+        self.conv1 = nn.Conv1d(in_channels, hidden_channels, kernel_size, stride)
+        self.relu = nn.ReLU()
+        self.conv2 = nn.Conv1d(hidden_channels, hidden_channels * 2, 1)
+        self.glu = nn.GLU(dim=1)
+
+    def forward(self, x: torch.Tensor):
+        # x shape: [batch_size, 1, num_samples]
+        x = self.conv1.forward(x)
+        # x shape: [batch_size, hidden_channels, new_num_samples]
+        x = self.relu(x)
+        x = self.conv2.forward(x)
+        # x shape: [batch_size, hidden_channels*2, new_num_samples]
+        x = self.glu(x)
+        # x shape: [batch_size, hidden_channels, new_num_samples]
+        # new_num_samples = (num_samples-kernel_size) // stride + 1
+        return x
+
+
+class DownSampling(nn.Module):
+    def __init__(self,
+                 num_layers: int,
+                 in_channels: int,
+                 hidden_channels: int,
+                 kernel_size: int,
+                 stride: int,
+                 ):
+        super(DownSampling, self).__init__()
+        self.num_layers = num_layers
+
+        down_sampling_block_list = list()
+        for idx in range(self.num_layers):
+            down_sampling_block = DownSamplingBlock(
+                in_channels=in_channels,
+                hidden_channels=hidden_channels,
+                kernel_size=kernel_size,
+                stride=stride,
+            )
+            down_sampling_block_list.append(down_sampling_block)
+            in_channels = hidden_channels
+
+        self.down_sampling_block_list = nn.ModuleList(modules=down_sampling_block_list)
+
+    def forward(self, x: torch.Tensor):
+        # x shape: [batch_size, channels, num_samples]
+        skip_connection_list = list()
+        for down_sampling_block in self.down_sampling_block_list:
+            x = down_sampling_block.forward(x)
+            skip_connection_list.append(x)
+        # x shape: [batch_size, hidden_channels, num_samples**]
+        return x, skip_connection_list
+
+
+class UpSamplingBlock(nn.Module):
+    def __init__(self,
+                 out_channels: int,
+                 hidden_channels: int,
+                 kernel_size: int,
+                 stride: int,
+                 do_relu: bool = True,
+                 ):
+        super(UpSamplingBlock, self).__init__()
+        self.do_relu = do_relu
+
+        self.conv1 = nn.Conv1d(hidden_channels, hidden_channels * 2, 1)
+        self.glu = nn.GLU(dim=1)
+        self.convt = nn.ConvTranspose1d(hidden_channels, out_channels, kernel_size, stride)
+        self.relu = nn.ReLU()
+
+    def forward(self, x: torch.Tensor):
+        # x shape: [batch_size, hidden_channels*2, num_samples]
+        x = self.conv1.forward(x)
+        # x shape: [batch_size, hidden_channels, num_samples]
+        x = self.glu(x)
+        # x shape: [batch_size, hidden_channels, num_samples]
+        x = self.convt.forward(x)
+        # x shape: [batch_size, hidden_channels, new_num_samples]
+        # new_num_samples = (num_samples - 1) * stride + kernel_size
+        if self.do_relu:
+            x = self.relu(x)
+        return x
+
+
+class UpSampling(nn.Module):
+    def __init__(self,
+                 num_layers: int,
+                 out_channels: int,
+                 hidden_channels: int,
+                 kernel_size: int,
+                 stride: int,
+                 ):
+        super(UpSampling, self).__init__()
+        self.num_layers = num_layers
+
+        up_sampling_block_list = list()
+        for idx in range(self.num_layers-1):
+            up_sampling_block = UpSamplingBlock(
+                out_channels=hidden_channels,
+                hidden_channels=hidden_channels,
+                kernel_size=kernel_size,
+                stride=stride,
+                do_relu=True,
+            )
+            up_sampling_block_list.append(up_sampling_block)
+        else:
+            up_sampling_block = UpSamplingBlock(
+                out_channels=out_channels,
+                hidden_channels=hidden_channels,
+                kernel_size=kernel_size,
+                stride=stride,
+                do_relu=False,
+            )
+            up_sampling_block_list.append(up_sampling_block)
+        self.up_sampling_block_list = nn.ModuleList(modules=up_sampling_block_list)
+
+    def forward(self, x: torch.Tensor, skip_connection_list: List[torch.Tensor]):
+        skip_connection_list = skip_connection_list[::-1]
+
+        # x shape: [batch_size, channels, num_samples]
+        for idx, up_sampling_block in enumerate(self.up_sampling_block_list):
+            skip_x = skip_connection_list[idx]
+            x = x + skip_x
+            # x = x + skip_x[:, :, :x.size(2)]
+            x = up_sampling_block.forward(x)
+        return x
+
+
+def get_padding_length(length, num_layers: int, kernel_size: int, stride: int):
+    for _ in range(num_layers):
+        if length < kernel_size:
+            length = 1
+        else:
+            length = 1 + np.ceil((length - kernel_size) / stride)
+
+    for _ in range(num_layers):
+        length = (length - 1) * stride + kernel_size
+
+    padded_length = int(length)
+    return padded_length
+
+
+class NXDenoise(nn.Module):
+    def __init__(self, config: NXDenoiseConfig):
+        super().__init__()
+        self.config = config
+
+        self.down_sampling = DownSampling(
+            num_layers=config.down_sampling_num_layers,
+            in_channels=config.down_sampling_in_channels,
+            hidden_channels=config.down_sampling_hidden_channels,
+            kernel_size=config.down_sampling_kernel_size,
+            stride=config.down_sampling_stride,
+        )
+        self.causal_conv_in = CausalConv2dEncoder(
+            in_channels=config.causal_in_channels,
+            hidden_channels=config.causal_hidden_channels,
+            out_channels=config.causal_hidden_channels,
+            kernel_size=config.causal_kernel_size,
+            bias=config.causal_bias,
+            separable=config.causal_separable,
+            f_stride=config.causal_f_stride,
+            lookahead=0,
+            num_layers=config.causal_num_layers,
+        )
+        self.ts_transformer = TSTransformerEncoder(
+            input_size=config.down_sampling_hidden_channels,
+            hidden_size=config.tsfm_hidden_size,
+            attention_heads=config.tsfm_attention_heads,
+            num_blocks=config.tsfm_num_blocks,
+            dropout_rate=config.tsfm_dropout_rate,
+            max_time_relative_position=config.tsfm_max_time_relative_position,
+            max_freq_relative_position=config.tsfm_max_freq_relative_position,
+            chunk_size=config.tsfm_chunk_size,
+            num_left_chunks=config.tsfm_num_left_chunks,
+            num_right_chunks=config.tsfm_num_right_chunks,
+        )
+        self.causal_conv_out = CausalConv2dEncoder(
+            in_channels=config.causal_hidden_channels,
+            hidden_channels=config.causal_hidden_channels,
+            out_channels=config.causal_in_channels,
+            kernel_size=config.causal_kernel_size,
+            bias=config.causal_bias,
+            separable=config.causal_separable,
+            f_stride=config.causal_f_stride,
+            lookahead=0,
+            num_layers=config.causal_num_layers,
+        )
+        self.up_sampling = UpSampling(
+            num_layers=config.down_sampling_num_layers,
+            out_channels=config.down_sampling_in_channels,
+            hidden_channels=config.down_sampling_hidden_channels,
+            kernel_size=config.down_sampling_kernel_size,
+            stride=config.down_sampling_stride,
+        )
+
+    def forward(self, noisy_audios: torch.Tensor):
+        # noisy_audios shape: [batch_size, n_samples]
+        noisy_audios = torch.unsqueeze(noisy_audios, dim=1)
+        # noisy_audios shape: [batch_size, 1, n_samples]
+
+        n_samples = noisy_audios.shape[-1]
+        padded_length = get_padding_length(
+            n_samples,
+            num_layers=self.config.down_sampling_num_layers,
+            kernel_size=self.config.down_sampling_kernel_size,
+            stride=self.config.down_sampling_stride,
+        )
+        noisy_audios_padded = F.pad(input=noisy_audios, pad=(0, padded_length - n_samples), mode="constant", value=0)
+
+        # down sampling
+        bottle_neck, skip_connection_list = self.down_sampling.forward(noisy_audios_padded)
+        # bottle_neck shape: [batch_size, channels, time_steps]
+        bottle_neck = torch.transpose(bottle_neck, dim0=-2, dim1=-1)
+        # bottle_neck shape: [batch_size, time_steps, channels]
+        bottle_neck = torch.unsqueeze(bottle_neck, dim=1)
+        # bottle_neck shape: [batch_size, 1, time_steps, freq_dim]
+
+        # causal conv in
+        bottle_neck = self.causal_conv_in.forward(bottle_neck)
+        # bottle_neck shape: [batch_size, channels, time_steps, freq_dim]
+
+        # ts transformer
+        # bottle_neck shape: [batch_size, channels, time_steps, freq_dim]
+
+        # causal conv out
+        bottle_neck = self.causal_conv_out.forward(bottle_neck)
+        # bottle_neck shape: [batch_size, 1, time_steps, freq_dim]
+
+        # up sampling
+        bottle_neck = torch.squeeze(bottle_neck, dim=1)
+        # bottle_neck shape: [batch_size, time_steps, channels]
+        bottle_neck = torch.transpose(bottle_neck, dim0=-2, dim1=-1)
+        # bottle_neck shape: [batch_size, channels, time_steps]
+
+        enhanced_audios = self.up_sampling.forward(bottle_neck, skip_connection_list)
+
+        enhanced_audios = enhanced_audios[:, :, :n_samples]
+        # enhanced_audios shape: [batch_size, 1, n_samples]
+
+        enhanced_audios = torch.squeeze(enhanced_audios, dim=1)
+        # enhanced_audios shape: [batch_size, n_samples]
+
+        return enhanced_audios
+
+
+MODEL_FILE = "generator.pt"
+
+
+class NXDenoisePretrainedModel(NXDenoise):
+    def __init__(self,
+                 config: NXDenoiseConfig,
+                 ):
+        super(NXDenoisePretrainedModel, self).__init__(
+            config=config,
+        )
+        self.config = config
+
+    @classmethod
+    def from_pretrained(cls, pretrained_model_name_or_path, **kwargs):
+        config = NXDenoiseConfig.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 = NXDenoiseConfig()
+
+    # shape: [batch_size, channels, num_samples]
+    # min length: 94, stride: 32, 32 == 2**5
+    # x = torch.ones([4, 94])
+    # x = torch.ones([4, 126])
+    # x = torch.ones([4, 158])
+    # x = torch.ones([4, 190])
+    x = torch.ones([4, 16000])
+
+    model = NXDenoise(config)
+    enhanced_audios = model.forward(x)
+    print(enhanced_audios.shape)
+    return
+
+
+if __name__ == "__main__":
+    main()

toolbox/torchaudio/models/nx_denoise/transformers/__init__.py

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

toolbox/torchaudio/models/nx_denoise/transformers/attention.py

@@ -0,0 +1,263 @@
+#!/usr/bin/python3
+# -*- coding: utf-8 -*-
+import math
+from typing import Tuple
+
+import torch
+import torch.nn as nn
+
+
+class MultiHeadSelfAttention(nn.Module):
+    def __init__(self, n_head: int, n_feat: int, dropout_rate: float):
+        """
+        :param n_head: int. the number of heads.
+        :param n_feat: int. the number of features.
+        :param dropout_rate: float. dropout rate.
+        """
+        super().__init__()
+        assert n_feat % n_head == 0
+        # We assume d_v always equals d_k
+        self.d_k = n_feat // n_head
+        self.h = n_head
+        self.linear_q = nn.Linear(n_feat, n_feat)
+        self.linear_k = nn.Linear(n_feat, n_feat)
+        self.linear_v = nn.Linear(n_feat, n_feat)
+        self.linear_out = nn.Linear(n_feat, n_feat)
+        self.dropout = nn.Dropout(p=dropout_rate)
+
+    def forward_qkv(self,
+                    query: torch.Tensor,
+                    key: torch.Tensor,
+                    value: torch.Tensor
+                    ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        """
+        transform query, key and value.
+        :param query: torch.Tensor. query tensor. shape=(batch_size, time1, n_feat).
+        :param key: torch.Tensor. key tensor. shape=(batch_size, time2, n_feat).
+        :param value: torch.Tensor. value tensor. shape=(batch_size, time2, n_feat).
+        :return:
+        """
+        n_batch = query.size(0)
+        q = self.linear_q(query).view(n_batch, -1, self.h, self.d_k)
+        k = self.linear_k(key).view(n_batch, -1, self.h, self.d_k)
+        v = self.linear_v(value).view(n_batch, -1, self.h, self.d_k)
+        q = q.transpose(1, 2)  # (batch, head, time1, d_k)
+        k = k.transpose(1, 2)  # (batch, head, time2, d_k)
+        v = v.transpose(1, 2)  # (batch, head, time2, d_k)
+
+        return q, k, v
+
+    def forward_attention(self,
+                          value: torch.Tensor,
+                          scores: torch.Tensor,
+                          mask: torch.Tensor = torch.ones((0, 0, 0), dtype=torch.bool)
+                          ) -> torch.Tensor:
+        """
+        compute attention context vector.
+        :param value: torch.Tensor. transformed value. shape=(batch_size, n_head, time2, d_k).
+        :param scores: torch.Tensor. attention score. shape=(batch_size, n_head, time1, time2).
+        :param mask: torch.Tensor. mask. shape=(batch_size, 1, time2) or
+                (batch_size, time1, time2), (0, 0, 0) means fake mask.
+        :return: torch.Tensor. transformed value. (batch_size, time1, d_model).
+                weighted by the attention score (batch_size, time1, time2).
+        """
+        n_batch = value.size(0)
+        # NOTE: When will `if mask.size(2) > 0` be True?
+        #   1. onnx(16/4) [WHY? Because we feed real cache & real mask for the
+        #           1st chunk to ease the onnx export.]
+        #   2. pytorch training
+        if mask.size(2) > 0:  # time2 > 0
+            mask = mask.unsqueeze(1).eq(0)  # (batch, 1, *, time2)
+            # For last chunk, time2 might be larger than scores.size(-1)
+            mask = mask[:, :, :, :scores.size(-1)]  # (batch, 1, *, time2)
+            scores = scores.masked_fill(mask, -float('inf'))
+            attn = torch.softmax(scores, dim=-1).masked_fill(mask, 0.0)  # (batch, head, time1, time2)
+
+        # NOTE: When will `if mask.size(2) > 0` be False?
+        #   1. onnx(16/-1, -1/-1, 16/0)
+        #   2. jit (16/-1, -1/-1, 16/0, 16/4)
+        else:
+            attn = torch.softmax(scores, dim=-1)  # (batch, head, time1, time2)
+
+        p_attn = self.dropout(attn)
+        x = torch.matmul(p_attn, value)  # (batch, head, time1, d_k)
+        x = x.transpose(1, 2).contiguous().view(n_batch, -1, self.h * self.d_k)  # (batch, time1, n_feat)
+
+        return self.linear_out(x)  # (batch, time1, n_feat)
+
+    def forward(self,
+                x: torch.Tensor,
+                mask: torch.Tensor = torch.ones((0, 0, 0), dtype=torch.bool),
+                cache: torch.Tensor = torch.zeros((0, 0, 0, 0))
+                ) -> Tuple[torch.Tensor, torch.Tensor]:
+
+        q, k, v = self.forward_qkv(x, x, x)
+
+        if cache.size(0) > 0:
+            key_cache, value_cache = torch.split(
+                cache, cache.size(-1) // 2, dim=-1)
+            k = torch.cat([key_cache, k], dim=2)
+            v = torch.cat([value_cache, v], dim=2)
+        # NOTE: We do cache slicing in encoder.forward_chunk, since it's
+        #   non-trivial to calculate `next_cache_start` here.
+        new_cache = torch.cat((k, v), dim=-1)
+
+        scores = torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(self.d_k)
+        return self.forward_attention(v, scores, mask), new_cache
+
+
+class RelativeMultiHeadSelfAttention(nn.Module):
+
+    def __init__(self, n_head: int, n_feat: int, dropout_rate: float, max_relative_position: int = 5120):
+        """
+        :param n_head: int. the number of heads.
+        :param n_feat: int. the number of features.
+        :param dropout_rate: float. dropout rate.
+        :param max_relative_position: int. maximum relative position for relative position encoding.
+        """
+        super().__init__()
+        assert n_feat % n_head == 0
+        # We assume d_v always equals d_k
+        self.d_k = n_feat // n_head
+        self.h = n_head
+        self.linear_q = nn.Linear(n_feat, n_feat)
+        self.linear_k = nn.Linear(n_feat, n_feat)
+        self.linear_v = nn.Linear(n_feat, n_feat)
+        self.linear_out = nn.Linear(n_feat, n_feat)
+        self.dropout = nn.Dropout(p=dropout_rate)
+
+        # Relative position encoding
+        self.max_relative_position = max_relative_position
+        self.relative_position_k = nn.Parameter(torch.randn(max_relative_position * 2 + 1, self.d_k))
+
+    def forward_qkv(self,
+                    query: torch.Tensor,
+                    key: torch.Tensor,
+                    value: torch.Tensor
+                    ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        """
+        transform query, key and value.
+        :param query: torch.Tensor. query tensor. shape=(batch_size, time1, n_feat).
+        :param key: torch.Tensor. key tensor. shape=(batch_size, time2, n_feat).
+        :param value: torch.Tensor. value tensor. shape=(batch_size, time2, n_feat).
+        :return:
+        """
+        n_batch = query.size(0)
+        q = self.linear_q(query).view(n_batch, -1, self.h, self.d_k)
+        k = self.linear_k(key).view(n_batch, -1, self.h, self.d_k)
+        v = self.linear_v(value).view(n_batch, -1, self.h, self.d_k)
+        q = q.transpose(1, 2)  # (batch, head, time1, d_k)
+        k = k.transpose(1, 2)  # (batch, head, time2, d_k)
+        v = v.transpose(1, 2)  # (batch, head, time2, d_k)
+
+        return q, k, v
+
+    def forward_attention(self,
+                          value: torch.Tensor,
+                          scores: torch.Tensor,
+                          mask: torch.Tensor = None
+                          ) -> torch.Tensor:
+        """
+        compute attention context vector.
+        :param value: torch.Tensor. transformed value. shape=(batch_size, n_head, key_time_steps, d_k).
+        :param scores: torch.Tensor. attention score. shape=(batch_size, n_head, query_time_steps, key_time_steps).
+        :param mask: torch.Tensor. mask. shape=(batch_size, 1, key_time_steps) or (batch_size, query_time_steps, key_time_steps).
+        :return: torch.Tensor. transformed value. (batch_size, query_time_steps, d_model).
+                weighted by the attention score (batch_size, query_time_steps, key_time_steps).
+        """
+        n_batch = value.size(0)
+        if mask is not None:
+            mask = mask.unsqueeze(1).eq(0)
+            # mask shape: [batch_size, 1, query_time_steps, key_time_steps]
+            scores = scores.masked_fill(mask, -float('inf'))
+            attn = torch.softmax(scores, dim=-1).masked_fill(mask, 0.0)
+        else:
+            attn = torch.softmax(scores, dim=-1)
+        # attn shape: [batch_size, n_head, query_time_steps, key_time_steps]
+
+        p_attn = self.dropout(attn)
+
+        x = torch.matmul(p_attn, value)
+        # x shape: [batch_size, n_head, query_time_steps, d_k]
+        x = x.transpose(1, 2)
+        # x shape: [batch_size, query_time_steps, n_head, d_k]
+
+        x = x.contiguous().view(n_batch, -1, self.h * self.d_k)  # (batch, time1, n_feat)
+        # x shape: [batch_size, query_time_steps, n_head * d_k]
+        # x shape: [batch_size, query_time_steps, n_feat]
+
+        x = self.linear_out(x)
+        # x shape: [batch_size, query_time_steps, n_feat]
+        return x
+
+    def relative_position_encoding(self, length: int) -> torch.Tensor:
+        """
+        Generate relative position encoding.
+        :param length: int. length of the sequence.
+        :return: torch.Tensor. relative position encoding. shape=(length, length, d_k).
+        """
+        range_vec = torch.arange(length)
+        distance_mat = range_vec.unsqueeze(0) - range_vec.unsqueeze(1)
+        distance_mat_clipped = torch.clamp(distance_mat, -self.max_relative_position, self.max_relative_position)
+        final_mat = distance_mat_clipped + self.max_relative_position
+        return final_mat
+
+    def forward(self,
+                x: torch.Tensor,
+                mask: torch.Tensor = None,
+                cache: torch.Tensor = None
+                ) -> Tuple[torch.Tensor, torch.Tensor]:
+        # attention! self attention.
+
+        q, k, v = self.forward_qkv(x, x, x)
+        # q k v shape: [batch_size, self.h, query_time_steps, self.d_k]
+
+        if cache is not None:
+            key_cache, value_cache = torch.split(
+                cache, cache.size(-1) // 2, dim=-1)
+            k = torch.cat([key_cache, k], dim=2)
+            v = torch.cat([value_cache, v], dim=2)
+
+        # new_cache shape: [batch_size, self.h, time_steps, self.d_k * 2]
+        new_cache = torch.cat((k, v), dim=-1)
+
+        # native_scores shape: [batch_size, self.h, q_time_steps, k_time_steps]
+        native_scores = torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(self.d_k)
+
+        # Compute relative position encoding
+        q_length, k_length = q.size(2), k.size(2)
+        relative_position = self.relative_position_encoding(k_length)
+
+        relative_position = relative_position[-q_length:]
+
+        relative_position_k = self.relative_position_k[relative_position.view(-1)].view(q_length, k_length, -1)
+
+        relative_position_k = relative_position_k.unsqueeze(0).unsqueeze(0)  # (1, 1, q_length, k_length, d_k)
+        relative_position_k = relative_position_k.expand(q.size(0), q.size(1), -1, -1, -1)  # (batch, head, q_length, k_length, d_k)
+
+        relative_position_scores = torch.matmul(q.unsqueeze(3), relative_position_k.transpose(-2, -1)).squeeze(3) / math.sqrt(self.d_k)
+        # relative_position_scores shape: [batch_size, self.h, q_time_steps, k_time_steps]
+
+        # score
+        scores = native_scores + relative_position_scores
+
+        return self.forward_attention(v, scores, mask), new_cache
+
+
+def main():
+    rel_attention = RelativeMultiHeadSelfAttention(n_head=4, n_feat=256, dropout_rate=0.1)
+
+    x = torch.ones(size=(1, 200, 256), dtype=torch.float32)
+    xt, new_cache = rel_attention.forward(x, x, x)
+
+    # x = torch.ones(size=(1, 1, 256), dtype=torch.float32)
+    # cache = torch.ones(size=(1, 4, 199, 128), dtype=torch.float32)
+    # xt, new_cache = rel_attention.forward(x, x, x, cache=cache)
+
+    print(xt.shape)
+    print(new_cache.shape)
+    return
+
+
+if __name__ == '__main__':
+    main()

toolbox/torchaudio/models/nx_denoise/transformers/mask.py

@@ -0,0 +1,74 @@
+#!/usr/bin/python3
+# -*- coding: utf-8 -*-
+import torch
+
+
+def make_pad_mask(lengths: torch.Tensor,
+                  max_len: int = 0,
+                  ) -> torch.Tensor:
+    batch_size = lengths.size(0)
+    max_len = max_len if max_len > 0 else lengths.max().item()
+    seq_range = torch.arange(
+        0,
+        max_len,
+        dtype=torch.int64,
+        device=lengths.device
+    )
+    seq_range_expand = seq_range.unsqueeze(0).expand(batch_size, max_len)
+    seq_length_expand = lengths.unsqueeze(-1)
+    mask = seq_range_expand >= seq_length_expand
+    return mask
+
+
+
+def subsequent_chunk_mask(
+        size: int,
+        chunk_size: int,
+        num_left_chunks: int = -1,
+        num_right_chunks: int = 0,
+        device: torch.device = torch.device("cpu"),
+) -> torch.Tensor:
+    """
+    Create mask for subsequent steps (size, size) with chunk size,
+    this is for streaming encoder
+
+    Examples:
+    > subsequent_chunk_mask(4, 2)
+    [[1, 1, 0, 0],
+     [1, 1, 0, 0],
+     [1, 1, 1, 1],
+     [1, 1, 1, 1]]
+
+    :param size: int. size of mask.
+    :param chunk_size: int. size of chunk.
+    :param num_left_chunks: int. number of left chunks. <0: use full chunk. >=0 use num_left_chunks.
+    :param num_right_chunks: int. number of right chunks.
+    :param device: torch.device. "cpu" or "cuda" or torch.Tensor.device.
+    :return: torch.Tensor. mask
+    """
+
+    ret = torch.zeros(size, size, device=device, dtype=torch.bool)
+    for i in range(size):
+        if num_left_chunks < 0:
+            start = 0
+        else:
+            start = max((i // chunk_size - num_left_chunks) * chunk_size, 0)
+        ending = min((i // chunk_size + 1 + num_right_chunks) * chunk_size, size)
+        ret[i, start:ending] = True
+    return ret
+
+
+def main():
+    chunk_mask = subsequent_chunk_mask(size=8, chunk_size=2, num_left_chunks=2)
+    print(chunk_mask)
+
+    chunk_mask = subsequent_chunk_mask(size=8, chunk_size=2, num_left_chunks=2, num_right_chunks=1)
+    print(chunk_mask)
+
+    chunk_mask = subsequent_chunk_mask(size=9, chunk_size=2, num_left_chunks=2, num_right_chunks=1)
+    print(chunk_mask)
+    return
+
+
+if __name__ == '__main__':
+    main()

toolbox/torchaudio/models/nx_denoise/transformers/transformers.py

@@ -0,0 +1,479 @@
+#!/usr/bin/python3
+# -*- coding: utf-8 -*-
+from typing import Dict, Optional, Tuple, List, Union
+
+import torch
+import torch.nn as nn
+
+from toolbox.torchaudio.models.nx_clean_unet.transformers.mask import subsequent_chunk_mask
+from toolbox.torchaudio.models.nx_clean_unet.transformers.attention import MultiHeadSelfAttention, RelativeMultiHeadSelfAttention
+
+
+class PositionwiseFeedForward(nn.Module):
+    def __init__(self,
+                 input_dim: int,
+                 hidden_units: int,
+                 dropout_rate: float,
+                 activation: torch.nn.Module = torch.nn.ReLU()):
+        """
+        FeedForward are applied on each position of the sequence.
+        the output dim is same with the input dim.
+
+        :param input_dim: int. input dimension.
+        :param hidden_units: int. the number of hidden units.
+        :param dropout_rate: float. dropout rate.
+        :param activation: torch.nn.Module. activation function.
+        """
+        super(PositionwiseFeedForward, self).__init__()
+        self.w_1 = torch.nn.Linear(input_dim, hidden_units)
+        self.activation = activation
+        self.dropout = torch.nn.Dropout(dropout_rate)
+        self.w_2 = torch.nn.Linear(hidden_units, input_dim)
+
+    def forward(self, xs: torch.Tensor) -> torch.Tensor:
+        """
+        Forward function.
+        :param xs: torch.Tensor. input tensor. shape=(batch_size, max_length, dim).
+        :return: output tensor. shape=(batch_size, max_length, dim).
+        """
+        return self.w_2(self.dropout(self.activation(self.w_1(xs))))
+
+
+class TransformerBlock(nn.Module):
+    def __init__(self,
+                 input_dim: int,
+                 dropout_rate: float = 0.1,
+                 n_heads: int = 4,
+                 max_relative_position: int = 5120
+                 ):
+        super().__init__()
+        self.norm1 = nn.LayerNorm(input_dim, eps=1e-5)
+        self.attention = RelativeMultiHeadSelfAttention(
+            n_head=n_heads,
+            n_feat=input_dim,
+            dropout_rate=dropout_rate,
+            max_relative_position=max_relative_position,
+        )
+
+        self.dropout1 = nn.Dropout(dropout_rate)
+        self.norm2 = nn.LayerNorm(input_dim, eps=1e-5)
+        self.ffn = PositionwiseFeedForward(
+            input_dim=input_dim,
+            hidden_units=input_dim,
+            dropout_rate=dropout_rate
+        )
+        self.dropout2 = nn.Dropout(dropout_rate)
+        self.norm3 = nn.LayerNorm(input_dim, eps=1e-5)
+
+    def forward(
+            self,
+            x: torch.Tensor,
+            mask: torch.Tensor = None,
+            attention_cache: torch.Tensor = None,
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        """
+
+        :param x: torch.Tensor. shape=(batch_size, time, input_dim).
+        :param mask: torch.Tensor. mask tensor for the input. shape=(batch_size, time，time).
+        :param attention_cache: torch.Tensor. cache tensor of the KEY & VALUE
+                shape=(batch_size=1, head, cache_t1, d_k * 2), head * d_k == input_dim.
+        :return:
+                torch.Tensor: Output tensor (batch_size, time, input_dim).
+                torch.Tensor: att_cache tensor, (batch_size=1, head, cache_t1 + time, d_k * 2).
+        """
+        xt = self.norm1(x)
+
+        x_att, new_att_cache = self.attention.forward(
+            xt, mask=mask, cache=attention_cache
+        )
+        x = x + self.dropout1(xt)
+        xt = self.norm2(x)
+        xt = self.ffn.forward(xt)
+        x = x + self.dropout2(xt)
+
+        x = self.norm3(x)
+
+        return x, new_att_cache
+
+
+class TransformerEncoder(nn.Module):
+    """
+    https://github.com/wenet-e2e/wenet/blob/main/wenet/transformer/encoder.py#L364
+    """
+    def __init__(self,
+                 input_size: int = 64,
+                 hidden_size: int = 256,
+                 attention_heads: int = 4,
+                 num_blocks: int = 6,
+                 dropout_rate: float = 0.1,
+                 max_relative_position: int = 1024,
+                 chunk_size: int = 1,
+                 num_left_chunks: int = 128,
+                 num_right_chunks: int = 2,
+                 ):
+        super().__init__()
+        self.input_size = input_size
+        self.hidden_size = hidden_size
+
+        self.max_relative_position = max_relative_position
+        self.chunk_size = chunk_size
+        self.num_left_chunks = num_left_chunks
+        self.num_right_chunks = num_right_chunks
+
+        self.input_linear = nn.Linear(
+            in_features=self.input_size,
+            out_features=self.hidden_size,
+        )
+
+        self.encoder_layer_list = torch.nn.ModuleList([
+            TransformerBlock(
+                input_dim=hidden_size,
+                n_heads=attention_heads,
+                dropout_rate=dropout_rate,
+                max_relative_position=max_relative_position,
+            ) for _ in range(num_blocks)
+        ])
+
+        self.output_linear = nn.Linear(
+            in_features=self.hidden_size,
+            out_features=self.input_size,
+        )
+
+    def forward(self,
+                xs: torch.Tensor,
+                ):
+        """
+        :param xs: Tensor, shape: [batch_size, time_steps, input_size]
+        :return: Tensor, shape: [batch_size, time_steps, input_size]
+        """
+        batch_size, time_steps, _ = xs.shape
+        # xs shape: [batch_size, time_steps, input_size]
+        xs = self.input_linear.forward(xs)
+        # xs shape: [batch_size, time_steps, hidden_size]
+
+        chunk_masks = subsequent_chunk_mask(
+            size=time_steps,
+            chunk_size=self.chunk_size,
+            num_left_chunks=self.num_left_chunks,
+            num_right_chunks=self.num_right_chunks,
+        )
+        chunk_masks = chunk_masks.to(xs.device)
+        # chunk_masks shape: [time_steps, time_steps]
+        chunk_masks = torch.broadcast_to(chunk_masks, size=(batch_size, time_steps, time_steps))
+        # chunk_masks shape: [batch_size, time_steps, time_steps]
+
+        for encoder_layer in self.encoder_layer_list:
+            xs, _ = encoder_layer.forward(xs, chunk_masks)
+
+        # xs shape: [batch_size, time_steps, hidden_size]
+        xs = self.output_linear.forward(xs)
+        # xs shape: [batch_size, time_steps, input_size]
+
+        return xs
+
+    def forward_chunk(self,
+                      xs: torch.Tensor,
+                      max_att_cache_length: int,
+                      attention_cache: torch.Tensor = None,
+                      ) -> Tuple[torch.Tensor, torch.Tensor]:
+        """
+
+        :param xs:
+        :param max_att_cache_length:
+        :param attention_cache: Tensor, [num_layers, ...]
+        :return:
+        """
+        # xs shape: [batch_size, time_steps, input_size]
+        xs = self.input_linear.forward(xs)
+        # xs shape: [batch_size, time_steps, hidden_size]
+
+        r_att_cache = []
+        for idx, encoder_layer in enumerate(self.encoder_layer_list):
+            xs, new_att_cache = encoder_layer.forward(
+                x=xs, attention_cache=attention_cache[idx] if attention_cache is not None else None,
+            )
+            # new_att_cache shape: [batch_size, n_heads, time_steps, dim]
+            if new_att_cache.size(2) > max_att_cache_length:
+                begin = (self.num_left_chunks + self.num_right_chunks) * self.chunk_size
+                end = self.num_right_chunks * self.chunk_size
+                new_att_cache = new_att_cache[:, :, -begin:-end, :]
+            r_att_cache.append(new_att_cache)
+
+        r_att_cache = torch.stack(r_att_cache, dim=0)
+
+        # xs shape: [batch_size, time_steps, hidden_size]
+        xs = self.output_linear.forward(xs)
+        # xs shape: [batch_size, time_steps, input_size]
+
+        return xs, r_att_cache
+
+    def forward_chunk_by_chunk(
+            self,
+            xs: torch.Tensor,
+    ) -> torch.Tensor:
+
+        batch_size, time_steps, _ = xs.shape
+
+        # attention_cache shape: [num_blocks, attention_heads, self.num_left_chunks * self.chunk_size, n_heads * d_k * 2]
+        max_att_cache_length = (self.num_left_chunks + self.num_right_chunks) * self.chunk_size
+        attention_cache = None
+
+        outputs = []
+        for idx in range(0, time_steps, self.chunk_size):
+            begin = idx
+            end = begin + self.chunk_size * (self.num_right_chunks + 1)
+            chunk_xs = xs[:, begin:end, :]
+            # print(f"begin: {begin}, end: {end}, length: {chunk_xs.size(1)}")
+
+            ys, attention_cache = self.forward_chunk(
+                xs=chunk_xs,
+                max_att_cache_length=max_att_cache_length,
+                attention_cache=attention_cache,
+            )
+
+            # ys shape: [batch_size, self.chunk_size * (self.num_right_chunks + 1), input_size]
+            ys = ys[:, :self.chunk_size, :]
+
+            outputs.append(ys)
+
+        ys = torch.cat(outputs, 1)
+        return ys
+
+
+class TSTransformerBlock(nn.Module):
+    def __init__(self,
+                 input_dim: int,
+                 dropout_rate: float = 0.1,
+                 n_heads: int = 4,
+                 max_time_relative_position: int = 1024,
+                 max_freq_relative_position: int = 128,
+                 ):
+        super(TSTransformerBlock, self).__init__()
+        self.time_transformer = TransformerBlock(input_dim, dropout_rate, n_heads, max_time_relative_position)
+        self.freq_transformer = TransformerBlock(input_dim, dropout_rate, n_heads, max_freq_relative_position)
+
+    def forward(self,
+                x: torch.Tensor,
+                mask: torch.Tensor = None,
+                attention_cache: torch.Tensor = None,
+                ):
+        """
+
+        :param x: Tensor. shape: [batch_size, hidden_size, time_steps, input_size]
+        :param mask: Tensor. shape: [time_steps, time_steps]
+        :param attention_cache:
+        :return:
+        """
+        b, c, t, f = x.size()
+
+        mask = None if mask is None else torch.broadcast_to(mask, size=(b*f, t, t))
+
+        x = x.permute(0, 3, 2, 1).contiguous().view(b*f, t, c)
+        x_, new_att_cache = self.time_transformer.forward(x, mask, attention_cache)
+        x = x_ + x
+        x = x.view(b, f, t, c).permute(0, 2, 1, 3).contiguous().view(b*t, f, c)
+        x_, _ = self.freq_transformer.forward(x)
+        x = x_ + x
+        x = x.view(b, t, f, c).permute(0, 3, 1, 2)
+        return x, new_att_cache
+
+
+class TSTransformerEncoder(nn.Module):
+    def __init__(self,
+                 input_size: int = 64,
+                 hidden_size: int = 256,
+                 attention_heads: int = 4,
+                 num_blocks: int = 6,
+                 dropout_rate: float = 0.1,
+                 max_time_relative_position: int = 1024,
+                 max_freq_relative_position: int = 128,
+                 chunk_size: int = 1,
+                 num_left_chunks: int = 128,
+                 num_right_chunks: int = 2,
+                 ):
+        super().__init__()
+        self.input_size = input_size
+        self.hidden_size = hidden_size
+
+        self.max_time_relative_position = max_time_relative_position
+        self.max_freq_relative_position = max_freq_relative_position
+        self.chunk_size = chunk_size
+        self.num_left_chunks = num_left_chunks
+        self.num_right_chunks = num_right_chunks
+
+        self.input_linear = nn.Linear(
+            in_features=self.input_size,
+            out_features=self.hidden_size,
+        )
+
+        self.encoder_layer_list = torch.nn.ModuleList([
+            TSTransformerBlock(
+                input_dim=hidden_size,
+                n_heads=attention_heads,
+                dropout_rate=dropout_rate,
+                max_time_relative_position=max_time_relative_position,
+                max_freq_relative_position=max_freq_relative_position,
+            ) for _ in range(num_blocks)
+        ])
+
+        self.output_linear = nn.Linear(
+            in_features=self.hidden_size,
+            out_features=self.input_size,
+        )
+
+    def forward(self,
+                xs: torch.Tensor,
+                ):
+        """
+        :param xs: Tensor, shape: [batch_size, channels, time_steps, input_size]
+        :return: Tensor, shape: [batch_size, channels, time_steps, input_size]
+        """
+        batch_size, channels, time_steps, _ = xs.shape
+        # xs shape: [batch_size, channels, time_steps, input_size]
+        xs = xs.permute(0, 3, 2, 1)
+        # xs shape: [batch_size, input_size, time_steps, channels]
+        xs = self.input_linear.forward(xs)
+        # xs shape: [batch_size, input_size, time_steps, hidden_size]
+        xs = xs.permute(0, 3, 2, 1)
+        # xs shape: [batch_size, hidden_size, time_steps, input_size]
+
+        chunk_masks = subsequent_chunk_mask(
+            size=time_steps,
+            chunk_size=self.chunk_size,
+            num_left_chunks=self.num_left_chunks,
+            num_right_chunks=self.num_right_chunks,
+        )
+        chunk_masks = chunk_masks.to(xs.device)
+        # chunk_masks shape: [time_steps, time_steps]
+
+        for encoder_layer in self.encoder_layer_list:
+            xs, _ = encoder_layer.forward(xs, chunk_masks)
+        # xs shape: [batch_size, hidden_size, time_steps, input_size]
+        xs = xs.permute(0, 3, 2, 1)
+        # xs shape: [batch_size, input_size, time_steps, hidden_size]
+        xs = self.output_linear.forward(xs)
+        # xs shape: [batch_size, input_size, time_steps, channels]
+        xs = xs.permute(0, 3, 2, 1)
+        # xs shape: [batch_size, channels, time_steps, input_size]
+
+        return xs
+
+    def forward_chunk(self,
+                      xs: torch.Tensor,
+                      max_att_cache_length: int,
+                      attention_cache: torch.Tensor = None,
+                      ) -> Tuple[torch.Tensor, torch.Tensor]:
+        """
+
+        :param xs:
+        :param max_att_cache_length:
+        :param attention_cache: Tensor, shape: [num_layers, ...]
+        :return:
+        """
+        # xs shape: [batch_size, channels, time_steps, input_size]
+        xs = xs.permute(0, 3, 2, 1)
+        xs = self.input_linear.forward(xs)
+        xs = xs.permute(0, 3, 2, 1)
+        # xs shape: [batch_size, hidden_size, time_steps, input_size]
+
+        r_att_cache = []
+        for idx, encoder_layer in enumerate(self.encoder_layer_list):
+            xs, new_att_cache = encoder_layer.forward(
+                x=xs, attention_cache=attention_cache[idx] if attention_cache is not None else None,
+            )
+            # new_att_cache shape: [b*f, n_heads, time_steps, dim]
+            if new_att_cache.size(2) > max_att_cache_length:
+                begin = (self.num_left_chunks + self.num_right_chunks) * self.chunk_size
+                end = self.num_right_chunks * self.chunk_size
+                new_att_cache = new_att_cache[:, :, -begin:-end, :]
+            r_att_cache.append(new_att_cache)
+
+        r_att_cache = torch.stack(r_att_cache, dim=0)
+
+        # xs shape: [batch_size, hidden_size, time_steps, input_size]
+        xs = xs.permute(0, 3, 2, 1)
+        xs = self.output_linear.forward(xs)
+        xs = xs.permute(0, 3, 2, 1)
+        # xs shape: [batch_size, channels, time_steps, input_size]
+
+        return xs, r_att_cache
+
+    def forward_chunk_by_chunk(
+            self,
+            xs: torch.Tensor,
+    ) -> torch.Tensor:
+
+        batch_size, channels, time_steps, _ = xs.shape
+
+        max_att_cache_length = (self.num_left_chunks + self.num_right_chunks) * self.chunk_size
+        attention_cache = None
+
+        outputs = []
+        for idx in range(0, time_steps, self.chunk_size):
+            begin = idx
+            end = begin + self.chunk_size * (self.num_right_chunks + 1)
+            chunk_xs = xs[:, :, begin:end, :]
+            # chunk_xs shape: [batch_size, channels, self.chunk_size * (self.num_right_chunks + 1), input_size]
+
+            ys, attention_cache = self.forward_chunk(
+                xs=chunk_xs,
+                max_att_cache_length=max_att_cache_length,
+                attention_cache=attention_cache,
+            )
+            # ys shape: [batch_size, channels, self.chunk_size * (self.num_right_chunks + 1), input_size]
+            ys = ys[:, :, :self.chunk_size, :]
+
+            outputs.append(ys)
+
+        ys = torch.cat(outputs, dim=2)
+        return ys
+
+
+def main2():
+
+    encoder = TransformerEncoder(
+        input_size=64,
+        hidden_size=256,
+        attention_heads=4,
+        num_blocks=6,
+        dropout_rate=0.1,
+    )
+    print(encoder)
+
+    x = torch.ones([4, 200, 64])
+
+    x = torch.ones([4, 200, 64])
+    y = encoder.forward(xs=x)
+    print(y.shape)
+
+    x = torch.ones([4, 200, 64])
+    y = encoder.forward_chunk_by_chunk(xs=x)
+    print(y.shape)
+
+    return
+
+
+def main():
+
+    encoder = TSTransformerEncoder(
+        input_size=16,
+        hidden_size=64,
+        attention_heads=4,
+        num_blocks=4,
+        dropout_rate=0.1,
+    )
+    # print(encoder)
+
+    x = torch.ones([4, 16, 200, 32])
+    y = encoder.forward(xs=x)
+    print(y.shape)
+
+    x = torch.ones([4, 16, 200, 32])
+    y = encoder.forward_chunk_by_chunk(xs=x)
+    print(y.shape)
+
+    return
+
+
+if __name__ == '__main__':
+    main()

toolbox/torchaudio/models/nx_denoise/utils.py

@@ -0,0 +1,45 @@
+#!/usr/bin/python3
+# -*- coding: utf-8 -*-
+import torch
+import torch.nn as nn
+
+
+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)
+
+
+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
+
+
+if __name__ == '__main__':
+    pass

toolbox/torchaudio/models/nx_denoise/yaml/config.yaml

@@ -0,0 +1,51 @@
+model_name: "nx_clean_unet"
+
+sample_rate: 8000
+segment_size: 16000
+n_fft: 512
+win_size: 200
+hop_size: 80
+# 因为 hop_size 取 80，则相当于 stft 的时间步是 10ms 一步，所以降采样也考虑到差不多的分辨率。
+
+# 2**down_sampling_num_layers，
+# 例如 2**6=64 就意味着 64 个值在降采样之后是一个时间步，
+# 则一步是 64/sample_rate = 0.008秒。
+# 那么 tsfm_chunk_size=2 则为16ms，tsfm_chunk_size=4 则为32ms
+# 假设每次向左看1秒，向右看30ms，则：
+# tsfm_chunk_size=1，tsfm_num_left_chunks=128，tsfm_num_right_chunks=4
+# tsfm_chunk_size=2，tsfm_num_left_chunks=64，tsfm_num_right_chunks=2
+# tsfm_chunk_size=4，tsfm_num_left_chunks=32，tsfm_num_right_chunks=1
+down_sampling_num_layers: 6
+down_sampling_in_channels: 1
+down_sampling_hidden_channels: 64
+down_sampling_kernel_size: 4
+down_sampling_stride: 2
+
+causal_in_channels: 1
+causal_out_channels: 64
+causal_kernel_size: 3
+causal_bias: false
+causal_separable: true
+causal_f_stride: 1
+causal_num_layers: 3
+
+tsfm_hidden_size: 256
+tsfm_attention_heads: 8
+tsfm_num_blocks: 6
+tsfm_dropout_rate: 0.1
+tsfm_max_length: 512
+tsfm_chunk_size: 1
+tsfm_num_left_chunks: 128
+tsfm_num_right_chunks: 4
+
+discriminator_dim: 32
+discriminator_in_channel: 2
+
+compress_factor: 0.3
+
+batch_size: 4
+learning_rate: 0.0005
+adam_b1: 0.8
+adam_b2: 0.99
+lr_decay: 0.99
+seed: 1234