Problem Agnostic Speech Encoder (PASE)

This repository contains the official implementations of PASE and PASE+. These are speech waveform encoders trained in a self-supervised manner with the so called worker/minion framework. A PASE model can be used as a speech feature extractor or to pre-train an encoder for our desired end-task, like speech classification such as in ASR, seaker recognition, or emotion recognition, or speech generation such as in voice conversion or TTS.

Requirements

• PyTorch 1.0 or higher
• Torchvision 0.2 or higher
• To use data augmentation during training (recommended), you must build codec2 from source, then pip install pycodec2 (because pycodec2 needs the header files). You may also need to point LD_LIBRARY_PATH at \usr\local\lib for python to be able to load pycodec2 successfully.
• Install the requirements from requirements.txt: pip install -r requirements.txt

NOTE: Edit the cupy-cuda100 requirement in the file if needed depending on your CUDA version. Defaults to 10.0 now

Install

This framework can be installed locally by running:

python setup.py install

This will allow you to import PASE modules from anywhere.

Pre-trained Model

The PASE+ parameters used in our most recently published work can be found if you CLICK HERE. This ckpt file contains the encoder parameters only, without any worker. This ckpt named FE_e199.ckpt, and the configuration file cfg/frontend/PASE+.cfg let you build and use the encoder in the following simple manner:

from pase.models.frontend import wf_builder
pase = wf_builder('cfg/frontend/PASE+.cfg').eval()
pase.load_pretrained('FE_e199.ckpt', load_last=True, verbose=True)

# Now we can forward waveforms as Torch tensors
import torch
x = torch.randn(1, 1, 100000) # example with random noise to check shape
# y size will be (1, 256, 625), which are 625 frames of 256 dims each
y = pase(x)

The encoder can be inserted in any PyTorch model and fine-tuned, just like any other nn.Module.

Self-Supervised Training Do-It-Yourself (DIY)

Data preparation

The self-supervised training stage requires the following components to be specified to the training script:

• data root folder: contains files (or soft links to them) without subfolders in wav, mp3 or any Torchaudio-supported format.
• trainset statistics file to normalize each worker's output values, computed with the make_trainset_statistics.py script.
• dataset configuration data_cfg file: contains pointers to train/valid/test splits, among other info.
• front-end (encoder) configuration file: cfg/frontend/PASE+.cfg
• workers' configuration file: cfg/workers/workers+.cfg

Making the dataset config file

To make the dataset configuration file the following files have to be provided:

• training files list train_scp: contains a file name per line (without directory names), including .wav/mp3/etc. extension.
• test files list test_scp: contains a wav file name per line (without directory names), including .wav/mp3/etc. extension.
• dictionary with filename -> integer speaker class (speaker id) correspondence (same filenames as in train/test lists).

An example of each of these files can be found in the data/ folder of the repo. Build them based on your data files.

NOTE: The filename2spkclass dictionary is required to create a train/valid/test split which holds out some speakers from training, such that self-supervised training validation tracks the workers' losses with unseen identities (thus to truly generalize). Those labels, however, are not used during training for this is an unsupervised framework.

We use the following script to create our dataset configuration file (--cfg_file):

python unsupervised_data_cfg_librispeech.py --data_root data/LibriSpeech/wavs \
	--train_scp data/LibriSpeech/libri_tr.scp --test_scp data/LibriSpeech/libri_te.scp \
	--libri_dict data/LibriSpeech/libri_dict.npy --cfg_file data/librispeech_data.cfg

Making the trainset statistics file

The make_trainset_statistics.py script will load a certain amount of training batches with the config file we just generated, and will compute the normalization statistics for the workers to work properly in the self-supervised training. For PASE v0.1 we use this script as follows:

python make_trainset_statistics.py --data_root data/LibriSpeech/wavs \
	--data_cfg data/librispeech_data.cfg \
	--net_cfg cfg/workers/workers.cfg \
	--out_file data/librispeech_stats.pkl 

The file data/librispeech_stats.pkl will be generated. If this goes too slow, you may try with a smaller amount of training batches with the --max_batches 10 argument for example. The default is 20. Note that the --net_cfg cfg/workers+.cfg is supplied so that the script automatically retrieves the workers that will be active, and the statistics are specific to the workers.

To build the statistics file for PASE+ (recommended), then we simply use the new worker configuration cfg/workers/workers+.cfg:

python make_trainset_statistics.py --data_root data/LibriSpeech/wavs \
	--data_cfg data/librispeech_data.cfg \
	--net_cfg cfg/workers/workers+.cfg \
	--out_file data/librispeech_stats_pase+.pkl 

Training

To train PASE for 150 epochs, with the same hyper-parameters as those in the first published work, execute the following script:

python -u train.py --batch_size 32 --epoch 150 --save_path pase_ckpt --num_workers 4 \
	--net_cfg cfg/workers/workers.cfg --fe_cfg cfg/frontend/PASE.cfg \
	--data_cfg data/librispeech_data.cfg --min_lr 0.0005 --fe_lr 0.0005 \
	--data_root data/LibriSpeech/wavs/ --stats data/librispeech_stats.pkl --lrdec_step 30 --lrdecay 0.5

Note that data_root, stats and data_cfg are the mentioned data root folder, training statistics file and dataset configuration file (created in previous section). TensorboardX is used during training to dump stats information (stored in save_path folder, together with the model checkpoints). The learning rates min_lr and fe_lr control the worker learning rates and the encoder learning rates respectively. The lrdec_step and lrdecay params control the learning rate decay factor and the periodic step at which it is applied, for all components (workers and PASE).

To replicate PASE+ training, execute the following:

python -u  train.py --batch_size 16 --epoch 400 --save_path pase+_ckpt \
	       --num_workers 4 --warmup 10000000 --net_cfg cfg/workers/workers+.cfg \
	       --fe_cfg cfg/frontend/PASE+.cfg --data_cfg data/librispeech_data.cfg \
	       --min_lr 0.0005 --fe_lr 0.001 --data_root data/LibriSpeech/wavs/ \
	       --dtrans_cfg cfg/distortions/pase+.cfg \
	       --stats data/librispeech_stats_pase+.pkl \
	       --chunk_size 32000 \
	       --tensorboard False \
	       --backprop_mode base\
	       --random_scale True\
	       --lr_mode poly

Note that the --lr_mode allows to choose a different learning rate scheduler. In the poly case, a polynomial scheduler updates the LR to reach zero in the end of the programmed epochs.

The --dtrans_cfg flag controls the pointer to the configuration of data augmentation distortions in the form of additive noises, reverberations, etc.

Distortions Configuration

The configuration for the distortions (supplied with the --dtrans_cfg argument) allows to control the probability of a distortion being active for a sample in the batch. Hence, distortions are applied on the fly and independently, although with a hard-coded order as programmed in file pase/transforms.py (i.e. Reverb happens before Additive, etc.). Note that there are possible distortions:

• Overlap: activated with overlap_p > 0 . This overlaps random chunks of speech from the selected directory of speech samples, emulating background speakers with the specified SNRs in overlap_snrs (picked randomly).
• Additive noise: activated with noises_p > 0. Selects a noise file from the specified directories and applies a random SNR out of the possible values.
• Amplitude clipping: activated with clip_p > 0. Clips the waveform amplitude on values beyond a specified percentage of the maximum peak (e.g. 0.1value means clamp all values exceeding on absolute amplitude the value 0.1 x max_asbsolute_amplitude).
• Waveform chopping: activated with chop_p > 0. Chops continuous sections of speech by building windows randomly sized following a Gaussian pdf with the values specified as tuples in the chop_factors array. For instance, [0.05, 0.025] means sampling a window of size 0.05 sec on average with 0.025 sec standard deviation. Many Gaussian parameterizations can be supplied to have windows of different sizes on average, which are then sampled uniformly random.
• Waveform resampling: activated with downsample_p > 0. Resample the signal to make it narrowband.
• Frequency band-drop: activated with bandrop_p > 0. Apply random bandpass filters to equalize the spectrogram per bands.
• Reverberation: activated with reverb_p > 0.

Each distortion has a set of parameters that can be controlled, like the impulse response files used to emulate reverberation or pointers to the directories where additive noises are found and the SNRs to be applied randomly. The file cfg/distortions/pase+.cfg exemplifies all the possible options to be controlled for the different distortions.

If no --dtrans_cfg file is provided, the waveforms are loaded as-is without any change except for a possible random scaling in case `--random_scal