Implement Wave-U-Net by PyTorch, and migrate it to the speech enhancement.
# 确保 CUDA 的 bin 目录添加到 PATH 环境变量中
# 通过附加 LD_LIBRARY_PATH 环境变量来安装 CUDA 附带的 CUPTI
export PATH="/usr/local/cuda-10.0/bin:$PATH"
export LD_LIBRARY_PATH="/usr/local/cuda-10.0/lib64:$LD_LIBRARY_PATH"
# 安装 Anaconda,以清华镜像源,python 3.6.5为例
wget https://mirrors.tuna.tsinghua.edu.cn/anaconda/archive/Anaconda3-5.2.0-Linux-x86_64.sh
chmod a+x Anaconda3-5.2.0-Linux-x86_64.sh
./Anaconda3-5.2.0-Linux-x86_64.sh # 按 f 翻页,默认安装在 ~/anaconda 目录下,安装过程会提示修改 PATH 变量
# Create env
conda create -n wave-u-net python=3
conda activate wave-u-net
# Install deps
conda install pytorch torchvision cudatoolkit=10.0 -c pytorch
conda install tensorflow-gpu
conda install matplotlib
pip install tqdm librosa
pip install pystoi # for STOI metric
pip install pesq # for PESQ metric
# 配置好环境与依赖之后,可以拉取代码
git clone https://github.com/haoxiangsnr/Wave-U-Net-for-Speech-Enhancement.git
当前项目有三个入口文件:
- 用于训练模型的入口文件:
train.py
- 用于增强带噪语音的入口文件:
enhancement.py
- 用于测试模型降噪能力的入口文件(TODO):
test.py
使用 train.py
训练模型,它接收三个命令行参数:
-h
,显示帮助信息-C, --config
,指定训练所需的配置文件-R, --resume
,从最近一次保存的模型断点处继续训练
语法:python train.py [-h] -C CONFIG [-R]
例如:
python train.py -C config/train.json
# 训练模型所用的配置文件为 config/train.json
# 使用所有的 GPU 进行训练
python train.py -C config/train.json -R
# 训练模型所用的配置文件为 config/train.json
# 使用所有的 GPU 从最近一次保存的模型断点继续训练
CUDA_VISIBLE_DEVICES=1,2 python train.py -C config/train.json
# 训练模型所用的配置文件为 config/train.json
# 使用 1,2 号索引的GPU进行训练
CUDA_VISIBLE_DEVICES=-1 python train.py -C config/train.json
# 训练模型所用的配置文件为 config/train.json
# 使用 CPU 进行训练
补充:
- 一般将训练所需要的配置文件放置于
config/train/
目录下 - 训练配置文件中的参数见“参数说明”部分
- 配置文件的文件名即是实验名
使用 enhancement.py
来增强带噪语音,它接收以下参数:
-h, --help
,显示帮助信息-C, --config
,指定增强语音所用的模型,以及被增强的数据集。-D, --device
,增强所用的 GPU 索引,-1 表示使用 CPU-O, --output_dir
,指定在哪里存储增强后的语音,需要确保这个目录提前存在-M, --model_checkpoint_path
,模型断点的路径,拓展名为 .tar 或 .pth
语法:python enhancement.py [-h] -C CONFIG [-D DEVICE] -O OUTPUT_DIR -M MODEL_CHECKPOINT_PATH
例如:
python enhancement.py -C config/enhancement/unet_basic.json -D 0 -O enhanced -M /media/imucs/DataDisk/haoxiang/Experiment/Wave-U-Net-for-Speech-Enhancement/smooth_l1_loss/checkpoints/model_0020.pth
# 增强语音所用的配置文件为 config/enhancement/unet_basic.json,使用这个文件可以指定增强所需的模型以及数据集信息
# 使用索引为 0 的 GPU
# 输出的目录为 enhanced/,该目录需要提前新建好
# 指定模型断点的路径
python enhancement.py -C config/enhancement/unet_basic.json -D -1 -O enhanced -M /media/imucs/DataDisk/haoxiang/Experiment/Wave-U-Net-for-Speech-Enhancement/smooth_l1_loss/checkpoints/model_0020.tar
# 使用 CPU 来增强语音
补充:
- 一般将增强所需要的配置文件放置于
config/enhancement/
目录下 - 增强配置文件中的参数见“参数说明”部分
TODO
训练中产生的所有日志信息都会存储至config["save_location"]/<config_filename>/
目录下。假设用于训练的配置文件为config/train/sample_16384.json
,sample_16384.json
中save_location
参数的值为/home/UNet/
,那么当前实验训练过程中产生的日志会存储在 /home/UNet/sample_16384/
目录下。
该目录会包含以下内容:
logs/
目录: 存储 Tensorboard 相关的数据,包含损失曲线,波形文件,语音文件等checkpoints/
目录: 存储模型的所有断点,后续可从这些断点处重启训练或进行语音增强config.json
文件: 训练配置文件的备份
在训练过程中可以使用 tensorboard
来启动一个静态的前端服务器,可视化相关目录中的日志数据:
tensorboard --logdir config["save_location"]/<config_filename>/
# 可使用 --port 指定 tensorboard 静态服务器的启动端口
tensorboard --logdir config["save_location"]/<config_filename>/ --port <port>
# 例如,配置文件中的 "save_location" 参数为 "/home/happy/Experiments",配置文件名为 "train_config.json",修改默认端口为 6000
# 可使用如下命令:
tensorboard --logdir /home/happy/Experiments/train_config --port 6000
在项目运行过程,会产生多个目录,均有不同的用途:
- 主目录:当前 README.md 所在的目录,存储着所有源代码
- 训练目录:训练配置文件中的
config["save_location"]
目录,存储当前项目的所有实验日志和模型断点 - 实验目录:
config["save_location"]/<实验名>/
目录,存储着某一次实验的日志信息
config/train/<实验名>.json
,训练过程中产生的日志信息会存放在config["save_location"]/<实验名>/
目录下
{
"seed": 0, // 保证实验可重复性的随机种子
"description": "...", // 实验描述,后续会显示在 Tensorboard 中
"root_dir": "~/Experiments/Wave-U-Net", //存放实验结果的目录
"cudnn_deterministic": false,
"trainer": { // 训练过程
"module": "trainer.trainer", // 训练器模型的文件
"main": "Trainer", // 训练器模型的具体类
"epochs": 1200, // 训练的上限
"save_checkpoint_interval": 10, // 保存模型断点的间隔
"validation":{
"interval": 10, // 验证的间隔
"find_max": true, // 当 find_max 为 true 时,如果计算出的评价指标为已知的最大值,就会将当前轮次的模型断点另外缓存一份
"custon": {
"visualize_audio_limit": 20, // 验证时可视化音频的间隔,之所以设置这个参数,是因为可视化音频比较慢
"visualize_waveform_limit": 20, // 验证时可视化波形的间隔,之所以设置这个参数,是因为可视化波形比较慢
"visualize_spectrogram_limit": 20, //验证可视化频谱的间隔,之所以设置这个参数,是因为可视化频谱比较慢
"sample_length": 16384 //采样点数
}
}
},
"model": {
"module": "model.unet_basic", // 训练使用的模型文件
"main": "Model", // 训练模型的具体类
"args": {} // 传给模型类的参数
},
"loss_function": {
"module": "model.loss", // 损失函数的模型文件
"main": "mse_loss", // 损失函数模型的具体类
"args": {} // 传给模型类的参数
},
"optimizer": {
"lr": 0.001,
"beta1": 0.9,
"beat2": 0.009
},
"train_dataset": {
"module": "dataset.waveform_dataset", // 存放训练集类模型的文件
"main": "Dataset", // 训练集模型的具体类
"args": { // 传递给训练集类的参数,详见具体的训练集类
"dataset": "~/Datasets/SEGAN_Dataset/train_dataset.txt",
"limit": null,
"offset": 0,
"sample_length": 16384,
"mode":"train"
}
},
"validation_dataset": {
"module": "dataset.waveform_dataset",
"main": "Dataset",
"args": {
"dataset": "~/Datasets/SEGAN_Dataset/test_dataset.txt",
"limit": 400,
"offset": 0,
"mode":"validation"
}
},
"train_dataloader": {
"batch_size": 120,
"num_workers": 40, // 开启多少个线程对数据进行预处理
"shuffle": true,
"pin_memory":true
}
}
config/enhancement/*.json
{
"model": {
"module": "model.unet_basic", // 放置模型的文件
"main": "UNet",// 文件内的具体模型类
"args": {} // 传给模型类的参数
},
"dataset": {
"module": "dataset.waveform_dataset", // 增强使用的数据集类
"main": "WaveformDataset", // 传递给数据集类的参数,详见具体的训练集类
"args": {
"dataset": "/home/imucs/Datasets/2019-09-03-timit_train-900_test-50/enhancement.txt",
"limit": 400,
"offset": 0,
"sample_length": 16384
}
}
}
在增强时,存储数据集路径的 txt 文件仅仅指定带噪语音的路径即可,类似这样:
# enhancement.txt
/home/imucs/tmp/UNet_and_Inpainting/0001_babble_-7dB_Clean.wav
/home/imucs/tmp/UNet_and_Inpainting/0001_babble_-7dB_Enhanced_Inpainting_200.wav
/home/imucs/tmp/UNet_and_Inpainting/0001_babble_-7dB_Enhanced_Inpainting_270.wav
/home/imucs/tmp/UNet_and_Inpainting/0001_babble_-7dB_Enhanced_UNet.wav
/home/imucs/tmp/UNet_and_Inpainting/0001_babble_-7dB_Mixture.wav
- 使用全长语音进行验证
- 增强脚本
- 测试脚本