Home

Awesome

Wave-U-Net (Pytorch)

<a href="https://replicate.ai/f90/wave-u-net-pytorch"><img src="https://img.shields.io/static/v1?label=Replicate&message=Demo and Docker Image&color=darkgreen" height=20></a>

Improved version of the Wave-U-Net for audio source separation, implemented in Pytorch.

Click here for the original Wave-U-Net implementation in Tensorflow. You can find more information about the model and results there as well.

Improvements

Installation

GPU strongly recommended to avoid very long training times.

Option 1: Direct install (recommended)

System requirements:

Clone the repository:

git clone https://github.com/f90/Wave-U-Net-Pytorch.git

Recommended: Create a new virtual environment to install the required Python packages into, then activate the virtual environment:

virtualenv --python /usr/bin/python3.6 waveunet-env
source waveunet-env/bin/activate

Install all the required packages listed in the requirements.txt:

pip3 install -r requirements.txt

Option 2: Singularity

We also provide a Singularity container which allows you to avoid installing the correct Python, CUDA and other system libraries, however we don't provide specific advice on how to run the container and so only do this if you have to or know what you are doing (since you need to mount dataset paths to the container etc.)

To pull the container, run

singularity pull shub://f90/Wave-U-Net-Pytorch

Then run the container from the directory where you cloned this repository to, using the commands listed further below in this readme.

Download datasets

To directly use the pre-trained models we provide for download to separate your own songs, now skip directly to the last section, since the datasets are not needed in that case.

To start training your own models, download the full MUSDB18HQ dataset and extract it into a folder of your choice. It should have two subfolders: "test" and "train" as well as a README.md file.

You can of course use your own datasets for training, but for this you would need to modify the code manually, which will not be discussed here. However, we provide a loading function for the normal MUSDB18 dataset as well.

Training the models

To train a Wave-U-Net, the basic command to use is

python3.6 train.py --dataset_dir /PATH/TO/MUSDB18HQ 

where the path to MUSDB18HQ dataset needs to be specified, which contains the train and test subfolders.

Add more command line parameters as needed:

For more config options, see train.py.

Training progress can be monitored by using Tensorboard on the respective log_dir. After training, the model is evaluated on the MUSDB18HQ test set, and SDR/SIR/SAR metrics are reported for all instruments and written into both the Tensorboard, and in more detail also into a results.pkl file in the checkpoint_dir

<a name="test"></a> Test trained models on songs!

We provide the default model in a pre-trained form as download so you can separate your own songs right away.

Downloading our pretrained models

Download our pretrained model here. Extract the archive into the checkpoints subfolder in this repository, so that you have one subfolder for each model (e.g. REPO/checkpoints/waveunet)

Run pretrained model

To apply our pretrained model to any of your own songs, simply point to its audio file path using the input_path parameter:

python3.6 predict.py --load_model checkpoints/waveunet/model --input "audio_examples/Cristina Vane - So Easy/mix.mp3"

By default, output is written where the input music file is located, using the original file name plus the instrument name as output file name. Use --output to customise the output directory.

To run your own model: