FSPEN (Unofficial Implementation)

Description

This repository contains an unofficial PyTorch implementation of FSPEN (Frequency-domain Sub-band Phase Estimation Network).

Paper: FSPEN: AN ULTRA-LIGHTWEIGHT NETWORK FOR REAL TIME SPEECH ENAHNCMENT
Status: Under development.

Installation

Pip

# clone project
git clone https://github.com/iliasslasri/fspen
cd fspen

# [OPTIONAL] create conda environment
conda create -n myenv python=3.9
conda activate myenv

pip install \
  "torch==2.0.1+cu118" \
  "torchvision==0.15.2+cu118" \
  "torchaudio==2.0.2+cu118" \
  "lightning==2.0.9" \
  "torchmetrics==0.11.4" \
  "numpy<2.0" \
  "pesq" \
  "hydra-colorlog" \
  --extra-index-url https://download.pytorch.org/whl/cu118

How to run

Train model with default configuration

# train on CPU
python src/train.py trainer=cpu

# train on GPU
python src/train.py trainer=gpu

Train model with chosen experiment configuration from configs/experiment/

python src/train.py experiment=experiment_name.yaml

You can override any parameter from command line like this

python src/train.py trainer.max_epochs=20 data.batch_size=64

To eval with the correct model config:

python eval.py \
  --config-dir ./logs/run_name/runs/XX_XX_XX/.hydra/ \
  --config-name config.yaml \
  ckpt_path="CKPT_PATH"

Credits

• lightning-hydra-template for the main structure of the code