WaveGrad
谷歌大脑高保真WaveGrad声码器的实现(PyTorch)(论文)。GitHub上首个实现高质量6次迭代生成的项目。
状态
- 已记录API。
- 高保真生成。
- 多次迭代推理支持(低迭代次数稳定)。
- 支持混合精度,稳定且快速的训练。
- 支持分布式训练。
- 在单个12GB GPU上以96批次大小成功运行训练。
- 支持CLI推理。
- 灵活的架构配置,适用于您自己的数据。
- 在流行的GPU和CPU设备上估计RTF(见下文)。
- 在RTX 2080 Ti上,100及以下迭代次数的推理速度快于实时。6次迭代推理速度快于论文中报道的。
- 并行网格搜索以找到最佳噪声计划。
- 上传了不同迭代次数生成的样本(见
generated_samples
文件夹)。 - 在22KHz LJSpeech数据集上带噪声计划的预训练检查点。
实时因子 (RTF)
参数数量: 15,810,401
模型 | 稳定 | RTX 2080 Ti | Tesla K80 | Intel Xeon 2.3GHz* |
---|---|---|---|---|
1000次迭代 | + | 9.59 | - | - |
100次迭代 | + | 0.94 | 5.85 | - |
50次迭代 | + | 0.45 | 2.92 | - |
25次迭代 | + | 0.22 | 1.45 | - |
12次迭代 | + | 0.10 | 0.69 | 4.55 |
6次迭代 | + | 0.04 | 0.33 | 2.09 |
*注意: 使用了旧版的Intel Xeon CPU。
关于
WaveGrad是一个通过估计数据密度梯度生成波形的条件模型,具有与WaveNet相似的采样质量。这个声码器既不是GAN,也不是正则流,也不是经典的自回归模型。声码器的主要概念基于去噪扩散概率模型(DDPM),利用朗之万动力学和评分匹配框架。此外,与经典的DDPM相比,WaveGrad在朗之万动力学迭代采样方案方面实现了超快速收敛(6次迭代,甚至可能更少)。
安装
- 克隆此仓库:
git clone https://github.com/ivanvovk/WaveGrad.git
cd WaveGrad
- 安装依赖:
pip install -r requirements.txt
训练
1 准备数据
- 制作您的音频数据的训练和测试文件列表,格式和
filelists
文件夹中的相同。 - 在
configs
文件夹中创建一个配置文件*。
*注意: 如果您打算更改STFT的hop_length
,请确保您的上采样factors
在配置中的乘积等于您的新hop_length
。
2 单GPU和分布式GPU训练
- 打开
runs/train.sh
脚本并指定可见的GPU设备和配置文件的路径。如果指定多个GPU,训练将以分布式模式运行。 - 运行
sh runs/train.sh
3 Tensorboard和日志记录
要跟踪您的训练进程,请运行tensorboard:tensorboard --logdir=logs/YOUR_LOGDIR_FOLDER
。所有日志信息和检查点都将保存在logs/YOUR_LOGDIR_FOLDER
。logdir
在配置文件中指定。
4 噪声计划网格搜索
一旦模型训练完成,在notebooks/inference.ipynb
中进行所需迭代次数的最佳计划网格搜索*。代码支持并行性,因此您可以指定多个作业数以加速搜索。
*注意: 网格搜索只对少量迭代(如6或7次)是必要的。对于较大数量,尝试Fibonacci序列benchmark.fibonacci(...)
初始化:我使用它进行25次迭代,它效果很好。从好的25次迭代计划开始,例如,您可以通过复制元素来构建一个更高阶的计划。
预训练模型的噪声计划
- 6次迭代计划通过网格搜索获得。在此基础上,我手动找到一个稍好些的近似方案。
- 7次迭代计划以同样方式获得。
- 12次迭代计划以同样方式获得。
- 25次迭代计划使用Fibonacci序列
benchmark.fibonacci(...)
获得。 - 50次迭代计划通过重复25次迭代计划中的元素获得。
- 100次迭代计划以同样方式获得。
- 1000次迭代计划以同样方式获得。
推理
CLI
将您的梅尔频谱图放入某个文件夹。制作一个文件列表。然后用自己的参数运行此命令:
sh runs/inference.sh -c <your-config> -ch <your-checkpoint> -ns <your-noise-schedule> -m <your-mel-filelist> -v "yes"
Jupyter Notebook
更多推理细节见notebooks/inference.ipynb
。在那里您还可以找到如何为模型设置噪声计划并进行最佳方案的网格搜索。
其他
生成的音频
生成的音频例子见generated_samples
文件夹。如果为后者找到最佳计划,1000次迭代和6次迭代推理之间的质量下降是不可忽略的。
预训练检查点
您可以通过此Google Drive链接找到一个22KHz LJSpeech的预训练检查点*。
*注意: 上传的检查点是一个包含单个键'model'
的dict
。
重要细节、问题和评论
- 在训练过程中,WaveGrad使用1000次迭代和从范围(1e-6, 0.01)线性缩放的betas的默认噪声计划。对于推理,您可以设置具有较少迭代次数的其他计划。仔细调整betas,输出质量真的高度依赖于它。
- 默认情况下,模型以混合精度方式运行。批量大小与论文中修改的不同(256 -> 96),因为作者在TPU上训练了他们的模型。
- 在单个GPU上经过约1万次训练迭代(1-2小时)后,模型表现出对50次迭代推理的良好生成。总训练时间约为1-2天(完全收敛)。
- 在某个时刻,训练可能开始表现异常(损失爆炸),因此我引入了学习率(LR)调度和梯度剪裁。如果您的数据损失爆炸,请尝试稍微降低LR调度器gamma。这应该会有所帮助。
- 默认情况下,您的STFT的跳跃长度(hop length)等于300(因此总上采样因子)。其他情况尚未测试,但您可以尝试。请记住,总上采样因子应仍等于您的新跳跃长度。
更新历史
- (新: 2020/10/24) 大更新。分布式训练和混合精度支持。更正的位置编码。推理的CLI支持。并行网格搜索。模型大小显著减少。
- 新的NVIDIA Tesla K80 GPU卡(在Google Colab服务中流行)和CPU Intel Xeon 2.3GHz的新RTF信息。
- 大更新。新的6次迭代生成的样本示例。新的噪声计划设置API。添加了最佳计划网格搜索代码。
- 通过引入更智能的学习率调度器改进训练。获得高保真合成。
- 稳定的训练和多次迭代推理。支持6次迭代噪声计划。
- 稳定的训练和固定迭代推理,显著减少了背景静态噪声。解决了所有位置编码问题。
- 稳定训练25次、50次和1000次固定迭代模型。未发现位置编码(C=5000来自论文)的线性缩放(错误)。
- 稳定训练25次、50次和1000次固定迭代模型。修正了位置编码降尺度。并行片段采样被全梅尔采样取代。
- (发布,GitHub上首个)。并行片段采样和破损位置编码降尺度。来自并行片段生成的拼接点击声导致质量差。
参考
- Nanxin Chen 等人,WaveGrad: Estimating Gradients for Waveform Generation
- Jonathan Ho 等人,Denoising Diffusion Probabilistic Models
- 去噪扩散概率模型仓库(TensorFlow实现),从中采用了扩散计算