Setup & configuration — Repository, commands and the YAML

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

Download the repository
Open the repository page, press the green Code button and choose Download as ZIP, then extract it anywhere convenient.
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.
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.

Repository layout after step 3

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

Config → .yaml editable

Everything about the model and how it trains: architecture dimensions, chunk size, learning rate, epochs, which instruments to target.

Where most of your tuning happens.
Lives at the repo root, next to train.py.

Weights → .ckpt millions of numbers

The model itself. You never edit it — you load it, train further, and let the trainer write new checkpoints.

Lives in results/.
One file, plus its matching config.

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

Two practical notes

Some browsers display a config file instead of downloading it. If that happens, open it in a text editor (Notepad++ works fine), copy the contents, and save the file yourself with a .yaml extension.
Every model published with the repository is open-source. Check the licence of the specific checkpoint before you build anything commercial on top of it.

Rofomers have their own loss

Mel-Band and BS Roformers train with an additional spectrogram loss rather than a pure waveform loss. That is why you never need to add one manually unless you are deliberately training for fullness — see Chapter 10.

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.

Local runs: use an administrator shell

On Windows, run these commands in a Command Prompt or PowerShell opened as administrator. Package installs and some file operations will fail silently or half-way without it.

Installation

powershell · one-time setup

# 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.

powershell · start training

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.

Make it a two-click run

Paste the command into Notepad, edit the values there, and save it as run.bat in the repository root. From then on you double-click it, or type .\run.bat in the shell. No more re-typing the argument list every time — and no more typos in paths.

Inference: using the model

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

powershell · inference

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.

Cloud only · OpenBLAS

On rented instances you may hit an OpenBLAS threading error during training. Fix it for the current session with:

bash · cloud

export OPENBLAS_NUM_THREADS=1

It only applies to the terminal session you run it in — see Chapter 12 for where it fits in the cloud sequence.

Argument reference — training

Argument	What it does
`--model_type`	Architecture to train: `mdx23c`, `htdemucs`, `segm_models`, `mel_band_roformer`, `bs_roformer`, `swin_upernet`, `bandit`.
`--config_path`	Path to the YAML config for that model.
`--start_check_point`	Initial checkpoint to start from (fine-tuning). Omit it to train from scratch.
`--results_path`	Folder for outputs — both `.ckpt` files and run metadata.
`--data_path`	Path to your dataset folder.
`--dataset_type`	`1`, `2`, `3` or `4` — which layout your data follows. See the dataset types documentation.
`--valid_path`	Path to the validation dataset folder.
`--num_workers`	How many CPU workers load and pre-process audio in parallel.
`--pin_memory`	Keeps host memory page-locked so transfers to the GPU are faster; worth enabling with several workers.
`--seed`	Seeds the randomness in the run — experiment with values for reproducibility.
`--device_ids`	List of GPU IDs to use; normally just `0`.
`--use_multistft_loss`	Multi-STFT loss (spectrogram based) — the setting behind fullness models.
`--use_mse_loss`	Default MSE loss, waveform based.
`--use_l1_loss`	L1 loss, waveform based.
`--wandb_key`	Weights & Biases API key, for live run dashboards.
`--pre_valid`	Runs a validation pass before training begins.
`--metrics`	Metrics to compute each validation pass: `sdr`, `l1_freq`, `si_sdr`, `neg_log_wmse`, `aura_stft`, `aura_mrstft`, `bleedless`, `fullness`.
`--metric_for_scheduler`	Which metric the learning-rate scheduler watches (same options as above).
`--train_lora`	Train with LoRA (Low-Rank Adaptation) instead of full weights.
`--lora_checkpoint`	Starting checkpoint for LoRA weights.

Argument reference — inference

Argument	What it does
`--model_type`	Architecture, matching the checkpoint you loaded.
`--config_path`	Path to the config file.
`--start_check_point`	The checkpoint to run.
`--input_folder`	Folder holding the mixtures you want separated.
`--store_dir`	Where the resulting stems are written.
`--draw_spectro`	Also renders spectrogram images of the results; the value sets how many seconds of the track to draw (default `0`, off).
`--device_ids`	List of GPU IDs to use.
`--extract_instrumental`	Inverts the vocal output to produce an instrumental (and vice versa).
`--disable_detailed_pbar`	Turns off the detailed progress bar.
`--force_cpu`	Forces CPU inference even when CUDA is available.
`--flac_file`	Write FLAC instead of WAV.
`--pcm_type`	Bit depth for FLAC output: `PCM_16` or `PCM_24`.
`--use_tta`	Test-time augmentation (polarity and channel inversion). Roughly triples runtime, reduces noise and slightly improves quality.
`--lora_checkpoint`	Starting checkpoint for LoRA weights.

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 worth knowing: the audio chunk size for conversions.

chunk_size — the one audio setting you will change

Roformer configs describe their audio window in dim_t (frames), while the trainer wants chunk_size (samples). Convert with:

The formula

chunk_size = (dim_t − 1) × hop_length

dim_t	chunk_size (samples)	Audio window
`256`	`112455`	2.55 s
`801`	`352800`	8.00 s
`1101`	`485100`	11.00 s
`1333`	`587412`	13.32 s

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 in memory.

Schematic · chunk windows drawn to scale

The one audio setting worth changing. Roformer configs describe their window in frames (dim_t); the trainer wants samples (chunk_size). Bars are scaled to a 16-second excerpt, so what you are looking at is how much musical context each setting sees.

Chunk size & VRAM calculator

Type in your numbers and see what window you are actually asking for, and whether your card can hold it. Nothing is sent anywhere — the arithmetic runs in the page.

This calculator needs JavaScript. Without it, the same arithmetic by hand: chunk_size = (dim_t − 1) × hop_length, and seconds = chunk_size ÷ sample_rate.

dim_t (frames) hop_length sample rate batch_size grad accum steps GPU VRAM

dim_t preset

chunk_size—

window—

effective batch—

estimated VRAM—

—

How this is derived. The conversion is exact: chunk_size = (dim_t − 1) × hop_length. The memory figure is an estimate scaled linearly from the one data point the guide has — a Roformer at batch_size: 4 with a 13.32 s window filling a 140 GB H200 — so it guesses that VRAM grows with batch size × window length. Real attention cost grows faster than the window does, which means wide windows land worse than this predicts; use_amp and use_torch_checkpoint move it the other way. Use it to decide what to try first, not to predict an exact number.

The training block

config.yaml · training section

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

Key	Meaning
`batch_size`	How many audio samples are used to update the weights. Higher values need more VRAM.
`gradient_accumulation_steps`	Simulates a larger batch without the VRAM cost: effective batch becomes `batch_size × steps`. Weights are not updated on the in-between steps.
`grad_clip`	Gradient clipping threshold. `0` leaves it off.
`instruments`	The stems inside your dataset. Must match the folders you actually created.
`lr`	Learning rate. `1.0e-05` is the common default; `5.0e-06` is the other frequent choice for fine-tuning.
`patience`	How many epochs the scheduler waits without improvement before dropping the learning rate.
`reduce_factor`	How much the learning rate is multiplied by when it drops (`0.95` = −5%).
`target_instrument`	The stem the model focuses on. Set to `null` when you are declaring stems explicitly.
`num_epochs`	How many times the full `num_steps` cycle runs. Can be changed when resuming from a checkpoint.
`num_steps`	Weight updates per epoch (divided by gradient accumulation). Validation runs each time this count is reached. Also resumable.
`optimizer`	Optimiser for training; `adam` is the default.
`ema_momentum`	Momentum of the exponential moving average of the weights — a smoother, usually more stable copy of the model.
`q`	Quantile-style parameter used by the loss weighting.
`coarse_loss_clip`	Clips extreme loss values early in training to avoid a first-epoch blow-up.
`other_fix`	Needed on multisong datasets to check that `other` really is an instrumental.
`use_amp`	Mixed-precision (float16) training. Usually should be `true` — it is a large speed and VRAM win.
`use_torch_checkpoint`	Gradient checkpointing: trades a little speed for notably lower VRAM use. Enable when you hit out-of-memory errors.

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.

Set inference overlap to 1 while training

num_overlap in the inference section of the config controls how many passes are averaged at inference time. During training it only costs time — always set it to 1 for the training run, and raise it later when you actually separate audio.

Optimal parameters take time

There is no magic config. Expect to find your numbers by iterating: change one thing, train, listen, compare metrics. Add use_torch_checkpoint if you genuinely need the VRAM saving — that is, when you are seeing CUDA_OutOfMemory errors.