Notes on the ESPnet enh Training Task

Preface

This is the common recipe for ESPnet2 speech enhancement frontend.
This is the general recipe for the ESPnet2 speech enhancement front end.

This post is my study notes on the speech enhancement part of the ESPnet speech processing toolkit. The first draft was generated with the help of claude-3.5-sonnet, and I will keep updating it on that basis and adding my own understanding.

Below are the relevant ESPnet toolkit links

ESPnet

espnet installation

enh.sh

enh.sh official documentation

The location of this enh.sh in espnet: egs2/TEMPLATE/enh1/enh.sh , 13 stages are included.

Training task workflow

  • Choose a dataset
  • Choose a config file (specific parameters can be changed)
  • Run the script

Take the classic dataset wsj0_2mix as an example, wsj0_2mix; in the conf directory’s tuning subdirectory choose a config. I chose train_enh_rnn_tf.yaml, which is used to train an RNN-based speech separation model, where the tf suffix in this config file name stands for Time-Frequency domain

Then run run.sh, for example:

1
./run.sh --stage 1 --stop_stage 6 --conf conf/tuning/train_enh_rnn_tf.yaml

Below is the full content of run.sh (pasted here because it is short)

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
#!/usr/bin/env bash
# Set bash to 'debug' mode, it will exit on :
# -e 'error', -u 'undefined variable', -o ... 'error in pipeline', -x 'print commands',
set -e
set -u
set -o pipefail

min_or_max=min # "min" or "max". This is to determine how the mixtures are generated in local/data.sh.
sample_rate=8k


train_set="tr_${min_or_max}_${sample_rate}"
valid_set="cv_${min_or_max}_${sample_rate}"
test_sets="tt_${min_or_max}_${sample_rate} "

./enh.sh \
--train_set "${train_set}" \
--valid_set "${valid_set}" \
--test_sets "${test_sets}" \
--fs "${sample_rate}" \
--lang en \
--ngpu 1 \
--local_data_opts "--sample_rate ${sample_rate} --min_or_max ${min_or_max}" \
--enh_config conf/tuning/train_enh_dprnn_tasnet.yaml \
"$@"

Note that the enh.sh being called is in fact the enh.sh in TEMPLATE (analyzed in detail below)

1) Monitoring the training process:

  • View the log: tail -f exp/enh_train_*/train.log
  • View key metrics: grep "loss:" exp/enh_train_*/train.log

2) Evaluating the model

1
2
./run.sh --stage 7 --stop_stage 8 \
--conf conf/tuning/train_enh_rnn_tf.yaml

View the evaluation results:

  • Results are saved in exp/enh_train_*/RESULTS.txt
  • Includes metrics such as SI-SNR and SDR

3) Using the model

Enhance a single audio file:

1
2
3
4
5
python -m espnet2.bin.enh_inference \
--audio_file /path/to/mixed.wav \
--config exp/enh_train_*/config.yaml \
--model_file exp/enh_train_*/valid.acc.best.pth \
--output_dir ./enhanced

Get the enhanced audio:

  • The enhanced results are saved in the ./enhanced directory
  • The separated result for each speaker is saved separately

Notes

1) Resuming after training is interrupted:

  • Just run the same command again
  • ESPnet automatically loads the latest checkpoint

2) Common problems:

  • Out of memory: reduce batch_size
  • Out of GPU memory: reduce batch_size or use gradient accumulation
  • Training does not converge: adjust the learning rate or check the data preprocessing

3) Suggestions:

  • Test the pipeline on a small dataset first
  • Keep the config files and logs safe
  • Back up experiment results regularly

Config file analysis

Basic training parameters

1
2
3
4
5
6
optim: adam  # 优化器选择:adam优化器
init: xavier_uniform # 参数初始化方式:xavier均匀分布初始化
max_epoch: 100 # 最大训练轮数
batch_type: folded # 批次类型:folded表示按序列长度折叠
batch_size: 8 # 每批次样本数
num_workers: 4 # 数据加载器的并行工作进程数

Optimizer configuration

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
optim_conf:
lr: 1.0e-03 # 初始学习率
eps: 1.0e-08 # 数值稳定性参数
weight_decay: 1.0e-7 # L2正则化系数

# 早停耐心值:验证集性能多少轮未改善就停止
patience: 10

# 验证集调度器判断标准
val_scheduler_criterion:
- valid
- loss

# 最佳模型保存标准
best_model_criterion:
- - valid
- si_snr # 尺度不变信噪比
- max # 最大化
- - valid
- loss # 损失值
- min # 最小化

# 保存最好的模型数量
keep_nbest_models: 1

# 学习率调度器:当验证集性能不再提升时降低学习率
scheduler: reducelronplateau
scheduler_conf:
mode: min # 监控模式:最小化
factor: 0.7 # 学习率降低因子
patience: 1 # 调度器耐心值

Loss function configuration

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
# A list for criterions
# The overlall loss in the multi-task learning will be:
# loss = weight_1 * loss_1 + ... + weight_N * loss_N
# The default `weight` for each sub-loss is 1.0
criterions:
# 第一个损失函数:均方误差(MSE)
- name: mse
conf:
compute_on_mask: True # 在掩码上计算
mask_type: PSM # 相位敏感掩码
wrapper: pit # 用PIT(排列不变训练)包装
wrapper_conf:
weight: 1.0 # 损失权重

# 第二个损失函数:L1损失
- name: l1
conf:
compute_on_mask: False # 在波形上计算
wrapper: pit
wrapper_conf:
weight: 1.0
independent_perm: False # 使用前一个criterion的排列顺序

# 第三个损失函数:SI-SNR损失
- name: si_snr
conf:
eps: 1.0e-7 # 数值稳定性参数
wrapper: pit
wrapper_conf:
weight: 5.0 # 较大权重表示更重视此损失
independent_perm: False

Model architecture configuration

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
encoder: stft # STFT编码器配置
encoder_conf:
n_fft: 256 # FFT点数
hop_length: 128 # 帧移

# STFT解码器配置
decoder: stft
decoder_conf:
n_fft: 256
hop_length: 128

# 分离器配置:RNN架构
separator: rnn
separator_conf:
rnn_type: blstm # 双向LSTM
num_spk: 2 # 说话人数量
nonlinear: relu # 激活函数
layer: 3 # RNN层数
unit: 896 # 隐层单元数
dropout: 0.5 # Dropout比率

Analysis of enh.sh

Configuration before Stage 1

Basic settings

  1. bash debug mode settings

    1
    2
    3
    set -e        # 遇到错误就退出
    set -u # 使用未定义变量时报错
    set -o pipefail # 管道中任一命令失败则整个管道失败
  2. Helper functions

    1
    2
    3
    4
    5
    6
    7
    8
    # 日志函数:打印时间戳和调用位置信息
    log() {
    local fname=${BASH_SOURCE[1]##*/}
    echo -e "$(date '+%Y-%m-%dT%H:%M:%S') (${fname}:${BASH_LINENO[0]}:${FUNCNAME[1]}) $*"
    }

    # 求最小值函数:用于计算并行作业数
    min()

Note: one thing that puzzles me is why the log is not redirected to a file. Doesn’t just echoing it get very long?

Required parameters

  1. Dataset related
  • --train_set: name of the training set
  • --valid_set: name of the validation set
  • --test_sets: list of test set names

Optional parameters

  1. Basic configuration parameters

    1
    2
    3
    4
    5
    6
    7
    8
    9
    10
    stage=1                # 处理开始的阶段
    stop_stage=10000 # 处理结束的阶段
    skip_data_prep=false # 是否跳过数据准备阶段
    skip_train=false # 是否跳过训练阶段
    skip_eval=false # 是否跳过推理和评估阶段
    skip_packing=true # 是否跳过打包阶段
    skip_upload_hf=true # 是否跳过上传到HuggingFace阶段
    ngpu=1 # GPU数量(0表示使用CPU)
    num_nodes=1 # 节点数量
    nj=32 # 并行作业数
  2. Feature extraction parameters

    1
    2
    3
    4
    5
    feats_type=raw        # 特征类型(raw或fbank_pitch)
    audio_format=flac # 音频格式:wav,flac等
    fs=16k # 采样率
    min_wav_duration=0.1 # 最短音频长度(秒)
    max_wav_duration=20 # 最长音频长度(秒)
  3. Enhancement model parameters

    1
    2
    3
    4
    5
    6
    7
    8
    9
    enh_exp=             # 增强实验目录路径
    enh_tag= # 增强模型训练结果目录的后缀
    enh_config= # 增强模型训练配置
    enh_args= # 增强模型训练的额外参数
    ref_num=2 # 参考信号数量(等于说话人数量)
    inf_num= # 模型输出的推理结果数量
    noise_type_num=1 # 输入音频中的噪声类型数量
    dereverb_ref_num=1 # 去混响参考信号数量
    is_tse_task=false # 是否为目标说话人提取任务
  4. Training data parameters

    1
    2
    3
    use_dereverb_ref=false   # 是否使用去混响参考信号
    use_noise_ref=false # 是否使用噪声参考信号
    variable_num_refs=false # 是否使用可变数量的参考信号
  5. Inference and evaluation parameters

    1
    2
    3
    inference_args="--normalize_output_wav true --output_format wav"  # 推理参数
    inference_model=valid.loss.ave.pth # 推理使用的模型文件
    scoring_protocol="STOI SDR SAR SIR SI_SNR" # 评分指标

Detailed analysis of each Stage

Stage 1: Data preparation

  • Function: prepare the training, validation and test datasets
  • Execution: calls the local/data.sh script to process the data
  • Key code:
    1
    2
    3
    4
    5
    if [ ${stage} -le 1 ] && [ ${stop_stage} -ge 1 ]; then
    log "Stage 1: Data preparation for data/${train_set}, data/${valid_set}, etc."
    # [Task dependent] 需要为新语料库创建data.sh
    local/data.sh ${local_data_opts}
    fi
  • Important notes:
    • This stage is task dependent; a corresponding data.sh script needs to be created for the specific corpus
    • The local_data_opts parameter can be passed to data.sh to customize the data processing
  • Outputs:
    • Generated under data/${train_set}, data/${valid_set}, etc.:
      • wav.scp: mapping of audio file paths
      • utt2spk: utterance-to-speaker mapping
      • spk2utt: speaker-to-utterance mapping
      • mix.scp: list of mixed audio files
      • ref.scp: list of reference audio files

Note: at first I spent a long time looking for where data.sh was, and later found it inside the specific dataset (see the training task workflow above)

data.sh
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
#!/usr/bin/env bash

# 设置bash的错误处理
set -e # 遇到错误就退出
set -u # 使用未定义变量时报错
set -o pipefail # 管道中任一命令失败则整个管道失败

# 定义日志函数
log() {
local fname=${BASH_SOURCE[1]##*/}
echo -e "$(date '+%Y-%m-%dT%H:%M:%S') (${fname}:${BASH_LINENO[0]}:${FUNCNAME[1]}) $*"
}

# 帮助信息
help_message=$(cat << EOF
Usage: $0 [--min_or_max <min/max>] [--sample_rate <8k/16k>]
optional argument:
[--min_or_max]: min (Default), max # 混合方式:最小或最大
[--sample_rate]: 8k (Default), 16k # 采样率:8k或16k
EOF
)

# 导入数据库配置
. ./db.sh

# 设置路径变量
wsj_full_wav=$PWD/data/wsj0/wsj0_wav # WSJ0原始音频路径
wsj_2mix_wav=$PWD/data/wsj0_mix/2speakers # 双说话人混合音频路径
wsj_2mix_scripts=$PWD/data/wsj0_mix/scripts # 混合脚本路径

# 设置文本相关变量
other_text=data/local/other_text/text # 其他文本数据路径
nlsyms=data/nlsyms.txt # 非语言符号文件
min_or_max=min # 默认混合方式为min
sample_rate=8k # 默认采样率为8k

# 解析命令行参数
. utils/parse_options.sh

# 检查WSJ0和WSJ1数据集路径是否存在
if [ ! -e "${WSJ0}" ]; then
log "Fill the value of 'WSJ0' of db.sh"
exit 1
fi
if [ ! -e "${WSJ1}" ]; then
log "Fill the value of 'WSJ1' of db.sh"
exit 1
fi

# 设置数据集名称
train_set="tr_"${min_or_max}_${sample_rate} # 训练集
train_dev="cv_"${min_or_max}_${sample_rate} # 验证集
recog_set="tt_"${min_or_max}_${sample_rate} # 测试集

### WSJ0混合数据处理部分 ###
# 下载混合脚本并创建双说话人混合音频
local/wsj0_create_mixture.sh ${wsj_2mix_scripts} ${WSJ0} ${wsj_full_wav} \
${wsj_2mix_wav} || exit 1;

# 准备WSJ0_2mix数据集
local/wsj0_2mix_data_prep.sh --min-or-max ${min_or_max} --sample-rate ${sample_rate} \
${wsj_2mix_wav}/wav${sample_rate}/${min_or_max} ${wsj_2mix_scripts} ${wsj_full_wav} || exit 1;

### 创建参考音频的.scp文件 ###
# 为每个数据集创建说话人1和说话人2的scp文件
for folder in ${train_set} ${train_dev} ${recog_set}; do
sed -e 's/\/mix\//\/s1\//g' ./data/$folder/wav.scp > ./data/$folder/spk1.scp
sed -e 's/\/mix\//\/s2\//g' ./data/$folder/wav.scp > ./data/$folder/spk2.scp
done

### WSJ语料库处理部分 ###
# 准备WSJ数据
log "local/wsj_data_prep.sh ${WSJ0}/??-{?,??}.? ${WSJ1}/??-{?,??}.?"
local/wsj_data_prep.sh ${WSJ0}/??-{?,??}.? ${WSJ1}/??-{?,??}.?

# 格式化WSJ数据
log "local/wsj_format_data.sh"
local/wsj_format_data.sh

# 创建wsj目录并移动相关数据
log "mkdir -p data/wsj"
mkdir -p data/wsj
log "mv data/{dev_dt_*,local,test_dev*,test_eval*,train_si284} data/wsj"
mv data/{dev_dt_*,local,test_dev*,test_eval*,train_si284} data/wsj

# 准备额外的文本数据
log "Prepare text from lng_modl dir..."
mkdir -p "$(dirname ${other_text})"

# 处理语言模型训练数据
zcat ${WSJ1}/13-32.1/wsj1/doc/lng_modl/lm_train/np_data/{87,88,89}/*.z | \
grep -v "<" | tr "[:lower:]" "[:upper:]" | \
awk '{ printf("wsj1_lng_%07d %s\n",NR,$0) } ' > ${other_text}

# 创建非语言符号文件
log "Create non linguistic symbols: ${nlsyms}"
cut -f 2- data/wsj/train_si284/text | tr " " "\n" | sort | uniq | grep "<" > ${nlsyms}
cat ${nlsyms}

Start
├─ Check the WSJ dataset paths
├─ Generate mixed speech data
├─ Create speaker separation files
├─ Prepare the raw WSJ data
├─ Process additional text
└─ Extract non-linguistic symbols
End

Stage 2: Speed perturbation

  • Function: apply speed perturbation augmentation to the training data
  • Condition: only executed when speed_perturb_factors is set and dereverberation references are not used
  • Processing: perturb the audio at different speeds to generate augmented data
  • Outputs:
    • Generated under the data/${train_set}_sp directory:
      • The perturbed audio files and the corresponding config files
      • Updated wav.scp, utt2spk, spk2utt and other files

Stage 3: Audio formatting

  • Function: process audio formats uniformly
  • Key code:
    1
    2
    3
    4
    5
    # 格式化wav.scp文件
    scripts/audio/format_wav_scp.sh --nj "${nj}" --cmd "${train_cmd}" \
    --out-filename "${spk}.scp" \
    --audio-format "${audio_format}" --fs "${fs}" ${_opts} \
    "data/${dset}/${spk}.scp" "${data_feats}${_suf}/${dset}"
  • Processing steps:
    1. Re-create the “wav.scp” file
    2. Unify the audio format and sampling rate
    3. Handle the multi-speaker case
    4. Support splitting according to a segments file
  • Outputs:
    • Under the ${data_feats}/${dset} directory:
      • Audio files in the unified format
      • Updated wav.scp file
      • The .scp file for each speaker

Stage 4: Data filtering

  • Function: remove audio data that does not meet the length requirements
  • Key code:
    1
    2
    3
    4
    5
    6
    7
    8
    9
    10
    # 计算最小和最大长度(样本数)
    _fs=$(python3 -c "import humanfriendly as h;print(h.parse_size('${fs}'))")
    _min_length=$(python3 -c "print(int(${min_wav_duration} * ${_fs}))")
    _max_length=$(python3 -c "print(int(${max_wav_duration} * ${_fs}))")

    # 根据长度筛选数据
    <"${data_feats}/org/${dset}/utt2num_samples" \
    awk -v min_length="${_min_length}" -v max_length="${_max_length}" \
    '{ if ($2 > min_length && $2 < max_length ) print $0; }' \
    >"${data_feats}/${dset}/utt2num_samples"
  • Processing steps:
    1. Convert durations into numbers of samples
    2. Filter the audio by number of samples
    3. Update the related scp files
  • Outputs:
    • Under the ${data_feats}/${dset} directory:
      • The filtered utt2num_samples file
      • Updated wav.scp, spk.scp and other files
      • Audio entries that do not meet the length requirements are removed

Stage 5: Statistics collection

  • Function: collect the statistics needed for training
  • Key code:
    1
    2
    3
    4
    5
    6
    7
    ${python} -m ${train_module} \
    --collect_stats true \
    ${_train_data_param} \
    ${_valid_data_param} \
    --train_shape_file "${_logdir}/train.JOB.scp" \
    --valid_shape_file "${_logdir}/valid.JOB.scp" \
    --output_dir "${_logdir}/stats.JOB"
  • Processing steps:
    1. Collect statistics of the training and validation data
    2. Generate shape files
    3. Aggregate the statistics
  • Outputs:
    • Under the ${_logdir} directory:
      • stats.JOB directory: contains the statistics
      • train.JOB.scp: shape information of the training data
      • valid.JOB.scp: shape information of the validation data
      • global_stats: global statistics

Stage 6: Model training

  • Function: run the training of the enhancement model
  • Key code:
    1
    2
    3
    4
    5
    6
    7
    8
    9
    10
    ${python} -m ${train_module} \
    ${_train_data_param} \
    ${_valid_data_param} \
    ${_train_shape_param} \
    ${_valid_shape_param} \
    ${_fold_length_param} \
    --resume true \
    --output_dir "${enh_exp}" \
    ${init_param:+--init_param $init_param} \
    ${_opts} ${enh_args}
  • Processing steps:
    1. Set up the training and validation data
    2. Configure the training parameters
    3. Support resuming from checkpoints
    4. Optional initialization from a pretrained model
  • Outputs:
    • Under the ${enh_exp} directory:
      • config.yaml: model config file
      • Model checkpoint files (*.pth)
      • trainer.log: training log
      • Validation results and curve plots

Stage 7: Inference

  • Function: use the trained model to enhance audio
  • Key code:
    1
    2
    3
    4
    5
    6
    7
    8
    ${python} -m ${infer_module} \
    --ngpu "${_ngpu}" \
    --fs "${fs}" \
    ${_data_param} \
    --key_file "${_logdir}"/keys.JOB.scp \
    --train_config "${enh_exp}"/config.yaml \
    --model_file "${enh_exp}"/"${inference_model}" \
    --output_dir "${_logdir}"/output.JOB
  • Processing steps:
    1. Load the trained model
    2. Run inference on the test sets
    3. Generate the enhanced audio
    4. Support GPU inference
  • Outputs:
    • Under the ${_logdir}/output.JOB directory:
      • enhanced.wav: the enhanced audio file
      • keys.JOB.scp: key-value pairs of the processed audio
      • Inference logs and result files

Stage 8: Scoring

  • Function: evaluate the enhancement effect
  • Key code:
    1
    2
    3
    4
    5
    6
    7
    ${python} -m espnet2.bin.enh_scoring \
    --key_file "${_logdir}"/keys.JOB.scp \
    --output_dir "${_logdir}"/output.JOB \
    ${_ref_scp} \
    ${_inf_scp} \
    --ref_channel ${ref_channel} \
    --flexible_numspk ${flexible_numspk}
  • Evaluation metrics:
    • STOI: speech intelligibility
    • SDR: signal-to-distortion ratio
    • SAR: signal-to-artifacts ratio
    • SIR: signal-to-interference ratio
    • SI_SNR: scale-invariant signal-to-noise ratio
  • Outputs:
    • Under the ${_logdir}/output.JOB directory:
      • scoring.txt: contains each scoring metric
      • score_stats: detailed scoring statistics
      • Score distribution plots for each metric

Stage 9-10: ASR evaluation

  • Function: evaluate the enhancement effect with an ASR model
  • Key code:
    1
    2
    3
    4
    5
    6
    7
    ${python} -m espnet2.bin.asr_inference \
    --ngpu "${_ngpu}" \
    --data_path_and_name_and_type "${_ddir}/wav.scp,speech,${_type}" \
    --key_file "${_logdir}"/keys.JOB.scp \
    --asr_train_config "${asr_exp}"/config.yaml \
    --asr_model_file "${asr_exp}"/"${inference_asr_model}" \
    --output_dir "${_logdir}"/output.JOB
  • Processing steps:
    1. Decode the enhanced audio with the ASR model
    2. Compute the character error rate (CER) or word error rate (WER)
    3. Generate the evaluation report
  • Outputs:
    • Under the ${_logdir}/output.JOB directory:
      • asr_inference.txt: ASR decoding results
      • text: the recognized text
      • wer.txt/cer.txt: error rate statistics

Stage 11: Model packing

  • Function: pack the trained model
  • Processing:
    • Pack the model files
    • Pack the config information
    • Generate the release package
  • Outputs:
    • Under the ${enh_exp}/pack directory:
      • model.zip: the packed model file
      • config.yaml: a copy of the config file
      • README.md: model documentation

Stage 12: Model upload

  • Function: upload the model to HuggingFace
  • Condition: executed when skip_upload_hf=false
  • Processing:
    • Prepare the files to upload
    • Configure the HuggingFace repository
    • Upload the model
  • Outputs:
    • In the HuggingFace repository:
      • The uploaded model files and config
      • Model card
      • Example code and usage instructions

Summary

enh1.sh is a complete speech enhancement pipeline script, covering the whole process from data preparation to model training and evaluation. With properly configured parameters, every part of the pipeline can be controlled flexibly. When using it, pay special attention to the following:

  1. The names of the training, validation and test sets must be provided
  2. Set the number of GPUs and parallel jobs reasonably according to your needs
  3. The execution flow can be controlled via stage and stop_stage
  4. The evaluation stages offer multiple evaluation methods, including objective metrics and ASR evaluation

Translated from the Chinese original.

Welcome to my other publishing channels

中文