update

examples/spectrum_dfnet_aishell/run.sh

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

examples/spectrum_dfnet_aishell/step_1_prepare_data.py

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

examples/spectrum_dfnet_aishell/step_2_train_model.py

@@ -0,0 +1,459 @@
+#!/usr/bin/python3
+# -*- coding: utf-8 -*-
+"""
+https://github.com/WenzheLiu-Speech/awesome-speech-enhancement
+"""
+import argparse
+import json
+import logging
+from logging.handlers import TimedRotatingFileHandler
+import os
+import platform
+from pathlib import Path
+import random
+import sys
+import shutil
+from typing import List
+
+pwd = os.path.abspath(os.path.dirname(__file__))
+sys.path.append(os.path.join(pwd, "../../"))
+
+import numpy as np
+import torch
+import torch.nn as nn
+from torch.nn import functional as F
+from torch.utils.data.dataloader import DataLoader
+import torchaudio
+from tqdm import tqdm
+
+from toolbox.torch.utils.data.dataset.denoise_excel_dataset import DenoiseExcelDataset
+from toolbox.torchaudio.models.spectrum_dfnet.configuration_spectrum_dfnet import SpectrumDfNetConfig
+from toolbox.torchaudio.models.spectrum_dfnet.modeling_spectrum_dfnet import SpectrumDfNetPretrainedModel
+
+
+def get_args():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--train_dataset", default="train.xlsx", type=str)
+    parser.add_argument("--valid_dataset", default="valid.xlsx", type=str)
+
+    parser.add_argument("--max_epochs", default=100, type=int)
+
+    parser.add_argument("--batch_size", default=64, type=int)
+    parser.add_argument("--learning_rate", default=1e-4, type=float)
+    parser.add_argument("--num_serialized_models_to_keep", default=10, type=int)
+    parser.add_argument("--patience", default=5, type=int)
+    parser.add_argument("--serialization_dir", default="serialization_dir", type=str)
+    parser.add_argument("--seed", default=0, type=int)
+
+    parser.add_argument("--config_file", default="config.yaml", type=str)
+
+    args = parser.parse_args()
+    return args
+
+
+def logging_config(file_dir: str):
+    fmt = "%(asctime)s - %(name)s - %(levelname)s  %(filename)s:%(lineno)d >  %(message)s"
+
+    logging.basicConfig(format=fmt,
+                        datefmt="%m/%d/%Y %H:%M:%S",
+                        level=logging.INFO)
+    file_handler = TimedRotatingFileHandler(
+        filename=os.path.join(file_dir, "main.log"),
+        encoding="utf-8",
+        when="D",
+        interval=1,
+        backupCount=7
+    )
+    file_handler.setLevel(logging.INFO)
+    file_handler.setFormatter(logging.Formatter(fmt))
+    logger = logging.getLogger(__name__)
+    logger.addHandler(file_handler)
+
+    return logger
+
+
+class CollateFunction(object):
+    def __init__(self,
+                 n_fft: int = 512,
+                 win_length: int = 200,
+                 hop_length: int = 80,
+                 window_fn: str = "hamming",
+                 irm_beta: float = 1.0,
+                 epsilon: float = 1e-8,
+                 ):
+        self.n_fft = n_fft
+        self.win_length = win_length
+        self.hop_length = hop_length
+        self.window_fn = window_fn
+        self.irm_beta = irm_beta
+        self.epsilon = epsilon
+
+        self.complex_transform = torchaudio.transforms.Spectrogram(
+            n_fft=self.n_fft,
+            win_length=self.win_length,
+            hop_length=self.hop_length,
+            power=None,
+            window_fn=torch.hamming_window if window_fn == "hamming" else torch.hann_window,
+        )
+        self.transform = torchaudio.transforms.Spectrogram(
+            n_fft=self.n_fft,
+            win_length=self.win_length,
+            hop_length=self.hop_length,
+            power=2.0,
+            window_fn=torch.hamming_window if window_fn == "hamming" else torch.hann_window,
+        )
+
+    @staticmethod
+    def make_unfold_snr_db(x: torch.Tensor, n_time_steps: int = 3):
+        batch_size, channels, freq_dim, time_steps = x.shape
+
+        # kernel: [freq_dim, n_time_step]
+        kernel_size = (freq_dim, n_time_steps)
+
+        # pad
+        pad = n_time_steps // 2
+        x = torch.concat(tensors=[
+            x[:, :, :, :pad],
+            x,
+            x[:, :, :, -pad:],
+        ], dim=-1)
+
+        x = F.unfold(
+            input=x,
+            kernel_size=kernel_size,
+        )
+        # x shape: [batch_size, fold, time_steps]
+        return x
+
+    def __call__(self, batch: List[dict]):
+        speech_complex_spec_list = list()
+        mix_complex_spec_list = list()
+        speech_irm_list = list()
+        snr_db_list = list()
+        for sample in batch:
+            noise_wave: torch.Tensor = sample["noise_wave"]
+            speech_wave: torch.Tensor = sample["speech_wave"]
+            mix_wave: torch.Tensor = sample["mix_wave"]
+            # snr_db: float = sample["snr_db"]
+
+            noise_spec = self.transform.forward(noise_wave)
+            speech_spec = self.transform.forward(speech_wave)
+
+            speech_complex_spec = self.complex_transform.forward(speech_wave)
+            mix_complex_spec = self.complex_transform.forward(mix_wave)
+
+            # noise_irm = noise_spec / (noise_spec + speech_spec)
+            speech_irm = speech_spec / (noise_spec + speech_spec + self.epsilon)
+            speech_irm = torch.pow(speech_irm, self.irm_beta)
+
+            # noise_spec, speech_spec, mix_spec, speech_irm
+            # shape: [freq_dim, time_steps]
+
+            snr_db: torch.Tensor = 10 * torch.log10(
+                speech_spec / (noise_spec + self.epsilon)
+            )
+            snr_db = torch.clamp(snr_db, min=self.epsilon)
+
+            snr_db_ = torch.unsqueeze(snr_db, dim=0)
+            snr_db_ = torch.unsqueeze(snr_db_, dim=0)
+            snr_db_ = self.make_unfold_snr_db(snr_db_, n_time_steps=3)
+            snr_db_ = torch.squeeze(snr_db_, dim=0)
+            # snr_db_ shape: [fold, time_steps]
+
+            snr_db = torch.mean(snr_db_, dim=0, keepdim=True)
+            # snr_db shape: [1, time_steps]
+
+            speech_complex_spec_list.append(speech_complex_spec)
+            mix_complex_spec_list.append(mix_complex_spec)
+            speech_irm_list.append(speech_irm)
+            snr_db_list.append(snr_db)
+
+        speech_complex_spec_list = torch.stack(speech_complex_spec_list)
+        mix_complex_spec_list = torch.stack(mix_complex_spec_list)
+        speech_irm_list = torch.stack(speech_irm_list)
+        snr_db_list = torch.stack(snr_db_list)  # shape: (batch_size, time_steps, 1)
+
+        speech_complex_spec_list = speech_complex_spec_list[:, :-1, :]
+        mix_complex_spec_list = mix_complex_spec_list[:, :-1, :]
+        speech_irm_list = speech_irm_list[:, :-1, :]
+
+        # speech_complex_spec_list shape: [batch_size, freq_dim, time_steps]
+        # mix_complex_spec_list shape: [batch_size, freq_dim, time_steps]
+        # speech_irm_list shape: [batch_size, freq_dim, time_steps]
+        # snr_db shape: [batch_size, 1, time_steps]
+
+        # assert
+        if torch.any(torch.isnan(speech_complex_spec_list)) or torch.any(torch.isinf(speech_complex_spec_list)):
+            raise AssertionError("nan or inf in speech_complex_spec_list")
+        if torch.any(torch.isnan(mix_complex_spec_list)) or torch.any(torch.isinf(mix_complex_spec_list)):
+            raise AssertionError("nan or inf in mix_complex_spec_list")
+        if torch.any(torch.isnan(speech_irm_list)) or torch.any(torch.isinf(speech_irm_list)):
+            raise AssertionError("nan or inf in speech_irm_list")
+        if torch.any(torch.isnan(snr_db_list)) or torch.any(torch.isinf(snr_db_list)):
+            raise AssertionError("nan or inf in snr_db_list")
+
+        return speech_complex_spec_list, mix_complex_spec_list, speech_irm_list, snr_db_list
+
+
+collate_fn = CollateFunction()
+
+
+def main():
+    args = get_args()
+
+    serialization_dir = Path(args.serialization_dir)
+    serialization_dir.mkdir(parents=True, exist_ok=True)
+
+    logger = logging_config(serialization_dir)
+
+    random.seed(args.seed)
+    np.random.seed(args.seed)
+    torch.manual_seed(args.seed)
+    logger.info("set seed: {}".format(args.seed))
+
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    n_gpu = torch.cuda.device_count()
+    logger.info("GPU available count: {}; device: {}".format(n_gpu, device))
+
+    # datasets
+    logger.info("prepare datasets")
+    train_dataset = DenoiseExcelDataset(
+        excel_file=args.train_dataset,
+        expected_sample_rate=8000,
+        max_wave_value=32768.0,
+    )
+    valid_dataset = DenoiseExcelDataset(
+        excel_file=args.valid_dataset,
+        expected_sample_rate=8000,
+        max_wave_value=32768.0,
+    )
+    train_data_loader = DataLoader(
+        dataset=train_dataset,
+        batch_size=args.batch_size,
+        shuffle=True,
+        # Linux 系统中可以使用多个子进程加载数据, 而在 Windows 系统中不能.
+        num_workers=0 if platform.system() == "Windows" else os.cpu_count() // 2,
+        collate_fn=collate_fn,
+        pin_memory=False,
+        # prefetch_factor=64,
+    )
+    valid_data_loader = DataLoader(
+        dataset=valid_dataset,
+        batch_size=args.batch_size,
+        shuffle=True,
+        # Linux 系统中可以使用多个子进程加载数据, 而在 Windows 系统中不能.
+        num_workers=0 if platform.system() == "Windows" else os.cpu_count() // 2,
+        collate_fn=collate_fn,
+        pin_memory=False,
+        # prefetch_factor=64,
+    )
+
+    # models
+    logger.info(f"prepare models. config_file: {args.config_file}")
+    config = SpectrumDfNetConfig.from_pretrained(
+        pretrained_model_name_or_path=args.config_file,
+        # num_labels=vocabulary.get_vocab_size(namespace="labels")
+    )
+    model = SpectrumDfNetPretrainedModel(
+        config=config,
+    )
+    model.to(device)
+    model.train()
+
+    # optimizer
+    logger.info("prepare optimizer, lr_scheduler, loss_fn, categorical_accuracy")
+    param_optimizer = model.parameters()
+    optimizer = torch.optim.Adam(
+        param_optimizer,
+        lr=args.learning_rate,
+    )
+    # lr_scheduler = torch.optim.lr_scheduler.StepLR(
+    #     optimizer,
+    #     step_size=2000
+    # )
+    lr_scheduler = torch.optim.lr_scheduler.MultiStepLR(
+        optimizer,
+        milestones=[10000, 20000, 30000, 40000, 50000], gamma=0.5
+    )
+
+    speech_mse_loss = nn.MSELoss(
+        reduction="mean",
+    )
+    irm_mse_loss = nn.MSELoss(
+        reduction="mean",
+    )
+    snr_mse_loss = nn.MSELoss(
+        reduction="mean",
+    )
+
+    # training loop
+    logger.info("training")
+
+    training_loss = 10000000000
+    evaluation_loss = 10000000000
+
+    model_list = list()
+    best_idx_epoch = None
+    best_metric = None
+    patience_count = 0
+
+    for idx_epoch in range(args.max_epochs):
+        total_loss = 0.
+        total_examples = 0.
+        progress_bar = tqdm(
+            total=len(train_data_loader),
+            desc="Training; epoch: {}".format(idx_epoch),
+        )
+
+        for batch in train_data_loader:
+            speech_complex_spec, mix_complex_spec, speech_irm, snr_db = batch
+            speech_complex_spec = speech_complex_spec.to(device)
+            mix_complex_spec = mix_complex_spec.to(device)
+            speech_irm_target = speech_irm.to(device)
+            snr_db_target = snr_db.to(device)
+
+            speech_spec_prediction, speech_irm_prediction, lsnr_prediction = model.forward(mix_complex_spec)
+            if torch.any(torch.isnan(speech_spec_prediction)) or torch.any(torch.isinf(speech_spec_prediction)):
+                raise AssertionError("nan or inf in speech_spec_prediction")
+            if torch.any(torch.isnan(speech_irm_prediction)) or torch.any(torch.isinf(speech_irm_prediction)):
+                raise AssertionError("nan or inf in speech_irm_prediction")
+            if torch.any(torch.isnan(lsnr_prediction)) or torch.any(torch.isinf(lsnr_prediction)):
+                raise AssertionError("nan or inf in lsnr_prediction")
+
+            speech_loss = speech_mse_loss.forward(speech_spec_prediction, speech_complex_spec)
+            irm_loss = irm_mse_loss.forward(speech_irm_prediction, speech_irm_target)
+
+            lsnr_prediction = (lsnr_prediction - config.lsnr_min) / (config.lsnr_max - config.lsnr_min)
+            snr_db_target = (snr_db_target - config.lsnr_min) / (config.lsnr_max - config.lsnr_min)
+            if torch.max(lsnr_prediction) > 1 or torch.min(lsnr_prediction) < 0:
+                raise AssertionError(f"expected lsnr_prediction between 0 and 1.")
+            if torch.max(snr_db_target) > 1 or torch.min(snr_db_target) < 0:
+                raise AssertionError(f"expected snr_db_target between 0 and 1.")
+
+            snr_loss = snr_mse_loss.forward(lsnr_prediction, snr_db_target)
+
+            if torch.any(torch.isnan(snr_loss)) or torch.any(torch.isinf(snr_loss)):
+                raise AssertionError("nan or inf in snr_loss")
+
+            loss = speech_loss + irm_loss + snr_loss
+
+            total_loss += loss.item()
+            total_examples += mix_complex_spec.size(0)
+
+            optimizer.zero_grad()
+            loss.backward()
+            optimizer.step()
+            lr_scheduler.step()
+
+            training_loss = total_loss / total_examples
+            training_loss = round(training_loss, 4)
+
+            progress_bar.update(1)
+            progress_bar.set_postfix({
+                "training_loss": training_loss,
+            })
+
+        total_loss = 0.
+        total_examples = 0.
+        progress_bar = tqdm(
+            total=len(valid_data_loader),
+            desc="Evaluation; epoch: {}".format(idx_epoch),
+        )
+        for batch in valid_data_loader:
+            speech_complex_spec, mix_complex_spec, speech_irm, snr_db = batch
+            speech_complex_spec = speech_complex_spec.to(device)
+            mix_complex_spec = mix_complex_spec.to(device)
+            speech_irm_target = speech_irm.to(device)
+            snr_db_target = snr_db.to(device)
+
+            with torch.no_grad():
+                speech_spec_prediction, speech_irm_prediction, lsnr_prediction = model.forward(mix_complex_spec)
+                if torch.any(torch.isnan(speech_spec_prediction)) or torch.any(torch.isinf(speech_spec_prediction)):
+                    raise AssertionError("nan or inf in speech_spec_prediction")
+                if torch.any(torch.isnan(speech_irm_prediction)) or torch.any(torch.isinf(speech_irm_prediction)):
+                    raise AssertionError("nan or inf in speech_irm_prediction")
+                if torch.any(torch.isnan(lsnr_prediction)) or torch.any(torch.isinf(lsnr_prediction)):
+                    raise AssertionError("nan or inf in lsnr_prediction")
+
+                speech_loss = speech_mse_loss.forward(speech_spec_prediction, speech_complex_spec)
+                irm_loss = irm_mse_loss.forward(speech_irm_prediction, speech_irm_target)
+
+                lsnr_prediction = (lsnr_prediction - config.lsnr_min) / (config.lsnr_max - config.lsnr_min)
+                snr_db_target = (snr_db_target - config.lsnr_min) / (config.lsnr_max - config.lsnr_min)
+                if torch.max(lsnr_prediction) > 1 or torch.min(lsnr_prediction) < 0:
+                    raise AssertionError(f"expected lsnr_prediction between 0 and 1.")
+                if torch.max(snr_db_target) > 1 or torch.min(snr_db_target) < 0:
+                    raise AssertionError(f"expected snr_db_target between 0 and 1.")
+
+                snr_loss = snr_mse_loss.forward(lsnr_prediction, snr_db_target)
+
+                loss = speech_loss + irm_loss + snr_loss
+
+                total_loss += loss.item()
+            total_examples += mix_complex_spec.size(0)
+
+            evaluation_loss = total_loss / total_examples
+            evaluation_loss = round(evaluation_loss, 4)
+
+            progress_bar.update(1)
+            progress_bar.set_postfix({
+                "evaluation_loss": evaluation_loss,
+            })
+
+        # save path
+        epoch_dir = serialization_dir / "epoch-{}".format(idx_epoch)
+        epoch_dir.mkdir(parents=True, exist_ok=False)
+
+        # save models
+        model.save_pretrained(epoch_dir.as_posix())
+
+        model_list.append(epoch_dir)
+        if len(model_list) >= args.num_serialized_models_to_keep:
+            model_to_delete: Path = model_list.pop(0)
+            shutil.rmtree(model_to_delete.as_posix())
+
+        # save metric
+        if best_metric is None:
+            best_idx_epoch = idx_epoch
+            best_metric = evaluation_loss
+        elif evaluation_loss < best_metric:
+            best_idx_epoch = idx_epoch
+            best_metric = evaluation_loss
+        else:
+            pass
+
+        metrics = {
+            "idx_epoch": idx_epoch,
+            "best_idx_epoch": best_idx_epoch,
+            "training_loss": training_loss,
+            "evaluation_loss": evaluation_loss,
+            "learning_rate": optimizer.param_groups[0]["lr"],
+        }
+        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/spectrum_dfnet_aishell/step_3_evaluation.py

@@ -0,0 +1,270 @@
+#!/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
+
+from toolbox.torchaudio.models.spectrum_unet_irm.modeling_spectrum_unet_irm import SpectrumUnetIRMPretrainedModel
+
+
+def get_args():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--valid_dataset", default="valid.xlsx", type=str)
+    parser.add_argument("--model_dir", default="serialization_dir/best", type=str)
+    parser.add_argument("--evaluation_audio_dir", default="evaluation_audio_dir", type=str)
+
+    parser.add_argument("--limit", default=10, type=int)
+
+    args = parser.parse_args()
+    return args
+
+
+def logging_config():
+    fmt = "%(asctime)s - %(name)s - %(levelname)s  %(filename)s:%(lineno)d >  %(message)s"
+
+    logging.basicConfig(format=fmt,
+                        datefmt="%m/%d/%Y %H:%M:%S",
+                        level=logging.INFO)
+    stream_handler = logging.StreamHandler()
+    stream_handler.setLevel(logging.INFO)
+    stream_handler.setFormatter(logging.Formatter(fmt))
+
+    logger = logging.getLogger(__name__)
+
+    return logger
+
+
+def mix_speech_and_noise(speech: np.ndarray, noise: np.ndarray, snr_db: float):
+    l1 = len(speech)
+    l2 = len(noise)
+    l = min(l1, l2)
+    speech = speech[:l]
+    noise = noise[:l]
+
+    # np.float32, value between (-1, 1).
+
+    speech_power = np.mean(np.square(speech))
+    noise_power = speech_power / (10 ** (snr_db / 10))
+
+    noise_adjusted = np.sqrt(noise_power) * noise / np.sqrt(np.mean(noise ** 2))
+
+    noisy_signal = speech + noise_adjusted
+
+    return noisy_signal
+
+
+stft_power = torchaudio.transforms.Spectrogram(
+    n_fft=512,
+    win_length=200,
+    hop_length=80,
+    power=2.0,
+    window_fn=torch.hamming_window,
+)
+
+
+stft_complex = torchaudio.transforms.Spectrogram(
+    n_fft=512,
+    win_length=200,
+    hop_length=80,
+    power=None,
+    window_fn=torch.hamming_window,
+)
+
+
+istft = torchaudio.transforms.InverseSpectrogram(
+    n_fft=512,
+    win_length=200,
+    hop_length=80,
+    window_fn=torch.hamming_window,
+)
+
+
+def enhance(mix_spec_complex: torch.Tensor, speech_irm_prediction: torch.Tensor):
+    mix_spec_complex = mix_spec_complex.detach().cpu()
+    speech_irm_prediction = speech_irm_prediction.detach().cpu()
+
+    mask_speech = speech_irm_prediction
+    mask_noise = 1.0 - speech_irm_prediction
+
+    speech_spec = mix_spec_complex * mask_speech
+    noise_spec = mix_spec_complex * mask_noise
+
+    speech_wave = istft.forward(speech_spec)
+    noise_wave = istft.forward(noise_spec)
+
+    return speech_wave, noise_wave
+
+
+def save_audios(noise_wave: torch.Tensor,
+                speech_wave: torch.Tensor,
+                mix_wave: torch.Tensor,
+                speech_wave_enhanced: torch.Tensor,
+                noise_wave_enhanced: torch.Tensor,
+                output_dir: str,
+                sample_rate: int = 8000,
+                ):
+    basename = uuid.uuid4().__str__()
+    output_dir = Path(output_dir) / basename
+    output_dir.mkdir(parents=True, exist_ok=True)
+
+    filename = output_dir / "noise_wave.wav"
+    torchaudio.save(filename, noise_wave, sample_rate)
+    filename = output_dir / "speech_wave.wav"
+    torchaudio.save(filename, speech_wave, sample_rate)
+    filename = output_dir / "mix_wave.wav"
+    torchaudio.save(filename, mix_wave, sample_rate)
+
+    filename = output_dir / "speech_wave_enhanced.wav"
+    torchaudio.save(filename, speech_wave_enhanced, sample_rate)
+    filename = output_dir / "noise_wave_enhanced.wav"
+    torchaudio.save(filename, noise_wave_enhanced, sample_rate)
+
+    return output_dir.as_posix()
+
+
+def main():
+    args = get_args()
+
+    logger = logging_config()
+
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    n_gpu = torch.cuda.device_count()
+    logger.info("GPU available count: {}; device: {}".format(n_gpu, device))
+
+    logger.info("prepare model")
+    model = SpectrumUnetIRMPretrainedModel.from_pretrained(
+        pretrained_model_name_or_path=args.model_dir,
+    )
+    model.to(device)
+    model.eval()
+
+    # optimizer
+    logger.info("prepare loss_fn")
+    irm_mse_loss = nn.MSELoss(
+        reduction="mean",
+    )
+    snr_mse_loss = nn.MSELoss(
+        reduction="mean",
+    )
+
+    logger.info("read excel")
+    df = pd.read_excel(args.valid_dataset)
+
+    total_loss = 0.
+    total_examples = 0.
+    progress_bar = tqdm(total=len(df), desc="Evaluation")
+    for idx, row in df.iterrows():
+        noise_filename = row["noise_filename"]
+        noise_offset = row["noise_offset"]
+        noise_duration = row["noise_duration"]
+
+        speech_filename = row["speech_filename"]
+        speech_offset = row["speech_offset"]
+        speech_duration = row["speech_duration"]
+
+        snr_db = row["snr_db"]
+
+        noise_wave, _ = librosa.load(
+            noise_filename,
+            sr=8000,
+            offset=noise_offset,
+            duration=noise_duration,
+        )
+        speech_wave, _ = librosa.load(
+            speech_filename,
+            sr=8000,
+            offset=speech_offset,
+            duration=speech_duration,
+        )
+        mix_wave: np.ndarray = mix_speech_and_noise(
+            speech=speech_wave,
+            noise=noise_wave,
+            snr_db=snr_db,
+        )
+        noise_wave = torch.tensor(noise_wave, dtype=torch.float32)
+        speech_wave = torch.tensor(speech_wave, dtype=torch.float32)
+        mix_wave: torch.Tensor = torch.tensor(mix_wave, dtype=torch.float32)
+
+        noise_wave = noise_wave.unsqueeze(dim=0)
+        speech_wave = speech_wave.unsqueeze(dim=0)
+        mix_wave = mix_wave.unsqueeze(dim=0)
+
+        noise_spec: torch.Tensor = stft_power.forward(noise_wave)
+        speech_spec: torch.Tensor = stft_power.forward(speech_wave)
+        mix_spec: torch.Tensor = stft_power.forward(mix_wave)
+
+        noise_spec = noise_spec[:, :-1, :]
+        speech_spec = speech_spec[:, :-1, :]
+        mix_spec = mix_spec[:, :-1, :]
+
+        mix_spec_complex: torch.Tensor = stft_complex.forward(mix_wave)
+        # mix_spec_complex shape: [batch_size, freq_dim (257), time_steps, 2]
+
+        speech_irm = speech_spec / (noise_spec + speech_spec)
+        speech_irm = torch.pow(speech_irm, 1.0)
+
+        snr_db: torch.Tensor = 10 * torch.log10(
+            speech_spec / (noise_spec + 1e-8)
+        )
+        snr_db = torch.mean(snr_db, dim=1, keepdim=True)
+        # snr_db shape: [batch_size, 1, time_steps]
+
+        mix_spec = mix_spec.to(device)
+        speech_irm_target = speech_irm.to(device)
+        snr_db_target = snr_db.to(device)
+
+        with torch.no_grad():
+            speech_irm_prediction, lsnr_prediction = model.forward(mix_spec)
+            irm_loss = irm_mse_loss.forward(speech_irm_prediction, speech_irm_target)
+            # snr_loss = snr_mse_loss.forward(lsnr_prediction, snr_db_target)
+            # loss = irm_loss + 0.1 * snr_loss
+            loss = irm_loss
+
+        # mix_spec_complex shape: [batch_size, freq_dim (257), time_steps, 2]
+        # speech_irm_prediction shape: [batch_size, freq_dim (256), time_steps]
+        batch_size, _, time_steps = speech_irm_prediction.shape
+        speech_irm_prediction = torch.concat(
+            [
+                speech_irm_prediction,
+                0.5*torch.ones(size=(batch_size, 1, time_steps), dtype=speech_irm_prediction.dtype).to(device)
+            ],
+            dim=1,
+        )
+        # speech_irm_prediction shape: [batch_size, freq_dim (257), time_steps]
+        speech_wave_enhanced, noise_wave_enhanced = enhance(mix_spec_complex, speech_irm_prediction)
+        save_audios(noise_wave, speech_wave, mix_wave, speech_wave_enhanced, noise_wave_enhanced, args.evaluation_audio_dir)
+
+        total_loss += loss.item()
+        total_examples += mix_spec.size(0)
+
+        evaluation_loss = total_loss / total_examples
+        evaluation_loss = round(evaluation_loss, 4)
+
+        progress_bar.update(1)
+        progress_bar.set_postfix({
+            "evaluation_loss": evaluation_loss,
+        })
+
+        if idx > args.limit:
+            break
+
+    return
+
+
+if __name__ == '__main__':
+    main()

examples/spectrum_dfnet_aishell/yaml/config.yaml

@@ -0,0 +1,38 @@
+model_name: "spectrum_unet_irm"
+
+# spec
+sample_rate: 8000
+n_fft: 512
+win_length: 200
+hop_length: 80
+
+spec_bins: 256
+
+# model
+conv_channels: 64
+conv_kernel_size_input:
+  - 3
+  - 3
+conv_kernel_size_inner:
+  - 1
+  - 3
+conv_lookahead: 0
+
+convt_kernel_size_inner:
+  - 1
+  - 3
+
+encoder_emb_skip_op: "none"
+encoder_emb_linear_groups: 16
+encoder_emb_hidden_size: 256
+
+lsnr_max: 30
+lsnr_min: -15
+
+decoder_emb_num_layers: 3
+decoder_emb_skip_op: "none"
+decoder_emb_linear_groups: 16
+decoder_emb_hidden_size: 256
+
+# runtime
+use_post_filter: true

toolbox/torchaudio/models/mpnet/__init__.py

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

toolbox/torchaudio/models/mpnet/modeling_mpnet.py

@@ -0,0 +1,9 @@
+#!/usr/bin/python3
+# -*- coding: utf-8 -*-
+"""
+https://huggingface.co/spaces/LeeSangHoon/HierSpeech_TTS/blob/main/denoiser/generator.py
+"""
+
+
+if __name__ == '__main__':
+    pass

toolbox/torchaudio/models/spectrum_dfnet/__init__.py

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

toolbox/torchaudio/models/spectrum_dfnet/configuration_spectrum_dfnet.py

@@ -0,0 +1,107 @@
+#!/usr/bin/python3
+# -*- coding: utf-8 -*-
+from typing import Tuple
+
+from toolbox.torchaudio.configuration_utils import PretrainedConfig
+
+
+class SpectrumDfNetConfig(PretrainedConfig):
+    def __init__(self,
+                 sample_rate: int = 8000,
+                 n_fft: int = 512,
+                 win_length: int = 200,
+                 hop_length: int = 80,
+
+                 spec_bins: int = 256,
+
+                 conv_channels: int = 64,
+                 conv_kernel_size_input: Tuple[int, int] = (3, 3),
+                 conv_kernel_size_inner: Tuple[int, int] = (1, 3),
+                 conv_lookahead: int = 0,
+
+                 convt_kernel_size_inner: Tuple[int, int] = (1, 3),
+
+                 embedding_hidden_size: int = 256,
+                 encoder_combine_op: str = "concat",
+
+                 encoder_emb_skip_op: str = "none",
+                 encoder_emb_linear_groups: int = 16,
+                 encoder_emb_hidden_size: int = 256,
+
+                 encoder_linear_groups: int = 32,
+
+                 lsnr_max: int = 30,
+                 lsnr_min: int = -15,
+                 norm_tau: float = 1.,
+
+                 decoder_emb_num_layers: int = 3,
+                 decoder_emb_skip_op: str = "none",
+                 decoder_emb_linear_groups: int = 16,
+                 decoder_emb_hidden_size: int = 256,
+
+                 df_decoder_hidden_size: int = 256,
+                 df_num_layers: int = 2,
+                 df_order: int = 5,
+                 df_bins: int = 96,
+                 df_gru_skip: str =  "grouped_linear",
+                 df_decoder_linear_groups: int =  16,
+                 df_pathway_kernel_size_t: int = 5,
+                 df_lookahead: int = 2,
+
+                 use_post_filter: bool = False,
+                 **kwargs
+                 ):
+        super(SpectrumDfNetConfig, self).__init__(**kwargs)
+        # transform
+        self.sample_rate = sample_rate
+        self.n_fft = n_fft
+        self.win_length = win_length
+        self.hop_length = hop_length
+
+        # spectrum
+        self.spec_bins = spec_bins
+
+        # conv
+        self.conv_channels = conv_channels
+        self.conv_kernel_size_input = conv_kernel_size_input
+        self.conv_kernel_size_inner = conv_kernel_size_inner
+        self.conv_lookahead = conv_lookahead
+
+        self.convt_kernel_size_inner = convt_kernel_size_inner
+
+        self.embedding_hidden_size = embedding_hidden_size
+
+        # encoder
+        self.encoder_emb_skip_op = encoder_emb_skip_op
+        self.encoder_emb_linear_groups = encoder_emb_linear_groups
+        self.encoder_emb_hidden_size = encoder_emb_hidden_size
+
+        self.encoder_linear_groups = encoder_linear_groups
+        self.encoder_combine_op = encoder_combine_op
+
+        self.lsnr_max = lsnr_max
+        self.lsnr_min = lsnr_min
+        self.norm_tau = norm_tau
+
+        # decoder
+        self.decoder_emb_num_layers = decoder_emb_num_layers
+        self.decoder_emb_skip_op = decoder_emb_skip_op
+        self.decoder_emb_linear_groups = decoder_emb_linear_groups
+        self.decoder_emb_hidden_size = decoder_emb_hidden_size
+
+        # df decoder
+        self.df_decoder_hidden_size = df_decoder_hidden_size
+        self.df_num_layers = df_num_layers
+        self.df_order = df_order
+        self.df_bins = df_bins
+        self.df_gru_skip = df_gru_skip
+        self.df_decoder_linear_groups = df_decoder_linear_groups
+        self.df_pathway_kernel_size_t = df_pathway_kernel_size_t
+        self.df_lookahead = df_lookahead
+
+        # runtime
+        self.use_post_filter = use_post_filter
+
+
+if __name__ == "__main__":
+    pass

toolbox/torchaudio/models/spectrum_dfnet/modeling_spectrum_dfnet.py

@@ -0,0 +1,920 @@
+#!/usr/bin/python3
+# -*- coding: utf-8 -*-
+import os
+import math
+from typing import Any, Callable, Dict, Iterable, List, Optional, Tuple, Union
+
+import numpy as np
+import torch
+import torch.nn as nn
+import torchaudio
+
+from toolbox.torchaudio.models.spectrum_dfnet.configuration_spectrum_dfnet import SpectrumDfNetConfig
+from toolbox.torchaudio.configuration_utils import CONFIG_FILE
+
+
+MODEL_FILE = "model.pt"
+
+
+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.Sequential):
+    def __init__(self,
+                 in_channels: int,
+                 out_channels: int,
+                 kernel_size: Union[int, Iterable[int]],
+                 fstride: int = 1,
+                 dilation: int = 1,
+                 fpad: bool = True,
+                 bias: bool = True,
+                 separable: bool = False,
+                 norm_layer: str = "batch_norm_2d",
+                 activation_layer: str = "relu",
+                 lookahead: int = 0
+                 ):
+        """
+        Causal Conv2d by delaying the signal for any lookahead.
+
+        Expected input format: [batch_size, channels, time_steps, spec_dim]
+
+        :param in_channels:
+        :param out_channels:
+        :param kernel_size:
+        :param fstride:
+        :param dilation:
+        :param fpad:
+        """
+        super(CausalConv2d, self).__init__()
+        kernel_size = (kernel_size, kernel_size) if isinstance(kernel_size, int) else tuple(kernel_size)
+
+        if fpad:
+            fpad_ = kernel_size[1] // 2 + dilation - 1
+        else:
+            fpad_ = 0
+
+        # for last 2 dim, pad (left, right, top, bottom).
+        pad = (0, 0, kernel_size[0] - 1 - lookahead, lookahead)
+
+        layers = list()
+        if any(x > 0 for x in pad):
+            layers.append(nn.ConstantPad2d(pad, 0.0))
+
+        groups = math.gcd(in_channels, out_channels) if separable else 1
+        if groups == 1:
+            separable = False
+        if max(kernel_size) == 1:
+            separable = False
+
+        layers.append(
+            nn.Conv2d(
+                in_channels,
+                out_channels,
+                kernel_size=kernel_size,
+                padding=(0, fpad_),
+                stride=(1, fstride),  # stride over time is always 1
+                dilation=(1, dilation),  # dilation over time is always 1
+                groups=groups,
+                bias=bias,
+            )
+        )
+
+        if separable:
+            layers.append(
+                nn.Conv2d(
+                    out_channels,
+                    out_channels,
+                    kernel_size=1,
+                    bias=False,
+                )
+            )
+
+        if norm_layer is not None:
+            norm_layer = norm_layer_dict[norm_layer]
+            layers.append(norm_layer(out_channels))
+
+        if activation_layer is not None:
+            activation_layer = activation_layer_dict[activation_layer]
+            layers.append(activation_layer())
+
+        super().__init__(*layers)
+
+    def forward(self, inputs):
+        for module in self:
+            inputs = module(inputs)
+        return inputs
+
+
+class CausalConvTranspose2d(nn.Sequential):
+    def __init__(self,
+                 in_channels: int,
+                 out_channels: int,
+                 kernel_size: Union[int, Iterable[int]],
+                 fstride: int = 1,
+                 dilation: int = 1,
+                 fpad: bool = True,
+                 bias: bool = True,
+                 separable: bool = False,
+                 norm_layer: str = "batch_norm_2d",
+                 activation_layer: str = "relu",
+                 lookahead: int = 0
+                 ):
+        """
+        Causal ConvTranspose2d.
+
+        Expected input format: [batch_size, channels, time_steps, spec_dim]
+        """
+        super(CausalConvTranspose2d, self).__init__()
+
+        kernel_size = (kernel_size, kernel_size) if isinstance(kernel_size, int) else kernel_size
+
+        if fpad:
+            fpad_ = kernel_size[1] // 2
+        else:
+            fpad_ = 0
+
+        # for last 2 dim, pad (left, right, top, bottom).
+        pad = (0, 0, kernel_size[0] - 1 - lookahead, lookahead)
+
+        layers = []
+        if any(x > 0 for x in pad):
+            layers.append(nn.ConstantPad2d(pad, 0.0))
+
+        groups = math.gcd(in_channels, out_channels) if separable else 1
+        if groups == 1:
+            separable = False
+
+        layers.append(
+            nn.ConvTranspose2d(
+                in_channels,
+                out_channels,
+                kernel_size=kernel_size,
+                padding=(kernel_size[0] - 1, fpad_ + dilation - 1),
+                output_padding=(0, fpad_),
+                stride=(1, fstride),  # stride over time is always 1
+                dilation=(1, dilation),  # dilation over time is always 1
+                groups=groups,
+                bias=bias,
+            )
+        )
+
+        if separable:
+            layers.append(
+                nn.Conv2d(
+                    out_channels,
+                    out_channels,
+                    kernel_size=1,
+                    bias=False,
+                )
+            )
+
+        if norm_layer is not None:
+            norm_layer = norm_layer_dict[norm_layer]
+            layers.append(norm_layer(out_channels))
+
+        if activation_layer is not None:
+            activation_layer = activation_layer_dict[activation_layer]
+            layers.append(activation_layer())
+
+        super().__init__(*layers)
+
+
+class GroupedLinear(nn.Module):
+
+    def __init__(self, input_size: int, hidden_size: int, groups: int = 1):
+        super().__init__()
+        # self.weight: Tensor
+        self.input_size = input_size
+        self.hidden_size = hidden_size
+        self.groups = groups
+        assert input_size % groups == 0, f"Input size {input_size} not divisible by {groups}"
+        assert hidden_size % groups == 0, f"Hidden size {hidden_size} not divisible by {groups}"
+        self.ws = input_size // groups
+        self.register_parameter(
+            "weight",
+            torch.nn.Parameter(
+                torch.zeros(groups, input_size // groups, hidden_size // groups), requires_grad=True
+            ),
+        )
+        self.reset_parameters()
+
+    def reset_parameters(self):
+        nn.init.kaiming_uniform_(self.weight, a=math.sqrt(5))  # type: ignore
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        # x: [..., I]
+        b, t, _ = x.shape
+        # new_shape = list(x.shape)[:-1] + [self.groups, self.ws]
+        new_shape = (b, t, self.groups, self.ws)
+        x = x.view(new_shape)
+        # The better way, but not supported by torchscript
+        # x = x.unflatten(-1, (self.groups, self.ws))  # [..., G, I/G]
+        x = torch.einsum("btgi,gih->btgh", x, self.weight)  # [..., G, H/G]
+        x = x.flatten(2, 3)  # [B, T, H]
+        return x
+
+    def __repr__(self):
+        cls = self.__class__.__name__
+        return f"{cls}(input_size: {self.input_size}, hidden_size: {self.hidden_size}, groups: {self.groups})"
+
+
+class SqueezedGRU_S(nn.Module):
+    """
+    SGE net: Video object detection with squeezed GRU and information entropy map
+    https://arxiv.org/abs/2106.07224
+    """
+
+    def __init__(
+        self,
+        input_size: int,
+        hidden_size: int,
+        output_size: Optional[int] = None,
+        num_layers: int = 1,
+        linear_groups: int = 8,
+        batch_first: bool = True,
+        skip_op: str = "none",
+        activation_layer: str = "identity",
+    ):
+        super().__init__()
+        self.input_size = input_size
+        self.hidden_size = hidden_size
+
+        self.linear_in = nn.Sequential(
+            GroupedLinear(
+                input_size=input_size,
+                hidden_size=hidden_size,
+                groups=linear_groups,
+            ),
+            activation_layer_dict[activation_layer](),
+        )
+
+        # gru skip operator
+        self.gru_skip_op = None
+
+        if skip_op == "none":
+            self.gru_skip_op = None
+        elif skip_op == "identity":
+            if not input_size != output_size:
+                raise AssertionError("Dimensions do not match")
+            self.gru_skip_op = nn.Identity()
+        elif skip_op == "grouped_linear":
+            self.gru_skip_op = GroupedLinear(
+                input_size=hidden_size,
+                hidden_size=hidden_size,
+                groups=linear_groups,
+            )
+        else:
+            raise NotImplementedError()
+
+        self.gru = nn.GRU(
+            input_size=hidden_size,
+            hidden_size=hidden_size,
+            num_layers=num_layers,
+            batch_first=batch_first,
+            bidirectional=False,
+        )
+
+        if output_size is not None:
+            self.linear_out = nn.Sequential(
+                GroupedLinear(
+                    input_size=hidden_size,
+                    hidden_size=output_size,
+                    groups=linear_groups,
+                ),
+                activation_layer_dict[activation_layer](),
+            )
+        else:
+            self.linear_out = nn.Identity()
+
+    def forward(self, inputs: torch.Tensor, h=None) -> Tuple[torch.Tensor, torch.Tensor]:
+        x = self.linear_in(inputs)
+
+        x, h = self.gru.forward(x, h)
+
+        x = self.linear_out(x)
+
+        if self.gru_skip_op is not None:
+            x = x + self.gru_skip_op(inputs)
+
+        return x, h
+
+
+class Add(nn.Module):
+    def forward(self, a, b):
+        return a + b
+
+
+class Concat(nn.Module):
+    def forward(self, a, b):
+        return torch.cat((a, b), dim=-1)
+
+
+class Encoder(nn.Module):
+    def __init__(self, config: SpectrumDfNetConfig):
+        super(Encoder, self).__init__()
+        self.embedding_input_size = config.conv_channels * config.spec_bins // 4
+        self.embedding_output_size = config.conv_channels * config.spec_bins // 4
+        self.embedding_hidden_size = config.embedding_hidden_size
+
+        self.spec_conv0 = CausalConv2d(
+            in_channels=1,
+            out_channels=config.conv_channels,
+            kernel_size=config.conv_kernel_size_input,
+            bias=False,
+            separable=True,
+            fstride=1,
+            lookahead=config.conv_lookahead,
+        )
+        self.spec_conv1 = CausalConv2d(
+            in_channels=config.conv_channels,
+            out_channels=config.conv_channels,
+            kernel_size=config.conv_kernel_size_inner,
+            bias=False,
+            separable=True,
+            fstride=2,
+            lookahead=config.conv_lookahead,
+        )
+        self.spec_conv2 = CausalConv2d(
+            in_channels=config.conv_channels,
+            out_channels=config.conv_channels,
+            kernel_size=config.conv_kernel_size_inner,
+            bias=False,
+            separable=True,
+            fstride=2,
+            lookahead=config.conv_lookahead,
+        )
+        self.spec_conv3 = CausalConv2d(
+            in_channels=config.conv_channels,
+            out_channels=config.conv_channels,
+            kernel_size=config.conv_kernel_size_inner,
+            bias=False,
+            separable=True,
+            fstride=1,
+            lookahead=config.conv_lookahead,
+        )
+
+        self.df_conv0 = CausalConv2d(
+            in_channels=2,
+            out_channels=config.conv_channels,
+            kernel_size=config.conv_kernel_size_input,
+            bias=False,
+            separable=True,
+        )
+        self.df_conv1 = CausalConv2d(
+            in_channels=config.conv_channels,
+            out_channels=config.conv_channels,
+            kernel_size=config.conv_kernel_size_inner,
+            bias=False,
+            separable=True,
+            fstride=2,
+        )
+        self.df_fc_emb = nn.Sequential(
+            GroupedLinear(
+                config.conv_channels * config.df_bins // 2,
+                self.embedding_input_size,
+                groups=config.encoder_linear_groups
+            ),
+            nn.ReLU(inplace=True)
+        )
+
+        if config.encoder_combine_op == "concat":
+            self.embedding_input_size *= 2
+            self.combine = Concat()
+        else:
+            self.combine = Add()
+
+        # emb_gru
+        if config.spec_bins % 8 != 0:
+            raise AssertionError("spec_bins should be divisible by 8")
+
+        self.emb_gru = SqueezedGRU_S(
+            self.embedding_input_size,
+            self.embedding_hidden_size,
+            output_size=self.embedding_output_size,
+            num_layers=1,
+            batch_first=True,
+            skip_op=config.encoder_emb_skip_op,
+            linear_groups=config.encoder_emb_linear_groups,
+            activation_layer="relu",
+        )
+
+        # lsnr
+        self.lsnr_fc = nn.Sequential(
+            nn.Linear(self.embedding_output_size, 1),
+            nn.Sigmoid()
+        )
+        self.lsnr_scale = config.lsnr_max - config.lsnr_min
+        self.lsnr_offset = config.lsnr_min
+
+    def forward(self,
+                feat_power: torch.Tensor,
+                feat_spec: torch.Tensor,
+                hidden_state: torch.Tensor = None,
+                ):
+        # feat_power shape: (batch_size, 1, time_steps, spec_dim)
+        e0 = self.spec_conv0.forward(feat_power)
+        e1 = self.spec_conv1.forward(e0)
+        e2 = self.spec_conv2.forward(e1)
+        e3 = self.spec_conv3.forward(e2)
+        # e0 shape: [batch_size, channels, time_steps, spec_dim]
+        # e1 shape: [batch_size, channels, time_steps, spec_dim // 2]
+        # e2 shape: [batch_size, channels, time_steps, spec_dim // 4]
+        # e3 shape: [batch_size, channels, time_steps, spec_dim // 4]
+
+        # feat_spec, shape: (batch_size, 2, time_steps, df_bins)
+        c0 = self.df_conv0(feat_spec)
+        c1 = self.df_conv1(c0)
+        # c0 shape: [batch_size, channels, time_steps, df_bins]
+        # c1 shape: [batch_size, channels, time_steps, df_bins // 2]
+
+        cemb = c1.permute(0, 2, 3, 1)
+        # cemb shape: [batch_size, time_steps, df_bins // 2, channels]
+        cemb = cemb.flatten(2)
+        # cemb shape: [batch_size, time_steps, df_bins // 2 * channels]
+        cemb = self.df_fc_emb(cemb)
+        # cemb shape: [batch_size, time_steps, spec_dim // 4 * channels]
+
+        # e3 shape: [batch_size, channels, time_steps, spec_dim // 4]
+        emb = e3.permute(0, 2, 3, 1)
+        # emb shape: [batch_size, time_steps, spec_dim // 4, channels]
+        emb = emb.flatten(2)
+        # emb shape: [batch_size, time_steps, spec_dim // 4 * channels]
+
+        emb = self.combine(emb, cemb)
+        # if concat; emb shape: [batch_size, time_steps, spec_dim // 4 * channels * 2]
+        # if add; emb shape: [batch_size, time_steps, spec_dim // 4 * channels]
+
+        emb, h = self.emb_gru.forward(emb, hidden_state)
+        # emb shape: [batch_size, time_steps, spec_dim // 4 * channels]
+        # h shape: [batch_size, 1, spec_dim]
+
+        lsnr = self.lsnr_fc(emb) * self.lsnr_scale + self.lsnr_offset
+        # lsnr shape: [batch_size, time_steps, 1]
+
+        return e0, e1, e2, e3, emb, c0, lsnr, h
+
+
+class Decoder(nn.Module):
+    def __init__(self, config: SpectrumDfNetConfig):
+        super(Decoder, self).__init__()
+
+        if config.spec_bins % 8 != 0:
+            raise AssertionError("spec_bins should be divisible by 8")
+
+        self.emb_in_dim = config.conv_channels * config.spec_bins // 4
+        self.emb_out_dim = config.conv_channels * config.spec_bins // 4
+        self.emb_hidden_dim = config.decoder_emb_hidden_size
+
+        self.emb_gru = SqueezedGRU_S(
+            self.emb_in_dim,
+            self.emb_hidden_dim,
+            output_size=self.emb_out_dim,
+            num_layers=config.decoder_emb_num_layers - 1,
+            batch_first=True,
+            skip_op=config.decoder_emb_skip_op,
+            linear_groups=config.decoder_emb_linear_groups,
+            activation_layer="relu",
+        )
+        self.conv3p = CausalConv2d(
+            in_channels=config.conv_channels,
+            out_channels=config.conv_channels,
+            kernel_size=1,
+            bias=False,
+            separable=True,
+            fstride=1,
+            lookahead=config.conv_lookahead,
+        )
+        self.convt3 = CausalConv2d(
+            in_channels=config.conv_channels,
+            out_channels=config.conv_channels,
+            kernel_size=config.conv_kernel_size_inner,
+            bias=False,
+            separable=True,
+            fstride=1,
+            lookahead=config.conv_lookahead,
+        )
+        self.conv2p = CausalConv2d(
+            in_channels=config.conv_channels,
+            out_channels=config.conv_channels,
+            kernel_size=1,
+            bias=False,
+            separable=True,
+            fstride=1,
+            lookahead=config.conv_lookahead,
+        )
+        self.convt2 = CausalConvTranspose2d(
+            in_channels=config.conv_channels,
+            out_channels=config.conv_channels,
+            kernel_size=config.convt_kernel_size_inner,
+            bias=False,
+            separable=True,
+            fstride=2,
+            lookahead=config.conv_lookahead,
+        )
+        self.conv1p = CausalConv2d(
+            in_channels=config.conv_channels,
+            out_channels=config.conv_channels,
+            kernel_size=1,
+            bias=False,
+            separable=True,
+            fstride=1,
+            lookahead=config.conv_lookahead,
+        )
+        self.convt1 = CausalConvTranspose2d(
+            in_channels=config.conv_channels,
+            out_channels=config.conv_channels,
+            kernel_size=config.convt_kernel_size_inner,
+            bias=False,
+            separable=True,
+            fstride=2,
+            lookahead=config.conv_lookahead,
+        )
+        self.conv0p = CausalConv2d(
+            in_channels=config.conv_channels,
+            out_channels=config.conv_channels,
+            kernel_size=1,
+            bias=False,
+            separable=True,
+            fstride=1,
+            lookahead=config.conv_lookahead,
+        )
+        self.conv0_out = CausalConv2d(
+            in_channels=config.conv_channels,
+            out_channels=1,
+            kernel_size=config.conv_kernel_size_inner,
+            activation_layer="sigmoid",
+            bias=False,
+            separable=True,
+            fstride=1,
+            lookahead=config.conv_lookahead,
+        )
+
+    def forward(self, emb, e3, e2, e1, e0) -> torch.Tensor:
+        # Estimates erb mask
+        b, _, t, f8 = e3.shape
+
+        # emb shape: [batch_size, time_steps, (freq_dim // 4) * conv_channels]
+        emb, _ = self.emb_gru(emb)
+        # emb shape: [batch_size, conv_channels, time_steps, freq_dim // 4]
+        emb = emb.view(b, t, f8, -1).permute(0, 3, 1, 2)
+        e3 = self.convt3(self.conv3p(e3) + emb)
+        # e3 shape: [batch_size, conv_channels, time_steps, freq_dim // 4]
+        e2 = self.convt2(self.conv2p(e2) + e3)
+        # e2 shape: [batch_size, conv_channels, time_steps, freq_dim // 2]
+        e1 = self.convt1(self.conv1p(e1) + e2)
+        # e1 shape: [batch_size, conv_channels, time_steps, freq_dim]
+        mask = self.conv0_out(self.conv0p(e0) + e1)
+        # mask shape: [batch_size, 1, time_steps, freq_dim]
+        return mask
+
+
+class DfDecoder(nn.Module):
+    def __init__(self, config: SpectrumDfNetConfig):
+        super(DfDecoder, self).__init__()
+
+        self.embedding_input_size = config.conv_channels * config.spec_bins // 4
+        self.df_decoder_hidden_size = config.df_decoder_hidden_size
+        self.df_num_layers = config.df_num_layers
+
+        self.df_order = config.df_order
+
+        self.df_bins = config.df_bins
+        self.df_out_ch = config.df_order * 2
+
+        self.df_convp = CausalConv2d(
+            config.conv_channels,
+            self.df_out_ch,
+            fstride=1,
+            kernel_size=(config.df_pathway_kernel_size_t, 1),
+            separable=True,
+            bias=False,
+        )
+        self.df_gru = SqueezedGRU_S(
+            self.embedding_input_size,
+            self.df_decoder_hidden_size,
+            num_layers=self.df_num_layers,
+            batch_first=True,
+            skip_op="none",
+            activation_layer="relu",
+        )
+
+        if config.df_gru_skip == "none":
+            self.df_skip = None
+        elif config.df_gru_skip == "identity":
+            if config.embedding_hidden_size != config.df_decoder_hidden_size:
+                raise AssertionError("Dimensions do not match")
+            self.df_skip = nn.Identity()
+        elif config.df_gru_skip == "grouped_linear":
+            self.df_skip = GroupedLinear(
+                self.embedding_input_size,
+                self.df_decoder_hidden_size,
+                groups=config.df_decoder_linear_groups
+            )
+        else:
+            raise NotImplementedError()
+
+        self.df_out: nn.Module
+        out_dim = self.df_bins * self.df_out_ch
+
+        self.df_out = nn.Sequential(
+            GroupedLinear(
+                input_size=self.df_decoder_hidden_size,
+                hidden_size=out_dim,
+                groups=config.df_decoder_linear_groups
+            ),
+            nn.Tanh()
+        )
+        self.df_fc_a = nn.Sequential(
+            nn.Linear(self.df_decoder_hidden_size, 1),
+            nn.Sigmoid()
+        )
+
+    def forward(self, emb: torch.Tensor, c0: torch.Tensor) -> torch.Tensor:
+        # emb shape: [batch_size, time_steps, df_bins // 4 * channels]
+        b, t, _ = emb.shape
+        df_coefs, _ = self.df_gru(emb)
+        if self.df_skip is not None:
+            df_coefs = df_coefs + self.df_skip(emb)
+        # df_coefs shape: [batch_size, time_steps, df_decoder_hidden_size]
+
+        # c0 shape: [batch_size, channels, time_steps, df_bins]
+        c0 = self.df_convp(c0)
+        # c0 shape: [batch_size, df_order * 2, time_steps, df_bins]
+        c0 = c0.permute(0, 2, 3, 1)
+        # c0 shape: [batch_size, time_steps, df_bins, df_order * 2]
+
+        df_coefs = self.df_out(df_coefs)  # [B, T, F*O*2], O: df_order
+        # df_coefs shape: [batch_size, time_steps, df_bins * df_order * 2]
+        df_coefs = df_coefs.view(b, t, self.df_bins, self.df_out_ch)
+        # df_coefs shape: [batch_size, time_steps, df_bins, df_order * 2]
+        df_coefs = df_coefs + c0
+        # df_coefs shape: [batch_size, time_steps, df_bins, df_order * 2]
+        return df_coefs
+
+
+class DfOutputReshapeMF(nn.Module):
+    """Coefficients output reshape for multiframe/MultiFrameModule
+
+    Requires input of shape B, C, T, F, 2.
+    """
+
+    def __init__(self, df_order: int, df_bins: int):
+        super().__init__()
+        self.df_order = df_order
+        self.df_bins = df_bins
+
+    def forward(self, coefs: torch.Tensor) -> torch.Tensor:
+        # [B, T, F, O*2] -> [B, O, T, F, 2]
+        new_shape = list(coefs.shape)
+        new_shape[-1] = -1
+        new_shape.append(2)
+        coefs = coefs.view(new_shape)
+        coefs = coefs.permute(0, 3, 1, 2, 4)
+        return coefs
+
+
+class Mask(nn.Module):
+    def __init__(self, use_post_filter: bool = False, eps: float = 1e-12):
+        super().__init__()
+        self.use_post_filter = use_post_filter
+        self.eps = eps
+
+    def post_filter(self, mask: torch.Tensor, beta: float = 0.02) -> torch.Tensor:
+        """
+        Post-Filter
+
+        A Perceptually-Motivated Approach for Low-Complexity, Real-Time Enhancement of Fullband Speech.
+        https://arxiv.org/abs/2008.04259
+
+        :param mask: Real valued mask, typically of shape [B, C, T, F].
+        :param beta: Global gain factor.
+        :return:
+        """
+        mask_sin = mask * torch.sin(np.pi * mask / 2)
+        mask_pf = (1 + beta) * mask / (1 + beta * mask.div(mask_sin.clamp_min(self.eps)).pow(2))
+        return mask_pf
+
+    def forward(self, spec: torch.Tensor, mask: torch.Tensor) -> torch.Tensor:
+        # spec shape: [batch_size, 1, time_steps, spec_bins, 2]
+
+        if not self.training and self.use_post_filter:
+            mask = self.post_filter(mask)
+
+        # mask shape: [batch_size, 1, time_steps, spec_bins]
+        mask = mask.unsqueeze(4)
+        # mask shape: [batch_size, 1, time_steps, spec_bins, 1]
+        return spec * mask
+
+
+class DeepFiltering(nn.Module):
+    def __init__(self,
+                 df_bins: int,
+                 df_order: int,
+                 lookahead: int = 0,
+                 ):
+        super(DeepFiltering, self).__init__()
+        self.df_bins = df_bins
+        self.df_order = df_order
+        self.need_unfold = df_order > 1
+        self.lookahead = lookahead
+
+        self.pad = nn.ConstantPad2d((0, 0, df_order - 1 - lookahead, lookahead), 0.0)
+
+    def spec_unfold(self, spec: torch.Tensor):
+        """
+        Pads and unfolds the spectrogram according to frame_size.
+        :param spec: complex Tensor, Spectrogram of shape [B, C, T, F].
+        :return: Tensor, Unfolded spectrogram of shape [B, C, T, F, N], where N: frame_size.
+        """
+        if self.need_unfold:
+            # spec shape: [batch_size, spec_bins, time_steps]
+            spec_pad = self.pad(spec)
+            # spec_pad shape: [batch_size, 1, time_steps_pad, spec_bins]
+            spec_unfold = spec_pad.unfold(2, self.df_order, 1)
+            # spec_unfold shape: [batch_size, 1, time_steps, spec_bins, df_order]
+            return spec_unfold
+        else:
+            return spec.unsqueeze(-1)
+
+    def forward(self,
+                spec: torch.Tensor,
+                coefs: torch.Tensor,
+                ):
+        # spec shape: [batch_size, 1, time_steps, spec_bins, 2]
+        spec_u = self.spec_unfold(torch.view_as_complex(spec))
+        # spec_u shape: [batch_size, 1, time_steps, spec_bins, df_order]
+
+        # coefs shape: [batch_size, df_order, time_steps, df_bins, 2]
+        coefs = torch.view_as_complex(coefs)
+        # coefs shape: [batch_size, df_order, time_steps, df_bins]
+        spec_f = spec_u.narrow(-2, 0, self.df_bins)
+        # spec_f shape: [batch_size, 1, time_steps, df_bins, df_order]
+
+        coefs = coefs.view(coefs.shape[0], -1, self.df_order, *coefs.shape[2:])
+        # coefs shape: [batch_size, 1, df_order, time_steps, df_bins]
+
+        spec_f = self.df(spec_f, coefs)
+        # spec_f shape: [batch_size, 1, time_steps, df_bins]
+
+        if self.training:
+            spec = spec.clone()
+        spec[..., :self.df_bins, :] = torch.view_as_real(spec_f)
+        # spec shape: [batch_size, 1, time_steps, spec_bins, 2]
+        return spec
+
+    @staticmethod
+    def df(spec: torch.Tensor, coefs: torch.Tensor) -> torch.Tensor:
+        """
+        Deep filter implementation using `torch.einsum`. Requires unfolded spectrogram.
+        :param spec: (complex Tensor). Spectrogram of shape [B, C, T, F, N].
+        :param coefs: (complex Tensor). Coefficients of shape [B, C, N, T, F].
+        :return: (complex Tensor). Spectrogram of shape [B, C, T, F].
+        """
+        return torch.einsum("...tfn,...ntf->...tf", spec, coefs)
+
+
+class SpectrumDfNet(nn.Module):
+    def __init__(self, config: SpectrumDfNetConfig):
+        super(SpectrumDfNet, self).__init__()
+        self.config = config
+        self.encoder = Encoder(config)
+        self.decoder = Decoder(config)
+
+        self.df_decoder = DfDecoder(config)
+        self.df_out_transform = DfOutputReshapeMF(config.df_order, config.df_bins)
+        self.df_op = DeepFiltering(
+            df_bins=config.df_bins,
+            df_order=config.df_order,
+            lookahead=config.df_lookahead,
+        )
+
+        self.mask = Mask(use_post_filter=config.use_post_filter)
+
+    def forward(self,
+                spec_complex: torch.Tensor,
+                ):
+        feat_power = torch.square(torch.abs(spec_complex))
+        feat_power = feat_power.unsqueeze(1).permute(0, 1, 3, 2)
+        # feat_power shape: [batch_size, spec_bins, time_steps]
+        # feat_power shape: [batch_size, 1, spec_bins, time_steps]
+        # feat_power shape: [batch_size, 1, time_steps, spec_bins]
+
+        # spec shape: [batch_size, spec_bins, time_steps]
+        feat_spec = torch.view_as_real(spec_complex)
+        # spec shape: [batch_size, spec_bins, time_steps, 2]
+        feat_spec = feat_spec.permute(0, 3, 2, 1)
+        # feat_spec shape: [batch_size, 2, time_steps, spec_bins]
+        feat_spec = feat_spec[..., :self.df_decoder.df_bins]
+        # feat_spec shape: [batch_size, 2, time_steps, df_bins]
+
+        # spec shape: [batch_size, spec_bins, time_steps]
+        spec = torch.unsqueeze(spec_complex, dim=1)
+        # spec shape: [batch_size, 1, spec_bins, time_steps]
+        spec = spec.permute(0, 1, 3, 2)
+        # spec shape: [batch_size, 1, time_steps, spec_bins]
+        spec = torch.view_as_real(spec)
+        # spec shape: [batch_size, 1, time_steps, spec_bins, 2]
+
+        e0, e1, e2, e3, emb, c0, lsnr, h = self.encoder.forward(feat_power, feat_spec)
+
+        mask = self.decoder.forward(emb, e3, e2, e1, e0)
+        # mask shape: [batch_size, 1, time_steps, spec_bins]
+        if torch.any(mask > 1) or torch.any(mask < 0):
+            raise AssertionError
+
+        spec_m = self.mask.forward(spec, mask)
+
+        # lsnr shape: [batch_size, time_steps, 1]
+        lsnr = torch.transpose(lsnr, dim0=2, dim1=1)
+        # lsnr shape: [batch_size, 1, time_steps]
+
+        df_coefs = self.df_decoder.forward(emb, c0)
+        df_coefs = self.df_out_transform(df_coefs)
+        # df_coefs shape: [batch_size, df_order, time_steps, df_bins, 2]
+
+        spec_e = self.df_op.forward(spec.clone(), df_coefs)
+        # spec_e shape: [batch_size, 1, time_steps, spec_bins, 2]
+
+        spec_e[..., self.df_decoder.df_bins:, :] = spec_m[..., self.df_decoder.df_bins:, :]
+        return spec_e, mask, lsnr
+
+
+class SpectrumDfNetPretrainedModel(SpectrumDfNet):
+    def __init__(self,
+                 config: SpectrumDfNetConfig,
+                 ):
+        super(SpectrumDfNetPretrainedModel, self).__init__(
+            config=config,
+        )
+
+    @classmethod
+    def from_pretrained(cls, pretrained_model_name_or_path, **kwargs):
+        config = SpectrumDfNetConfig.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():
+
+    transformer = torchaudio.transforms.Spectrogram(
+        n_fft=512,
+        win_length=200,
+        hop_length=80,
+        window_fn=torch.hamming_window,
+        power=None,
+    )
+
+    config = SpectrumDfNetConfig()
+    model = SpectrumDfNet(config=config)
+
+    inputs = torch.randn(size=(1, 16000), dtype=torch.float32)
+    spec_complex = transformer.forward(inputs)
+    spec_complex = spec_complex[:, :-1, :]
+
+    output = model.forward(spec_complex)
+    print(output[0].shape)
+    return
+
+
+if __name__ == '__main__':
+    main()