Skip to content

Latest commit

 

History

History
290 lines (207 loc) · 12.9 KB

training.md

File metadata and controls

290 lines (207 loc) · 12.9 KB
Raw

Training (Hybrid) Demucs

Install all the dependencies

You should install all the dependencies either with either Anaconda (using the env file environment-cuda.yml ) or pip, with requirements.txt.

Datasets

MusDB HQ

Note that we do not support MusDB non HQ training anymore. Get the Musdb HQ dataset, and update the path to it in two places:

  • The dset.musdb key inside conf/config.yaml.
  • The variable MUSDB_PATH inside tools/automix.py.

Create the fine tuning datasets

This is only for the MDX 2021 competition models

I use a fine tuning on a dataset crafted by remixing songs in a musically plausible way. The automix script will make sure that BPM, first beat and pitches are aligned. In the file tools/automix.py, edit OUTPATH to suit your setup, as well as the MUSDB_PATH to point to your copy of MusDB HQ. Then run

export NUMBA_NUM_THREADS=1; python3 -m tools.automix

Important: the script will show many errors, those are normals. They just indicate when two stems do not batch due to BPM or music scale difference.

Finally, edit the file conf/dset/auto_mus.yaml and replace dset.wav to the value of OUTPATH.

If you have a custom dataset, you can also uncomment the lines dset2 = ... and dset3 = ... to add your custom wav data and the test set of MusDB for Track B models. You can then replace the paths in conf/dset/auto_extra.yaml, conf/dset/auto_extra_test.yaml and conf/dset/aetl.yaml (this last one was using 10 mixes instead of 6 for each song).

Dataset metadata cache

Datasets are scanned the first time they are used to determine the files and their durations. If you change a dataset and need a rescan, just delete the metadata folder.

A short intro to Dora

I use Dora for all the of experiments (XPs) management. You should have a look at the Dora README to learn about the tool. Here is a quick summary of what to know:

  • An XP is a unique set of hyper-parameters with a given signature. The signature is a hash of those hyper-parameters. I will always refer to an XP with its signature, e.g. 9357e12e. We will see after that you can retrieve the hyper-params and re-rerun it in a single command.
  • In fact, the hash is defined as a delta between the base config and the one obtained with the config overrides you passed from the command line. This means you must never change the conf/**.yaml files directly., except for editing things like paths. Changing the default values in the config files means the XP signature won't reflect that change, and wrong checkpoints might be reused. I know, this is annoying, but the reason is that otherwise, any change to the config file would mean that all XPs ran so far would see their signature change.

Dora commands

Run tar xvf outputs.tar.gz. This will initialize the Dora XP repository, so that Dora knows which hyper-params match the signature like 9357e12e. Once you have done that, you should be able to run the following:

dora info -f 81de367c  # this will show the hyper-parameter used by a specific XP.
                       # Be careful some overrides might present twice, and the right most one
                       # will give you the right value for it.
dora run -d -f 81de367c   # run an XP with the hyper-parameters from XP 81de367c.
                          # `-d` is for distributed, it will use all available GPUs.
dora run -d -f 81de367c hdemucs.channels=32  # start from the config of XP 81de367c but change some hyper-params.
                                             # This will give you a new XP with a new signature (here 3fe9c332).

An XP runs from a specific folder based on its signature, by default under the outputs/ folder. You can safely interrupt a training and resume it, it will reuse any existing checkpoint, as it will reuse the same folder. If you made some change to the code and need to ignore a previous checkpoint you can use dora run --clear [RUN ARGS].

If you have a Slurm cluster, you can also use the dora grid command, e.g. dora grid mdx. Please refer to the Dora documentation for more information.

Hyper parameters

Have a look at conf/config.yaml for a list of all the hyper-parameters you can override. If you are not familiar with Hydra, go checkout their page to be familiar with how to provide overrides for your trainings.

Model architecture

A number of architectures are supported. You can select one with model=NAME, and have a look in [conf/config.yaml'(../conf/config.yaml) for each architecture specific hyperparams. Those specific params will be always prefixed with the architecture name when passing the override from the command line or in grid files. Here is the list of models:

  • demucs: original time-only Demucs.
  • hdemucs: Hybrid Demucs (v3).
  • torch_hdemucs: Same as Hybrid Demucs, but using torchaudio official implementation.
  • htdemucs: Hybrid Transformer Demucs (v4).

Storing config in files

As mentioned earlier, you should never change the base config files. However, you can use Hydra config groups in order to store variants you often use. If you want to create a new variant combining multiple hyper-params, copy the file conf/variant/example.yaml to conf/variant/my_variant.yaml, and then you can use it with

dora run -d variant=my_variant

Once you have created this file, you should not edit it once you have started training models with it.

Fine tuning

If a first model is trained, you can fine tune it with other settings (e.g. automix dataset) with

dora run -d -f 81de367c continue_from=81de367c dset=auto_mus variant=finetune

Note that you need both -f 81de367c and continue_from=81de367c. The first one indicates that the hyper-params of 81de367c should be used as a starting point for the config. The second indicates that the weights from 81de367c should be used as a starting point for the solver.

Model evaluation

Your model will be evaluated automatically with the new SDR definition from MDX every 20 epochs. Old style SDR (which is quite slow) will only happen at the end of training.

Model Export

In order to use your models with other commands (such as the demucs command for separation) you must export it. For that run

python3 -m tools.export 9357e12e [OTHER SIGS ...]  # replace with the appropriate signatures.

The models will be stored under release_models/. You can use them with the demucs separation command with the following flags:

demucs --repo ./release_models -n 9357e12e my_track.mp3

Bag of models

If you want to combine multiple models, potentially with different weights for each source, you can copy demucs/remote/mdx.yaml to ./release_models/my_bag.yaml. You can then edit the list of models (all models used should have been exported first) and the weights per source and model (list of list, outer list is over models, inner list is over sources). You can then use your bag of model as

demucs --repo ./release_models -n my_bag my_track.mp3

Model evaluation

You can evaluate any pre-trained model or bag of models using the following command:

python3 -m tools.test_pretrained -n NAME_OF_MODEL [EXTRA ARGS]

where NAME_OF_MODEL is either the name of the bag (e.g. mdx, repro_mdx_a), or a single Dora signature of one of the model of the bags. You can pass EXTRA ARGS to customize the test options, like the number of random shifts (e.g. test.shifts=2). This will compute the old-style SDR and can take quite bit of time.

For custom models that were trained locally, you will need to indicate that you wish to use the local model repositories, with the --repo ./release_models flag, e.g.,

python3 -m tools.test_pretrained --repo ./release_models -n my_bag

API to retrieve the model

You can retrieve officially released models in Python using the following API:

from demucs import pretrained
from demucs.apply import apply_model
bag = pretrained.get_model('htdemucs')    # for a bag of models or a named model
                                          # (which is just a bag with 1 model).
model = pretrained.get_model('955717e8')  # using the signature for single models.

bag.models                       # list of individual models
stems = apply_model(model, mix)  # apply the model to the given mix.

Model Zoo

Hybrid Transformer Demucs

The configuration for the Hybrid Transformer models are available in:

dora grid mmi --dry_run --init
dora grid mmi_ft --dry_run --init  # fined tuned on each sources.

We release in particular 955717e8, Hybrid Transformer Demucs using 5 layers, 512 channels, 10 seconds training segment length. We also release its fine tuned version, with one model for each source f7e0c4bc, d12395a8, 92cfc3b6, 04573f0d (drums, bass, other, vocals). The model 955717e8 is also named htdemucs, while the bag of models is provided as htdemucs_ft.

We also release 75fc33f5, a regular Hybrid Demucs trained on the same dataset, available as hdemucs_mmi.

Models from the MDX Competition 2021

Here is a short descriptions of the models used for the MDX submission, either Track A (MusDB HQ only) or Track B (extra training data allowed). Training happen in two stage, with the second stage being the fine tunining on the automix generated dataset. All the fine tuned models are available on our AWS repository (you can retrieve it with demucs.pretrained.get_model(SIG)). The bag of models are available by doing demucs.pretrained.get_model(NAME) with NAME begin either mdx (for Track A) or mdx_extra (for Track B).

Track A

The 4 models are:

  • 0d19c1c6: fine-tuned on automix dataset from 9357e12e
  • 7ecf8ec1: fine-tuned on automix dataset from e312f349
  • c511e2ab: fine-tuned on automix dataset from 81de367c
  • 7d865c68: fine-tuned on automix dataset from 80a68df8

The 4 initial models (before fine tuning are):

  • 9357e12e: 64ch time domain only improved Demucs, with new residual branches, group norm, and singular value penalty.
  • e312f349: 64ch time domain only improved, with new residual branches, group norm, and singular value penalty, trained with a loss that focus only on drums and bass.
  • 81de367c: 48ch hybrid model , with residual branches, group norm, singular value penalty penalty and amplitude spectrogram.
  • 80a68df8: same as b5559babb but using CaC and different random seed, as well different weigths per frequency bands in outermost layers.

The hybrid models are combined with equal weights for all sources except for the bass. 0d19c1c6 (time domain) is used for both drums and bass. 7ecf8ec1 is used only for the bass.

You can see all the hyper parameters at once with (one common line for all common hyper params, and then only shows the hyper parameters that differs), along with the DiffQ variants that are used for the mdx_q models:

dora grid mdx --dry_run --init
dora grid mdx --dry_run --init

Track B

  • e51eebcc
  • a1d90b5c
  • 5d2d6c55
  • cfa93e08

All the models are 48ch hybrid demucs with different random seeds. Two of them are using CaC, and two are using amplitude spectrograms with masking. All the models are combined with equal weights for all sources.

Things are a bit messy for Track B, there was a lot of fine tuning over different datasets. I won't describe the entire genealogy of models here, but all the information can be accessed with the dora info -f SIG command.

Similarly you can do (those will contain a few extra lines, for training without the MusDB test set as training, and extra DiffQ XPs):

dora grid mdx_extra --dry_run --init

Reproducibility and Ablation

I updated the paper to report numbers with a more homogeneous setup than the one used for the competition. On MusDB HQ, I still need to use a combination of time only and hybrid models to achieve the best performance. The experiments are provided in the grids repro.py and repro_ft._py for the fine tuning on the realistic mix datasets.

The new bag of models reaches an SDR of 7.64 (vs. 7.68 for the original track A model). It uses 2 time only models trained with residual branches, local attention and the SVD penalty, along with 2 hybrid models, with the same features, and using CaC representation. We average the performance of all the models with the same weight over all sources, unlike what was done for the original track A model. We trained for 600 epochs, against 360 before.

The new bag of model is available as part of the pretrained model as repro_mdx_a. The time only bag is named repro_mdx_a_time_only, and the hybrid only repro_mdx_a_hybrid_only. Checkout the paper for more information on the training.