Training runs

Training from scratch

You can train with no checkpoint at all. It is the right move when fine-tuning simply cannot give you the model you need — and it is also the slowest, least forgiving path.

Fine-tuning adjusts an existing model toward your data. Training from scratch builds the representation itself, which is what you want when your target is far away from anything published. In the source guide this was the route taken for a personal metal dataset: pre-train on a laptop to get something decent, then move to cloud GPUs with different parameters to push it further.

Shifting a model's focus between stems

A special case worth knowing: taking a model that separates vocals and turning it into one that separates instrumentals — without retraining everything from the ground up.

This one came from a real problem: needing to move a vocal model's target over to the instrumental. After a series of unsuccessful attempts, unwa suggested a transfer-learning approach and supplied the code for it.

The patch

for param in model.parameters():
    param.requires_grad = False

for param in model.mask_estimators.parameters():
    param.requires_grad = True

Place it just above line 203 in train.py, where the optimiser parameters are assembled:

optim_params = dict()

Figure 9.1 · the patched section in train.py

Why this works

unwa's reasoning: "in my experience, it is the MaskEstimator that changes significantly when the target is changed." Freezing the Mel-band projections and the Roformer blocks and training only the final masking modules is enough to move the model from one target to another — and it cuts training time dramatically.

Figure 9.2 · Mel-Roformer architecture

How it actually went

The run was taken to epoch 12 to complete the switch, plus 2 further epochs to confirm the model stayed on the instrumental target — the second stretch purely to be sure it was not about to drift back.

Figure 9.3 · over-cooked run

Figure 9.4 · the same run once settled

Fullness models and model fusion

Fullness measures how complete the target stem sounds. Training a model specifically for it takes three steps — and produces a specific, important artefact.

1. Patch the architecture switch in train.py

   Find the line beginning if args.model_type in ['mel_band_roformer', 'bs_roformer']: and place a 1 before either mel or bs — or search the file for anything containing roformer.

2. Use the multi-STFT loss

   Pass --use_multistft_loss as your loss argument on the training command.

3. Add the loss block to your config

   These values are a working starting point; they can be adjusted.

loss_multistft:
  fft_sizes:
    - 1024
    - 2048
    - 4096
  hop_sizes:
    - 512
    - 1024
    - 2048
  win_lengths:
    - 1024
    - 2048
    - 4096
  window: "hann_window"
  scale: "mel"
  n_bins: 128
  sample_rate: 44100
  perceptual_weighting: false
  w_sc: 0.0
  w_log_mag: 1.0
  w_lin_mag: 0.0
  w_phs: 0.0
  mag_distance: "L1"

Model fusion (Sucial's script)

Not strictly training, but part of the same toolbox: a fusion script that merges two checkpoints into one, available at huggingface.co/Sucial/…/model_fusion.py. Download the .py file and drop it into the MSST folder.

Replace the placeholder checkpoint names in the script with the actual YAML names of the checkpoints you want to combine.

Running the training locally

Local training is the simplest loop there is: open a shell, paste the command, let it run. Your GPU will spin hard — that is how you know it is working.

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

Figure 11.1 · console output during a local run

1 2 3 4 5 6

1. Epoch and learning rate — which pass this is, and the rate in force. It moves when the scheduler reacts to a stalled metric.
2. Training loss — how far this epoch's output is from the target. You want it trending down; the avg_loss line next to it is smoother than the instantaneous number.
3. Num overlap — how many inference passes were averaged to produce the numbers below. Set it to 1 while training.
4. Per-stem metric summary — the averaged verdict for the target stem: fullness, bleedless, l1_freq (and sdr if you requested it). This is the block to compare with the previous epoch.
5. Elapsed time — wall-clock seconds for the epoch, which is how you work out what the rest of the run will cost.
6. Store weights — where the checkpoint was written. This is the file you load later to separate audio.

How to read the log

While it runs

Training is a waiting game. Locally you cannot do anything GPU-heavy alongside it — no games, no renders, no video export — so plan something else for the hours. That changes on cloud hardware, where you are paying for the privilege of not caring about the machine.

Running the training in the cloud

A step-by-step walkthrough for renting a GPU on vast.ai. It is genuinely unintuitive, and the existing documentation is not written for this use case. Follow the steps in order.

• Step 01Upload to DriveRepo plus both datasets, with at least 100 GB free
• Step 02Connect the drivevast.ai → CLOUD → Connect Google Drive
• Step 03Rent a GPUPyTorch (cuDNN) template, CUDA 12.4, SSH + Jupyter
• Step 04Pull the dataCreate a dump folder, then download into /TRAINING_STUFF
• Step 05Open a terminalcd, install requirements, OpenBLAS export, train
• Step 06Training startsOnce “Collecting Metadata” appears, you are running

1. Put everything on Google Drive

   Upload the training repository and both datasets into a folder of your choosing, and make sure you have at least 100 GB of drive space free. vast.ai offers a direct connection to Drive under CLOUD on the site (below INSTANCES): click Connect Google Drive, name the connection, and authorise it.

2. Rent a GPU

   In the SEARCH section, pick the PyTorch (cuDNN Runtime) template with CUDA 12.4, SSH and Jupyter. Give the instance enough disk space — it is included in the hourly price, so larger storage costs more per hour.

3. Create a dumping folder, then pull your data down

   Once the instance appears under INSTANCES, wait for it to finish loading and enter it. Create a folder for the data (in the source guide it was simply called TRAINING_STUFF). Do not close the page: go back to INSTANCES and click the cloud icon on the instance — a message appears about uploading or downloading data from cloud providers. Enter the source folder and the destination path, for example /TRAINING_STUFF, and let it run.

4. Open a terminal and adapt the training command

   When the transfer finishes, your data is visible in the Jupyter file browser. Take the training command from earlier in the guide, point --data_path and --valid_path at the uploaded folders, and — this is the part people get wrong — put a slash both before and after the path:

   path syntax on the instanceCopy

   --data_path /TRAINING_STUFF/DATA/2stem_metal_dataset/
   --valid_path /TRAINING_STUFF/DATA/validation/

   Then open File → New → Terminal in Jupyter.

5. Run the four commands in sequence

   In the terminal, one after the other:

   • change directory into the repository
   • install the requirements
   • the OpenBLAS export (it surfaced at the metadata stage in the source run, and turned out to be unrelated to the actual problem — but it costs nothing to set)
   • the training command
6. Install the missing modules as they appear

   Expect No module named X. Install it, paste the training command again, repeat until you see the metadata collection stage — at which point training starts properly. The modules that came up in the source run: wandb, soundfile, auraloss, audiomentations, pedalboard, ml_collections, omegaconf, einops, beartype, rotary_embedding_torch.

ERROR: Failed to build installable wheels for some pyproject.toml
based projects (wxpython, diffq, pesq)