A jax port of auraloss
Project description
fouriax
Documentation: https://mikesol.github.io/fouriax
Source Code: https://github.com/mikesol/fouriax
PyPI: https://pypi.org/project/fouriax/
A jax port of auraloss
.
Installation
pip install fouriax
Usage
import jax
import fouriax.stft as stft
from jax.nn.initializers import lecun_normal
key = jax.random.PRNGKey(0)
key1, key2 = jax.random.split(key)
shape = (4, 4098, 1)
# Initialize the tensor using LeCun normal distribution
input = lecun_normal()(key1, shape)
target = lecun_normal()(key2, shape)
fft_sizes = [1024, 2048, 512]
hop_sizes = [120, 240, 50]
win_lengths = [600, 1200, 240]
params = [
stft.init_stft_params(x, y, z)
for x, y, z in zip(fft_sizes, hop_sizes, win_lengths)
]
loss = multi_resolution_stft_loss(params, input, target)
Loss functions
We categorize the loss functions as either time-domain or frequency-domain approaches. Additionally, we include perceptual transforms.
Loss function | Interface | Reference |
---|---|---|
Time domain | ||
Error-to-signal ratio (ESR) | fouriax.time.esr_loss() |
Wright & Välimäki, 2019 |
DC error (DC) | auraloss.time.DCLoss() |
Wright & Välimäki, 2019 |
Log hyperbolic cosine (Log-cosh) | fouriax.time.log_cosh_loss() |
Chen et al., 2019 |
Frequency domain | ||
Aggregate STFT | fouriax.freq.stft_loss() |
Arik et al., 2018 |
Multi-resolution STFT | fouriax.freq.multi_resolution_stft_loss() |
Yamamoto et al., 2019* |
Perceptual transforms | ||
FIR pre-emphasis filters | fouriax.perceptual.fir_filter() |
Wright & Välimäki, 2019 |
* Wang et al., 2019 also propose a multi-resolution spectral loss (that Engel et al., 2020 follow), but they do not include both the log magnitude (L1 distance) and spectral convergence terms, introduced in Arik et al., 2018, and then extended for the multi-resolution case in Yamamoto et al., 2019.
PVC
A partial port of core routines in Paul Koonce's PVC can be found in pvc.py
. This includes a novel FFT algorithm called fkt
(Fast Koonce Transform) that, combined with convert
, produces amplitude/frequency pairs for a given signal. This is often more attractive to use in loss functions than a garden-variety FFT because it provides better frequency information.
There is also a noscbank
method that allows for resynthesis. This can be used as a simple recurrent layer at the end of a network to do waveform synthesis.
Development
- Clone this repository
- Requirements:
- Poetry
- Python 3.7+
- Create a virtual environment and install the dependencies
poetry install
- Activate the virtual environment
poetry shell
Testing
pytest
Documentation
The documentation is automatically generated from the content of the docs directory and from the docstrings of the public signatures of the source code. The documentation is updated and published as a Github project page automatically as part each release.
Releasing
Trigger the Draft release workflow (press Run workflow). This will update the changelog & version and create a GitHub release which is in Draft state.
Find the draft release from the GitHub releases and publish it. When a release is published, it'll trigger release workflow which creates PyPI release and deploys updated documentation.
Pre-commit
Pre-commit hooks run all the auto-formatters (e.g. black
, isort
), linters (e.g. mypy
, flake8
), and other quality
checks to make sure the changeset is in good shape before a commit/push happens.
You can install the hooks with (runs for each commit):
pre-commit install
Or if you want them to run only for each push:
pre-commit install -t pre-push
Or if you want e.g. want to run all checks manually for all files:
pre-commit run --all-files
This project was generated using the wolt-python-package-cookiecutter template.
Project details
Download files
Download the file for your platform. If you're not sure which to choose, learn more about installing packages.