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  # optimizer choice: Adam optimizer
init: xavier_uniform # parameter initialization: Xavier uniform initialization
max_epoch: 100 # maximum number of training epochs
batch_type: folded # batch type: folded means batches are folded by sequence length
batch_size: 8 # samples per batch
num_workers: 4 # number of parallel data-loader worker processes

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 # initial learning rate
eps: 1.0e-08 # numerical stability term
weight_decay: 1.0e-7 # L2 regularization coefficient

# Early-stopping patience: stop after this many epochs without validation improvement
patience: 10

# Criterion the scheduler monitors on the validation set
val_scheduler_criterion:
- valid
- loss

# Criteria for saving the best models
best_model_criterion:
- - valid
- si_snr # scale-invariant signal-to-noise ratio
- max # maximize
- - valid
- loss # loss value
- min # minimize

# Number of best models to keep
keep_nbest_models: 1

# Learning-rate scheduler: lower the learning rate when validation performance stops improving
scheduler: reducelronplateau
scheduler_conf:
mode: min # monitoring mode: minimize
factor: 0.7 # learning-rate reduction factor
patience: 1 # scheduler patience

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:
# First loss function: mean squared error (MSE)
- name: mse
conf:
compute_on_mask: True # compute on the mask
mask_type: PSM # phase-sensitive mask
wrapper: pit # wrap with PIT (permutation invariant training)
wrapper_conf:
weight: 1.0 # loss weight

# Second loss function: L1 loss
- name: l1
conf:
compute_on_mask: False # compute on the waveform
wrapper: pit
wrapper_conf:
weight: 1.0
independent_perm: False # reuse the permutation from the previous criterion

# Third loss function: SI-SNR loss
- name: si_snr
conf:
eps: 1.0e-7 # numerical stability term
wrapper: pit
wrapper_conf:
weight: 5.0 # a larger weight puts more emphasis on this loss
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 configuration
encoder_conf:
n_fft: 256 # FFT size (number of points)
hop_length: 128 # frame shift

# STFT decoder configuration
decoder: stft
decoder_conf:
n_fft: 256
hop_length: 128

# Separator configuration: RNN architecture
separator: rnn
separator_conf:
rnn_type: blstm # bidirectional LSTM
num_spk: 2 # number of speakers
nonlinear: relu # activation function
layer: 3 # RNN layer count
unit: 896 # number of hidden units
dropout: 0.5 # Dropout rate

Analysis of enh.sh

Configuration before Stage 1

Basic settings

  1. bash debug mode settings

    1
    2
    3
    set -e        # exit on error
    set -u # error on use of an undefined variable
    set -o pipefail # a pipeline fails if any command in it fails
  2. Helper functions

    1
    2
    3
    4
    5
    6
    7
    8
    # Logging function: prints a timestamp and the caller location
    log() {
    local fname=${BASH_SOURCE[1]##*/}
    echo -e "$(date '+%Y-%m-%dT%H:%M:%S') (${fname}:${BASH_LINENO[0]}:${FUNCNAME[1]}) $*"
    }

    # Min function: used to compute the number of parallel jobs
    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                # stage to start processing from
    stop_stage=10000 # stage to stop processing at
    skip_data_prep=false # whether to skip the data preparation stages
    skip_train=false # whether to skip the training stages
    skip_eval=false # whether to skip the inference and evaluation stages
    skip_packing=true # whether to skip the packing stage
    skip_upload_hf=true # whether to skip the upload-to-HuggingFace stage
    ngpu=1 # GPU count (0 means use CPU)
    num_nodes=1 # number of nodes
    nj=32 # number of parallel jobs
  2. Feature extraction parameters

    1
    2
    3
    4
    5
    feats_type=raw        # feature type (raw or fbank_pitch)
    audio_format=flac # audio format: wav, flac, etc.
    fs=16k # sampling rate
    min_wav_duration=0.1 # minimum audio duration (seconds)
    max_wav_duration=20 # maximum audio duration (seconds)
  3. Enhancement model parameters

    1
    2
    3
    4
    5
    6
    7
    8
    9
    enh_exp=             # path to the enhancement experiment directory
    enh_tag= # suffix for the enhancement model training output directory
    enh_config= # enhancement model training config
    enh_args= # extra arguments for enhancement model training
    ref_num=2 # number of reference signals (equal to the number of speakers)
    inf_num= # number of inference outputs produced by the model
    noise_type_num=1 # number of noise types in the input audio
    dereverb_ref_num=1 # number of dereverberation reference signals
    is_tse_task=false # whether this is a target speaker extraction task
  4. Training data parameters

    1
    2
    3
    use_dereverb_ref=false   # whether to use dereverberation reference signals
    use_noise_ref=false # whether to use noise reference signals
    variable_num_refs=false # whether to use a variable number of reference signals
  5. Inference and evaluation parameters

    1
    2
    3
    inference_args="--normalize_output_wav true --output_format wav"  # inference arguments
    inference_model=valid.loss.ave.pth # model file used for inference
    scoring_protocol="STOI SDR SAR SIR SI_SNR" # scoring metrics

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] for a new corpus you need to create 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

# Set up bash error handling
set -e # exit on error
set -u # error on use of an undefined variable
set -o pipefail # a pipeline fails if any command in it fails

# Define the logging function
log() {
local fname=${BASH_SOURCE[1]##*/}
echo -e "$(date '+%Y-%m-%dT%H:%M:%S') (${fname}:${BASH_LINENO[0]}:${FUNCNAME[1]}) $*"
}

# Help message
help_message=$(cat << EOF
Usage: $0 [--min_or_max <min/max>] [--sample_rate <8k/16k>]
optional argument:
[--min_or_max]: min (Default), max # mixing mode: min or max
[--sample_rate]: 8k (Default), 16k # sampling rate: 8k or 16k
EOF
)

# Source the database configuration
. ./db.sh

# Set path variables
wsj_full_wav=$PWD/data/wsj0/wsj0_wav # WSJ0 original audio path
wsj_2mix_wav=$PWD/data/wsj0_mix/2speakers # two-speaker mixture audio path
wsj_2mix_scripts=$PWD/data/wsj0_mix/scripts # mixing scripts path

# Set text-related variables
other_text=data/local/other_text/text # path to the other text data
nlsyms=data/nlsyms.txt # non-linguistic symbols file
min_or_max=min # default mixing mode is min
sample_rate=8k # default sampling rate is 8k

# Parse command-line arguments
. utils/parse_options.sh

# Check that the WSJ0 and WSJ1 dataset paths exist
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

# Set dataset names
train_set="tr_"${min_or_max}_${sample_rate} # training set
train_dev="cv_"${min_or_max}_${sample_rate} # validation set
recog_set="tt_"${min_or_max}_${sample_rate} # test set

### WSJ0 mixture data processing ###
# Download the mixing scripts and create the two-speaker mixture audio
local/wsj0_create_mixture.sh ${wsj_2mix_scripts} ${WSJ0} ${wsj_full_wav} \
${wsj_2mix_wav} || exit 1;

# Prepare the WSJ0_2mix dataset
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;

### Create .scp files for the reference audio ###
# Create scp files for speaker 1 and speaker 2 for each dataset
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 corpus processing ###
# Prepare the WSJ data
log "local/wsj_data_prep.sh ${WSJ0}/??-{?,??}.? ${WSJ1}/??-{?,??}.?"
local/wsj_data_prep.sh ${WSJ0}/??-{?,??}.? ${WSJ1}/??-{?,??}.?

# Format the WSJ data
log "local/wsj_format_data.sh"
local/wsj_format_data.sh

# Create the wsj directory and move the related data into it
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

# Prepare additional text data
log "Prepare text from lng_modl dir..."
mkdir -p "$(dirname ${other_text})"

# Process the language model training data
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}

# Create the non-linguistic symbols file
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
    # Format the wav.scp file
    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
    # Compute the minimum and maximum lengths (in samples)
    _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}))")

    # Filter the data by length
    <"${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

中文