HS-TasNet 0.3.1

HS TasNet

HS-TasNet

Implementation of HS-TasNet, "Real-time Low-latency Music Source Separation using Hybrid Spectrogram-TasNet", proposed by the research team at L-Acoustics

Install

$ pip install HS-TasNet

Usage

import torch
from hs_tasnet import HSTasNet

model = HSTasNet()

audio = torch.randn(1, 2, 204800) # ~5 seconds of stereo

separated_audios, _ = model(audio)

assert separated_audios.shape == (1, 4, 2, 204800) # second dimension is the separated tracks

With the Trainer

# model

from hs_tasnet import HSTasNet, Trainer

model = HSTasNet()

# trainer

trainer = Trainer(
    model,
    dataset = None,               # add your in-house Dataset
    concat_musdb_dataset = True,  # concat the musdb dataset automatically
    batch_size = 2,
    max_steps = 2,
    cpu = True,
)

trainer()

# after much training
# inferencing

model.sounddevice_stream(
    duration_seconds = 2,
    return_reduced_sources = [0, 2]
)

# or from the exponentially smoothed model (in the trainer)

trainer.ema_model.sounddevice_stream(...)

# or you can load from a specific checkpoint

model.load('./checkpoints/path.to.desired.ckpt.pt')
model.sounddevice_stream(...)

# to load an HS-TasNet from any of the saved checkpoints, without having to save its hyperparameters, just run

model = HSTasNet.init_and_load_from('./checkpoints/path.to.desired.ckpt.pt')

Training script

First make sure dependencies are there by running

$ sh scripts/install.sh

Then make sure uv is installed

$ pip install uv

Finally run the following to train a newly initialized model on a small subset of MusDB, and make sure the loss goes down

$ uv run train.py

For distributed training, you just need to run accelerate config first, courtesy of accelerate from 🤗 but single machine is fine too

Experiment tracking

To enable online experiment monitoring / tracking, you need to have wandb installed and logged in

$ pip install wandb && wandb login

Then

$ uv run train.py --use-wandb

To wipe the previous checkpoints and evaluated results, append --clear-folders

Alternative RNNs

The architecture defaults to using PyTorch's LSTM (or GRU), but you can easily substitute it for any other module by passing an rnn_klass to the HSTasNet constructor, as long as it adheres to a specific interface (read alternative_rnns.py)

For example, to use the minGRU architecture:

import torch
from hs_tasnet import HSTasNet
from hs_tasnet.alternative_rnns import minGRUWrapper

model = HSTasNet(rnn_klass = minGRUWrapper)

audio = torch.randn(1, 2, 204800)
separated_audios, _ = model(audio)

Test

$ uv pip install '.[test]' --system

Then

$ pytest tests

Sponsors

This open sourced work is sponsored by Sweet Spot

Citations

@misc{venkatesh2024realtimelowlatencymusicsource,
    title    = {Real-time Low-latency Music Source Separation using Hybrid Spectrogram-TasNet},
    author   = {Satvik Venkatesh and Arthur Benilov and Philip Coleman and Frederic Roskam},
    year     = {2024},
    eprint   = {2402.17701},
    archivePrefix = {arXiv},
    primaryClass = {eess.AS},
    url      = {https://arxiv.org/abs/2402.17701},
}

@inproceedings{Feng2024WereRA,
    title   = {Were RNNs All We Needed?},
    author  = {Leo Feng and Frederick Tung and Mohamed Osama Ahmed and Yoshua Bengio and Hossein Hajimirsadegh},
    year    = {2024},
    url     = {https://api.semanticscholar.org/CorpusID:273025630}
}