Mini-BS-RoFormer-V2.5
Model for the music source separation task, decomposes audio into four stems: bass, drums, other, and vocals.
针对音乐音频分离的 BS-RoFormer 模型，将音频分离为 bass、drums、other、vocals。
Model Details
- 模型参数量，46.2M
- 权重精度，bf16
- 采样率，44100 Hz
- 音频分块长度，8 s
在 MUSDB18HQ 数据的 val 集上的性能（单位 SDR，越高越好）：
| tracks | Mini-BS-RoFormer-V2.5-46.2M | Mini-BS-RoFormer-V2-46.8M |
|---|---|---|
| overall | 10.09 | 10.03 |
| bass | 9.71 | 9.68 |
| drums | 10.48 | 10.58 |
| other | 9.19 | 8.99 |
| vocal | 10.97 | 10.86 |
Uses
使用的 transformers 库版本为 5.18.0，PyTorch 版本为 2.14.1。为了正常运行模型还需要安装库 soundfile、einops 和 librosa。
GPU 推理：
from transformers import AutoModel
import soundfile
import torch
import librosa
model_name = "HiDolen/Mini-BS-RoFormer-V2.5-46.2M"
model = AutoModel.from_pretrained(
model_name,
trust_remote_code=True,
)
model.to("cuda")
# 加载音频
file = "./input.wav"
waveform, sr = librosa.load(file, sr=44100, mono=False)
waveform = torch.tensor(waveform).float()
waveform = waveform.to("cuda")
# 进行推理
result = model.separate(
waveform,
batch_size=2,
verbose=True,
)
# 保存处理结果
for i in range(result.shape[0]):
soundfile.write(f"separated_stem_{i}.wav", result[i].cpu().numpy().T, 44100)
以上代码会分离出 bass、drums、other 和 vocal 四个轨道。
若想只分离人声和伴奏两轨，在最后保存音频时合并即可：
···
result = model.separate(
waveform,
batch_size=2,
verbose=True,
)
# 合并 bass、drums、other 作为伴奏
instrumental = result[0] + result[1] + result[2]
vocals = result[3]
result = torch.stack([instrumental, vocals], dim=0)
for i in range(result.shape[0]):
soundfile.write(f"separated_stem_{i}.wav", result[i].cpu().numpy().T, 44100)
Differences from the previous version
Mini-BS-RoFormer-V2.5 相比于 V2 的主要改进：
- 损失改为频域听阈加权 L1
- 利用注意力池化达成各 stem 共享 mask 解码器
- 去掉 MUSDB18HQ 数据集中的坏数据进行训练
- 使用更宽更浅的模型超参配置
Training Details
使用 MUSDB18HQ 数据集的 train 和 test 集进行训练。
学习率 3e-4，以 batch_size=8 训练 246k 步。
对 transformer 层使用 Muon 优化器，其他网络层使用 AdamW 优化器。
Acknowledgments
- https://github.com/lucidrains/BS-RoFormer
- https://arxiv.org/abs/2309.02612 (Music Source Separation with Band-Split RoPE Transformer)