A Gated Reccurent Unit that models IIR filters.
Explore the docs »
View Demo
·
Report Bug
·
Request Feature
Table of Contents
This repo contains python scripts for :
- An input-size agnostic GRU model : takes a buffer of audio samples and a normalized$^*$ cut-off frequency, and outputs the filtered buffer according to given cutoff frequency : model.py
- A script to generate a training dataset : create_dataset.py
- A training script for the model : train.py
- A script to monitor generalization curves during training : generalization_curves.py
- A model evaluation notebook to analyse the model's accuracy and quality with notably Bode graphs : model_eval.ipynb
- A script to export a trained model to the ONNX Runtime format export_to_onnx.py
- An example of C++ code using ONNX Runtime to run inference of the exported model : onnx_inference_test
Currently, the scripts provides ways to train a model to reproduce lowpass and highpass butterworth and chebyshev 1 filters.
normalize_freq in create_dataset.py.
Start by creating a virtual python environment and install the needed dependencies. Note that if you want to use your GPU to train your model, you need should check the dedicated Pytorch page and replace dedicated dependencies in requirements.txt accordingly.
- You can create the environment with the needed dependencies by running :
python -m venv .venv source .venv/bin/activate pip install -r requirements.txt # Check wanted pytorch package before running this line
The folders order1 and order2 contains already trained models :
- order1/butter_lowpass.pt :
- hidden size = 64
- num layers = 2
- sample rate = 48000
- minimum prefered buffer size = 128
- filter algo = butterworth
- filter type = lowpass
- order2/cheby1_ripple5_low.pt :
- hidden size = 128
- num layers = 2
- sample rate = 48000
- minimum prefered buffer size = 256
- filter algo = cheby1
- ripple value = 5
- filter type = lowpass
Here we present how to train your own model and the recommended parameters to use.
The dataset generation script produces samples of white noise and sine sweeps. It can be used with the following arguments :
usage: create_dataset.py [-h] [--sample_rate SAMPLE_RATE] [--buffer_size BUFFER_SIZE] [--max_buffer_amount MAX_BUFFER_AMOUNT]
options:
-h, --help show this help message and exit
--sample_rate SAMPLE_RATE
Generated signal sample-rate
--buffer_size BUFFER_SIZE
Generated signal will have buffer_size * max_buffer_amount samples
--max_buffer_amount MAX_BUFFER_AMOUNT
Generated signal will have buffer_size * max_buffer_amount samplesThe generated white noise has specific characteristics :
- It is bandlimited to the hearable domain, which is delimited by the sample rate
- Different types of ramps are applied to it so that the models learns different levels of sound amplitude :
- Increasing linear amplitude ramp
- Decreasing linear amplitude ramp
- Gaussian distribution amplitude ramp
- No ramp
White noise and sine sweeps are a great training material as they offer the possibility for the model to train on multi-harmonic signals. See Parker, Julian D. et al. “MODELLING OF NONLINEAR STATE-SPACE SYSTEMS USING A DEEP NEURAL NETWORK.”.
The dataset generation script will generate a folder dataset-n (where n is an integer incrementing by one) which will contain a set of inputs.
Example of command to generate a dataset :
$ python create_dataset.py --sample_rate=48000 --buffer_size=128 --max_buffer_amount=5A rule of thumb :
- Generate input training samples that last at least a second ; so for buffers of size 1024, with a sample rate of 48kHz,
$1024*46=48000$ , you can use--buffer_size 1024and--max_buffer_amount 46 - Use at least 300 cut-off frequencies
- Use the highest possible batch size your hardware can afford before getting out-of memory
Great result were obtained with a GRU model of hidden size 64 and of number of layers equal to 2. The amount of cutoff frequencies during training was 400, batch size was 8.
The sample rate used was 48kHz. The buffer size during training was 1024. The number of epochs was 50. The model was trained using a NVIDIA RTX 3050 Ti.
The obtained model produced the following bode plots :
Great result were obtained with a GRU model of hidden size 128 and of number of layers equal to 2. The amount of cutoff frequencies during training was 450, batch size was 8.
The sample rate used was 48kHz. The buffer size during training was 1024. The number of epochs was 50. The model was trained using a NVIDIA RTX PRO 5000 Blackwell.
The obtained model produced the following bode plots :
After sourcing the environment and generating a dataset, you can run the training script with :
$ python train.py --buffer_size=1024 --hidden_size=128 --num_layers=2 --filter_algo=butter --filter_type=lowpass --filter_order=2 --amount_of_fc=300 --batch_size=8A quick glance at the available parameters :
usage: train.py [-h] [--buffer_size BUFFER_SIZE] [--sample_rate SAMPLE_RATE] [--dataset DATASET] [--val_dataset VAL_DATASET] [--epochs EPOCHS]
[--hidden_size HIDDEN_SIZE] [--num_layers NUM_LAYERS] [--batch_size BATCH_SIZE] [--initial_lr INITIAL_LR] [--filter_algo {butter,cheby1}]
[--filter_type {lowpass,highpass}] [--filter_order FILTER_ORDER] [--cheby_ripple CHEBY_RIPPLE] [--amount_of_fc AMOUNT_OF_FC] [--notes NOTES]
options:
-h, --help show this help message and exit
--buffer_size BUFFER_SIZE
Amount of samples passed as input to the GRU model during forward pass
--sample_rate SAMPLE_RATE
Sample rate
--dataset DATASET Folder which should contain two subfolders ./inputs and ./expected, that represent the training dataset generated with
create_dataset.py
--val_dataset VAL_DATASET
Folder which should contain two subfolders ./inputs and ./expected, that represent the validation dataset generated with
create_dataset.py
--epochs EPOCHS The amount of epoch for training
--hidden_size HIDDEN_SIZE
The hidden size of the GRU
--num_layers NUM_LAYERS
The number of layers in the GRU
--batch_size BATCH_SIZE
The batch size in dataset (default is 8)
--initial_lr INITIAL_LR
The initial learning rate (default is 10**-3)
--filter_algo {butter,cheby1}
Filter's equation-based definition. Supported : ['butter', 'cheby1']
--filter_type {lowpass,highpass}
Filter's type. Supported : ['lowpass', 'highpass']
--filter_order FILTER_ORDER
Filter's order
--cheby_ripple CHEBY_RIPPLE
Chebyshev passband or stopband ripple value ; is ignored if specified for a filter that is not cheby1 or cheby2
--amount_of_fc AMOUNT_OF_FC
Expected signals will be filtered with frequencies ranging from 50 to 7500 Hz on a logarithmic scale. This option gives the amount
of cut-off frequencies to use. The more frequencies the better the generalization but the longer the training
--notes NOTES Notes about the training that will be contained in the checkpoints files
Run the Jupyter Notebook's first three cells. Then modify the fourth cell's parameters according to your model hyperparameters:
#### Adapt the following parameters to your model ####
hidden_size = 128
num_layers = 2
buffer_size = 512
sample_rate = 48000
filter_algo="cheby1"
order = 2
filter_type = "lowpass"
ripple = 5
######################################################Change the path to point to your model on line :
state = torch.load(<path to your model here>, map_location=torch.device("cpu"))And then run the cell.
Fifth and sixth cells are for testing the model on random white noise and sine sweeps.
- Finish README.md
- Add validation loss in training
- Implement other filters
- n-th order filters
- other types than lowpass
- other algo than butterworth
- Implement JUCE VST Plugin to use the model
- Try other inference engines
- RTNeural
- TF Lite
- Executorch
See the open issues for a full list of proposed features (and known issues).
Contributions are what make the open source community such an amazing place to learn, inspire, and create. Any contributions you make are greatly appreciated.
If you have a suggestion that would make this better, please fork the repo and create a pull request. You can also simply open an issue with the tag "enhancement". Don't forget to give the project a star! Thanks again!
- Fork the Project
- Create your Feature Branch (
git checkout -b feature/AmazingFeature) - Commit your Changes (
git commit -m 'Add some AmazingFeature') - Push to the Branch (
git push origin feature/AmazingFeature) - Open a Pull Request
Distributed under the MIT License. See LICENSE.txt for more information.
Nicolas Gry - grybouilli at outlook.fr
Project Link: https://github.com/grybouilli/GRU-IRR-Filter