Setup & configuration

The training repository

One repository does the work: ZFTurbo/Music-Source-Separation-Training, usually shortened to MSST. It trains, validates and runs inference on the same architectures that UVR and the separation web apps use.

Get the code

1. Download the repository

   Open the repository page, press the green Code button and choose Download as ZIP, then extract it anywhere convenient.

2. Pick the model you want to train

   Scroll to the Pre-trained models section, open the List of Pre-trained models, and grab the Config and the Weights for your chosen architecture — for example a Mel-Band Roformer, or Kimberley Jensen's Roformer as a base.

3. Place the two files correctly

   The checkpoint goes into a folder named results; the YAML goes at the repository root, next to train.py, valid.py and inference.py.

Music-Source-Separation-Training/
├── train.py
├── valid.py
├── inference.py
├── requirements.txt
├── config_vocals_mel_band_roformer.yaml   ← the config, at the root
└── results/
    └── model_mel_band_roformer.ckpt          ← the weights

Config vs. checkpoint

If you have used UVR before, you already know these two files: this is exactly the pairing it ships with.

Commands & arguments

Everything runs through Python scripts, so everything is a command with arguments. Copy the templates, replace the placeholders, and keep them somewhere you can paste from.

Installation

# 1. Install PyTorch — pick the command for your CUDA version
# https://pytorch.org/get-started/locally/
pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu124

# 2. Install the repository requirements (from the repo folder)
cd C:\path\to\Music-Source-Separation-Training
pip install -r requirements.txt

# 3. Install anything the previous step skipped
pip install name-of-missing-package

The training command

This is the command you will run the most. Every highlighted value has to be replaced with your own paths and choices.

python train.py
  --model_type mel_band_roformer
  --config_path config_vocals.yaml
  --results_path results/
  --data_path C:\path\to\dataset
  --dataset_type 2
  --num_workers 4
  --device_ids 0
  --start_check_point results/model.ckpt
  --valid_path C:\path\to\validation
  --metric_for_scheduler sdr
  --metrics sdr fullness bleedless

The valid options for --metrics and --metric_for_scheduler are listed in the argument tables below; what each one measures is explained in Appendix B.

Inference: using the model

Once you have a checkpoint you are happy with, inference separates files with it.

python inference.py
  --model_type mel_band_roformer
  --config_path config_vocals.yaml
  --start_check_point results/model.ckpt
  --input_folder input
  --store_dir separation_results

Create two folders first: input, where you drop the files you want separated, and separation_results, where the script writes the stems.

export OPENBLAS_NUM_THREADS=1

Argument reference — training

Argument reference — inference

Configuring the YAML

The config is where the run is actually shaped. Its defaults are sane, but the training block is the part worth understanding — and the part worth experimenting with.

A config has four sections: audio, model, training and inference. You rarely touch the first two, with one exception: the audio chunk size. Which block owns the rest of the spectrum settings — n_fft, hop_length, dim_f — depends on the architecture, and there is a table of who owns what below the calculator.

chunk_size — the one audio setting you will change

The trainer reads the window itself from audio.chunk_size, in samples — that key is the one every architecture family reads. Configs also describe the same window in frames (dim_t), and because every STFT here runs with center=True the two are one value in two units:

The conversions above assume hop_length: 441 — the step every published Roformer in this guide uses. A different hop gives a different sample count for the same dim_t, because the frame count is what follows from the window: floor(chunk_size ÷ hop_length) + 1. So the hop decides both the window and the memory, and it belongs to the model rather than to you — the hop_length section below the calculator has the details. Smaller chunks train faster and use less VRAM; larger chunks teach the model longer musical context. This is the main dial you change when a run will not fit.

hop_length and n_fft — which block of the config owns them

hop_length is the step between STFT frames, and it is not a free preference: it is baked into the spectrogram the weights were trained on, so change it and the checkpoint you are fine-tuning no longer matches its own input. Where you change it, though, depends on the architecture — and this is the mistake that costs people an afternoon, because a Roformer config and an mdx23c config both have an n_fft and a hop_length, in different blocks, and only one of the two counts in each.

The mechanism is plain in the source. mel_band_roformer and bs_roformer are constructed as MelBandRoformer(**config.model), and SCNet as SCNet(**config.model) — the model block is the only thing they see, and their STFT parameters are named there. mdx23c and DTTNet go the other way and are handed the whole audio block (STFT(config.audio)), so for them audio.n_fft, audio.hop_length and audio.dim_f are the real ones. What every family reads from audio is chunk_size, num_channels, sample_rate and min_mean_abs.

Two rules account for every config in this ecosystem, and which one is right depends entirely on where your config came from:

Every published Roformer checkpoint this guide tells you to fine-tune uses 441 — a tenth of the sample rate, so a frame lands exactly on 10 ms of audio at any full-band rate. The blank templates that ship with MSST, and the DTTNet config, use n_fft ÷ 4 — the ordinary 75 % overlap you would pick for a fresh STFT — and mdx23c halves that again. So “keep hop_length = n_fft ÷ 4” is good advice for a config you are writing from scratch and the wrong advice for a checkpoint you are continuing: load Kimberley Jensen's model with hop_length: 512 and its learned features no longer line up with the input you feed them.

That is why the calculator's hop rule is a choice rather than a constant, and why it carries an n_fft field: both STFT rules are measured from the window, and dim_f — the width the transformer actually reads — is half of it, so n_fft: 2048 is dim_f: 1024, 4096 is 2048. A wider window is therefore also a wider spectrogram and a heavier model, which is the one part of the memory estimate the calculator still does not scale for.

Because the hop decides how many frames the window divides into, and the frame count is what the model attends over, the hop is a memory setting as much as a spectral one. At a fixed chunk_size the relationship is exact: frames = floor(chunk_size ÷ hop_length) + 1, so chunk_size: 587412 is 1333 frames at hop_length: 441 and 1148 frames at 512 — the same 13.3 seconds of audio, seen through far fewer, wider steps. That is the direction the calculator's hop rule moves in: it holds the window and lets the frame count follow. Run it the other way — hold dim_t and let the window grow — and dim_t: 801 is an 8.00 s window at 441 but a 9.29 s one at 512. Which of the two you want depends on what you are doing: keeping chunk_size is how you switch conventions without changing the music the model hears, and keeping dim_t is what a fine-tune does when it inherits the frames its checkpoint was trained on.

Sample rate — the band each choice keeps

The sample rate is not a memory dial, but it does decide something you cannot undo later: everything above half the rate — the Nyquist frequency — is gone before training starts. A rate of 22 050 Hz keeps 0–11 025 Hz and throws away cymbals, sibilance and guitar harmonics: free for a bass model, fatal for a drum model.

Halving the rate halves what the loader reads and what the STFT transforms, and whether it halves the memory depends entirely on the hop — because frames = chunk_size ÷ hop_length + 1. At a fixed chunk_size a halved rate makes the same chunk twice as long, so you get twice the frames and the estimate doubles; halve chunk_size at the same time and it comes straight back to where it started. That is exactly what hop_length = rate ÷ 100 arranges without you thinking about it: frames per second = rate ÷ hop_length is 100 at every full-band rate, so the frame count follows the seconds of audio and nothing else. Under a fixed hop it does not: at hop_length: 512 a halved rate halves the frames per second, so the same 13.3 s window arrives as 574 frames instead of 1148 — half the memory, and each frame now steps 23.2 ms instead of 11.6 ms. One arithmetic wrinkle, and only under the rate ÷ 100 rule: 22 050 ÷ 100 is not a whole number, so the hop lands on 221 rather than 220.5 — 10.02 ms per frame instead of 10, which moves the window by tens of milliseconds — while 24 000 ÷ 100 = 240 lands exactly. Switch the rule to n_fft ÷ 4 and the wrinkle disappears: 512 at any rate.

Channels — stereo is architecture, not file format

Every Roformer published for this ecosystem is stereo, and the config says so twice: audio.num_channels: 2 and model.stereo: true. Those lines describe the model, not your files. The band-split front end turns each mel band's two channels into one longer sequence for the transformer to attend over, which is why a stereo run costs more than a mono one — and why the calculator now asks which you are training.

The distinction matters in both directions. If your source material is mono, you have not made training any cheaper: a stereo model still expects two channels, so the mono file arrives duplicated. And if you genuinely want a mono model, you change the architecture — num_channels: 1 together with stereo: false, trained on mono mixtures and targets — which halves the sequence and so most of the activation memory.

The training block

training:
  batch_size: 1
  gradient_accumulation_steps: 1
  grad_clip: 0
  instruments:
    - vocals
    - other
  lr: 1.0e-05
  patience: 2
  reduce_factor: 0.95
  target_instrument: vocals
  num_epochs: 1000
  num_steps: 1000
  optimizer: adam
  ema_momentum: 0.999
  q: 0.95
  coarse_loss_clip: false
  other_fix: true
  use_amp: true
  use_torch_checkpoint: true
  augmentation: false
  augmentation_type: null
  use_mp3_compress: false
  augmentation_mix: false
  augmentation_loudness: false
  augmentation_loudness_type: 1
  augmentation_loudness_min: 0
  augmentation_loudness_max: 0

Augmentation options

• augmentation — master switch for audiomentations and pedalboard augmentations.
• augmentation_type — which augmentation profile to use.
• augmentation_mix — mixes several stems of the same type with some probability, useful when your dataset is small.
• augmentation_loudness, _type, _min, _max — randomly change the loudness of each stem, with the range you specify.
• use_mp3_compress — deprecated; leave it false.