CoolFace
Modelpublic

HiDolen/Mini-BS-RoFormer-18M

sourceHugging Facemitupdated 9mo agoView on Hugging Face
6likes114downloads
Model Card

Model Card for Model ID

Model for the Music source separation task. Its implementation is referenced to the existing BS-RoFormer code.

针对音乐音频分离任务的模型。改编自 现有的 BS-RoFormer 模型代码。

Model Details

模型参数：

  • —depth = 8
  • —hidden_size = 256
  • —intermediate_size = 256 * 3

总参数量 17.9M，在 MUSDB18HQ 数据的 val 集上达到平均 SDR 9.0 的性能。分轨具体 SDR：

  • —bass，8.31
  • —drums，9.55
  • —other，8.14
  • —vocal，10.03

Uses

使用的 transformers 库版本为 4.55.4。为了正常运行模型还需要安装库 soudfile、einops 和 librosa。

GPU 推理：

python
from transformers import AutoModel
import soundfile
import torch
import librosa

model_name = "HiDolen/Mini-BS-RoFormer-18M"
model = AutoModel.from_pretrained(
    model_name,
    trust_remote_code=True,
)
model.to("cuda")

# 加载音频
file = "./Bruno Mars - Runaway Baby.mp3"
waveform, sr = librosa.load(file, sr=44100, mono=False)
waveform = torch.tensor(waveform).float()
waveform = waveform.to("cuda")

# 进行推理
result = model.separate(
    waveform,
    chunk_size=44100 * 6,
    overlap_size=44100 * 3,
    gap_size=0,
    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)

只分离伴奏人声 2 轨而不是分离 4 轨：

python
from transformers import AutoModel
import soundfile
import torch
import librosa

model_name = "HiDolen/Mini-BS-RoFormer-18M"
model = AutoModel.from_pretrained(
    model_name,
    trust_remote_code=True,
)
model.to("cuda")

# 加载音频
file = "./Bruno Mars - Runaway Baby.mp3"
waveform, sr = librosa.load(file, sr=44100, mono=False)
waveform = torch.tensor(waveform).float()
waveform = waveform.to("cuda")

# 进行推理
result = model.separate(
    waveform,
    chunk_size=44100 * 6,
    overlap_size=44100 * 3,
    gap_size=0,
    batch_size=2,
    verbose=True,
)

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)

Training Details

使用 MUSDB18HQ 数据进行训练。

不使用原论文中提到的 Multi-STFT 损失项以提高训练速度。

学习率 5e-4，以 batch_size=6 训练 200k 步。

训练时发现的一些技巧：

  • —数据增强不是越多越好。去除音高和拉伸增强能让收敛更快，也可以占用更少 CPU 资源
  • —Multi-STFT 损失可以去掉，不影响训练的同时能极大提高训练速度
  • —MaskEstimator 可以只有一个线性层，仍然能拟合且能节省大量参数
  • —适量减少 freq_transformer 层数几乎不影响整体性能。主要会影响人声这种频段丰富的元素

Acknowledgments

  • —https://github.com/lucidrains/BS-RoFormer
  • —https://arxiv.org/abs/2309.02612 (Music Source Separation with Band-Split RoPE Transformer)