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
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
# 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 model1
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
5python -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
./enhanceddirectory - 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 | optim: adam # optimizer choice: Adam optimizer |
Optimizer configuration
1 | optim_conf: |
Loss function configuration
1 | # A list for criterions |
Model architecture configuration
1 | encoder: stft # STFT encoder configuration |
Analysis of enh.sh
Configuration before Stage 1
Basic settings
bash debug mode settings
1
2
3set -e # exit on error
set -u # error on use of an undefined variable
set -o pipefail # a pipeline fails if any command in it failsHelper 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
- 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
Basic configuration parameters
1
2
3
4
5
6
7
8
9
10stage=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 jobsFeature extraction parameters
1
2
3
4
5feats_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)Enhancement model parameters
1
2
3
4
5
6
7
8
9enh_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 taskTraining data parameters
1
2
3use_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 signalsInference and evaluation parameters
1
2
3inference_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
5if [ ${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
- Generated under
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 |
|
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}_spdirectory:- The perturbed audio files and the corresponding config files
- Updated wav.scp, utt2spk, spk2utt and other files
- Generated under the
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:
- Re-create the “wav.scp” file
- Unify the audio format and sampling rate
- Handle the multi-speaker case
- 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
- Under the
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:
- Convert durations into numbers of samples
- Filter the audio by number of samples
- 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
- Under the
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:
- Collect statistics of the training and validation data
- Generate shape files
- 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
- Under the
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:
- Set up the training and validation data
- Configure the training parameters
- Support resuming from checkpoints
- 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
- Under the
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:
- Load the trained model
- Run inference on the test sets
- Generate the enhanced audio
- Support GPU inference
- Outputs:
- Under the
${_logdir}/output.JOBdirectory:- enhanced.wav: the enhanced audio file
- keys.JOB.scp: key-value pairs of the processed audio
- Inference logs and result files
- Under the
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.JOBdirectory:- scoring.txt: contains each scoring metric
- score_stats: detailed scoring statistics
- Score distribution plots for each metric
- Under the
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:
- Decode the enhanced audio with the ASR model
- Compute the character error rate (CER) or word error rate (WER)
- Generate the evaluation report
- Outputs:
- Under the
${_logdir}/output.JOBdirectory:- asr_inference.txt: ASR decoding results
- text: the recognized text
- wer.txt/cer.txt: error rate statistics
- Under the
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}/packdirectory:- model.zip: the packed model file
- config.yaml: a copy of the config file
- README.md: model documentation
- Under the
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
- In the HuggingFace repository:
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:
- The names of the training, validation and test sets must be provided
- Set the number of GPUs and parallel jobs reasonably according to your needs
- The execution flow can be controlled via stage and stop_stage
- The evaluation stages offer multiple evaluation methods, including objective metrics and ASR evaluation
Translated from the Chinese original.