Spleeter工具简单分析

源码地址:https://github.com/deezer/spleeter

pytorch版本:https://github.com/generalwave/spleeter.pytorch

前言

    spleeter是一款基于深度学习的多音轨分离软件。其使用python3.7基于tensorflow1.15版本编写。本文给出了spleeter的pytorch版本，并将官方预训练模型转化为pytorch模型。该版本经验证效果与官方版本一致。
   从功能上,目前预训练模型为2stems(分离出人声/伴奏)，4stems(分离出人声/伴奏/鼓/贝斯/其他)，5stems(人声/鼓/贝斯/钢琴/其他)。性能上，按照spleeter的官网解释,4stems在使用GPU加速的情况下可以达到100s长度的音乐1s分离完成。从效果上来看，spleeter的各项指标均优于目前的其他开源模型。笔者测试2stems的分离效果，人声和伴奏的分离结果都还不错，两者都比较干净。
   本文将介绍spleeter的设计原理，模型结构以及预测和训练的方法和在阅读代码以及实际使用时笔者的一些理解，如有错误，希望读者批评指正。同时，读者需要注意，本文所介绍的代码级别的内容均基于2stems的模型，以及默认参数设定情况的调用到的代码，其余情况未解释，另外由于spleeter一直在更新，本文行号不一定正确，仅供参考。

设计原理

    spleeter基于频域进行音轨分离。其网络结构中，每条音轨对应着一个unet网络结构。2stems对应着两个unet，4stems对应4个unet网络。unet的网络输入为音频幅度谱，输出为某条音轨的幅度谱。训练时损失函数为计算出音轨的幅度谱与标准幅度谱的L1距离。预测时稍有不同,通过多条音轨的幅度谱计算出每条音轨占据输入音频的能量比例，即每条音轨的mask，通过输入音频频谱乘以mask得到各个音轨的输出频谱，计算得到wav。
    由于spleeter的训练和预测的方法有些差异，因此，笔者将分别叙述两者。
   首先描述训练时的基本原理。如下图1,训练时的一组数据为(音乐,伴奏,人声)，要求三者在时间轴上尽量完全一致.提取三者频谱并计算出幅度谱。将音乐幅度谱分别输入到人声unet和伴奏unet中，得到预测出的人声unet和伴奏unet，分别计算预测结果和标准结果的L1距离并取均值,作为损失函数,其中伴奏unet和人声unet内部参数会随着数据输入不断更新，此处不再赘述。

   预测过程没有标准的人声和伴奏，只有音乐。在于预测出伴奏和人声的幅度谱之后，spleeter将两者分别进行平方,得到人声能量v_eng,和伴奏能量，a_eng,然后使用v_mask = v_eng/(v_eng+a_eng)计算出每个时刻人声在音乐的每个频带上的占比，以及使用a_mask=a_eng/(v_eng+a_eng)计算出伴奏每个时刻在音乐的每个频带上的占比。最后利用输入的音乐频谱分别乘以v_mask和a_mask得到人声和伴奏频谱，使用逆STFT得到人声和伴奏的语音。具体流程如下图:

   此处需要介绍的一个地方为，本文的unet中，并不是存粹的一个unet结构，其中使用unet网络结构，并在最后一层使用sigmoid激活函数，计算出人声或者伴奏的mask系数，此时通过输入幅度谱乘以该系数得到了伴奏或人声的幅度谱。
   笔者认为预测的时候直接获取每条音轨的mask系数乘以输入音频的频谱得到对应音轨的频谱。实验过够，从听觉来说效果还是不错的。这样的好处在于训练目标和预测目标一致，其次多条音轨之间互不影响，方便根据业务需求进行裁剪优化.

流程以及代码讲解

训练部分

   训练部分在代码上的核心逻辑与上述流程图并不一致。上述流程图描述的是单个音乐的处理过程，由于在训练的时候需要大量的数据。因此，重点在于数据集的处理，其中模型的搭建以及训练过程和所用语言和框架强相关，本文不作为重点，只是简述一下。
   本部分将分为两块叙述。第一为使用部分，讲解如何使用spleeter来进行训练。分为训练时需要的命令以及对应参数解析。第二部分为训练时的核心代码讲解。

第一部分,使用篇

python3 -m spleeter train -p configs/2stems/base_config.json -d example

-m: 这是python的写法，表示指定模块，在此处表示使用spleeter模块
-p: 指定配置文件位置,详细内容见附3
-d: 指定数据集位置
此处需要解释的是，该命令表示执行spleeter的训练模型，采用2stems的配置文件，数据集的位置为本项目目录的example。在运行之前，需要在example中准备好数据集，结构无要求。因为在配置文件中会指定训练数据集和测试数据集的索引文件。具体配置文件信息以及索引文件信息参考附3的解释即可。

第二部分,代码篇

本篇将首先解析训练时的代码整体结构，而后重点分析数据集的构建部分。

代码整体结构:

训练代码使用的是tensorflow的estimator的高级api接口。一共分为三部分，分别为构建模型，构建数据集，训练以及预测.
预备知识:
   模型构建: tensorflow使用tf.estimator.Estimator来构建estimator，其中model_fn参数需要给定一个函数，包含了对于训练/验证/预测时的构建的模型。三者在model_fn中，通过mode参数来进行区分[mode参数是Estimator自动填充的]。其中训练模型需要包含，损失函数(loss)和优化器.验证模型[用来在测试集合上做验证]需要给出损失函数定义[一般会看测试集上的loss],预测模型则不需要包含损失函数和优化器。
   数据处理: 使用tf.estimator.TrainSpec进行训练数据处理，其中input_fn参数需要给出一个函数，并且返回值是(features, labels)的格式,同理，对于验证集合,使用的是tf.estimator.EvalSpec.
   模型训练:通过tf.estimator.train_and_evaluate将参数传入即可训练，并且可以在分布式机器上直接训练，无需改动代码
   注意1: partial函数[https://wiki.jikexueyuan.com/project/explore-python/Functional/partial.html],在构建数据集的时候用到了。
   注意2: 在构建数据集的时候spleeter大量使用了dataset.map函数，其效果与python中的map相近，都是对每一个元素做统一的操作，其中dataset.map函数还支持使用多线程处理。
   注意3: 构建数据集的时候用到了dataset.cache()函数，其将数据集缓存到硬盘中，下次使用的时候可以直接调用.
   注意4: tensorflow estimator的api训练时，遍历数据集，如果数据集被