wavegrad

WaveGrad 项目介绍

WaveGrad 是由 Google Brain 团队设计的一个快速高质量神经声码器。它的架构旨在将对数规模的 Mel 频谱图通过迭代优化转换为波形，详细描述见WaveGrad: Estimating Gradients for Waveform Generation。

项目概况

WaveGrad 项目自 2020 年 10 月 15 日以来已经完成了以下功能：

稳定的训练，支持 22 kHz 和 24 kHz 的音频采样率。
高质量合成音频。
混合精度训练，提升计算效率。
多 GPU 训练，加快模型的训练速度。
自定义噪声调度，实现更快的推断速度。
命令行推断，方便用户直接使用。
编程推断 API，提供更灵活的调用方式。
PyPI 包，便于安装使用。
音频样本展示，提供直观的质量验证。
预训练模型下载，方便复用。
尚未完成预计算噪声调度。

音频样例

用户可以通过这里查看 24 kHz 的音频样例，以亲耳感受 WaveGrad 的合成质量。

预训练模型

WaveGrad 提供了 24 kHz 的预训练模型（183 MB），可以直接下载并使用。这一模型的 SHA256 校验值是 65e9366da318d58d60d2c78416559351ad16971de906e53b415836c068e335f3。

安装方法

用户可以通过以下命令使用 pip 安装 WaveGrad：

pip install wavegrad

或者从 GitHub 克隆代码并安装：

git clone https://github.com/lmnt-com/wavegrad.git
cd wavegrad
pip install .

训练详情

在开始训练之前，用户需要准备一个训练数据集。数据集内部可以有任意目录结构，只需要保证其中的 .wav 文件为 16 位单声道即可（例如 LJSpeech, VCTK 数据集）。默认情况下，WaveGrad 假定采样率为 22 kHz，如果需要修改采样率，可以编辑 params.py。

训练步骤如下：

python -m wavegrad.preprocess /path/to/dir/containing/wavs
python -m wavegrad /path/to/model/dir /path/to/dir/containing/wavs

# 开启另一个终端以监测训练进度：
tensorboard --logdir /path/to/model/dir --bind_all

经过大约 20000 步（在 2080 Ti GPU 上约需 1.5 小时），您应当能够听到清晰可辨的语音。

推断 API

WaveGrad 提供了便捷的推断 API，基本使用方法如下：

from wavegrad.inference import predict as wavegrad_predict

model_dir = '/path/to/model/dir'
spectrogram = # 获取格式为 [N,C,W] 的频谱图
audio, sample_rate = wavegrad_predict(spectrogram, model_dir)

# 生成的 `audio` 是一个格式为 [N,T] 的 GPU 张量。

如果您有自定义的噪声调度：

from wavegrad.inference import predict as wavegrad_predict

params = { 'noise_schedule': np.load('/path/to/noise_schedule.npy') }
model_dir = '/path/to/model/dir'
spectrogram = # 获取格式为 [N,C,W] 的频谱图
audio, sample_rate = wavegrad_predict(spectrogram, model_dir, params=params)

# `audio` 是 [N,T] 格式的 GPU 张量。

推断命令行工具

用户可以使用如下命令进行推断：

python -m wavegrad.inference /path/to/model /path/to/spectrogram -o output.wav

噪声调度

默认情况下，WaveGrad 使用 1000 次迭代来优化波形，这比实时合成稍慢。然而，WaveGrad 可以通过仅使用 6 次迭代而无需重新训练模型来实现高质量且快于实时的合成。

为了实现这一速度提升，用户需要寻找适合自己数据集的噪声调度。WaveGrad 提供了一个脚本来帮助用户进行搜索：

python -m wavegrad.noise_schedule /path/to/trained/model /path/to/preprocessed/validation/dataset
python -m wavegrad.inference /path/to/trained/model /path/to/spectrogram -n noise_schedule.npy -o output.wav

默认设置应该能够在较短的时间内提供不错的结果。如果用户希望找到更佳的噪声调度，或希望使用不同的推断迭代次数，可以运行 noise_schedule 脚本并使用 --help 选项查看更多配置选项。