The following sections provide greater details of the dataset, running training and inference, and the training results.

Scripts and sample code

The sample code for Tacotron 2 and WaveGlow has scripts specific to a particular model, located in directories ./tacotron2 and ./waveglow, as well as scripts common to both models, located in the ./common directory. The model-specific scripts are as follows:

• <model_name>/model.py - the model architecture, definition of forward and inference functions
• <model_name>/arg_parser.py - argument parser for parameters specific to a given model
• <model_name>/data_function.py - data loading functions
• <model_name>/loss_function.py - loss function for the model

The common scripts contain layer definitions common to both models (common/layers.py), some utility scripts (common/utils.py) and scripts for audio processing (common/audio_processing.py and common/stft.py). In the root directory ./ of this repository, the ./run.py script is used for training while inference can be executed with the ./inference.py script. The scripts ./models.py, ./data_functions.py and ./loss_functions.py call the respective scripts in the <model_name> directory, depending on what model is trained using the run.py script.

Parameters

In this section, we list the most important hyperparameters and command-line arguments, together with their default values that are used to train Tacotron 2 and WaveGlow models.

Shared parameters

• --epochs - number of epochs (Tacotron 2: 1501, WaveGlow: 1001)
• --learning-rate - learning rate (Tacotron 2: 1e-3, WaveGlow: 1e-4)
• --batch-size - batch size (Tacotron 2 FP16/FP32: 104/48, WaveGlow FP16/FP32: 10/4)
• --amp - use mixed precision training
• --cpu - use CPU with TorchScript for inference

Shared audio/STFT parameters

• --sampling-rate - sampling rate in Hz of input and output audio (22050)
• --filter-length - (1024)
• --hop-length - hop length for FFT, i.e., sample stride between consecutive FFTs (256)
• --win-length - window size for FFT (1024)
• --mel-fmin - lowest frequency in Hz (0.0)
• --mel-fmax - highest frequency in Hz (8.000)

Tacotron 2 parameters

• --anneal-steps - epochs at which to anneal the learning rate (500 1000 1500)
• --anneal-factor - factor by which to anneal the learning rate (FP16/FP32: 0.3/0.1)

WaveGlow parameters

• --segment-length - segment length of input audio processed by the neural network (8000)
• --wn-channels - number of residual channels in the coupling layer networks (512)

Command-line options

To see the full list of available options and their descriptions, use the -h or --help command line option, for example:

python train.py --help

The following example output is printed when running the sample:

Batch: 7/260 epoch 0
:::NVLOGv0.2.2 Tacotron2_PyT 1560936205.667271376 (/workspace/tacotron2/dllogger/logger.py:251) train_iter_start: 7
:::NVLOGv0.2.2 Tacotron2_PyT 1560936207.209611416 (/workspace/tacotron2/dllogger/logger.py:251) train_iteration_loss: 5.415428161621094
:::NVLOGv0.2.2 Tacotron2_PyT 1560936208.705905914 (/workspace/tacotron2/dllogger/logger.py:251) train_iter_stop: 7
:::NVLOGv0.2.2 Tacotron2_PyT 1560936208.706479311 (/workspace/tacotron2/dllogger/logger.py:251) train_iter_items/sec: 8924.00136085362
:::NVLOGv0.2.2 Tacotron2_PyT 1560936208.706998110 (/workspace/tacotron2/dllogger/logger.py:251) iter_time: 3.0393316745758057
Batch: 8/260 epoch 0
:::NVLOGv0.2.2 Tacotron2_PyT 1560936208.711485624 (/workspace/tacotron2/dllogger/logger.py:251) train_iter_start: 8
:::NVLOGv0.2.2 Tacotron2_PyT 1560936210.236668825 (/workspace/tacotron2/dllogger/logger.py:251) train_iteration_loss: 5.516331672668457

Getting the data

The Tacotron 2 and WaveGlow models were trained on the LJSpeech-1.1 dataset. This repository contains the ./scripts/prepare_dataset.sh script which will automatically download and extract the whole dataset. By default, data will be extracted to the ./LJSpeech-1.1 directory. The dataset directory contains a README file, a wavs directory with all audio samples, and a file metadata.csv that contains audio file names and the corresponding transcripts.

Dataset guidelines

The LJSpeech dataset has 13,100 clips that amount to about 24 hours of speech. Since the original dataset has all transcripts in the metadata.csv file, in this repository we provide file lists in the ./filelists directory that determine training and validation subsets; ljs_audio_text_train_filelist.txt is a test set used as a training dataset and ljs_audio_text_val_filelist.txt is a test set used as a validation dataset.

Multi-dataset

To use datasets different than the default LJSpeech dataset:

1. Prepare a directory with all audio files and pass it to the --dataset-path command-line option.

2. Add two text files containing file lists: one for the training subset (--training-files) and one for the validation subset (--validation files). The structure of the filelists should be as follows:

   `<audio file path>|<transcript>`
   

   The <audio file path> is the relative path to the path provided by the --dataset-path option.

Training process

The Tacotron2 and WaveGlow models are trained separately and independently. Both models obtain mel-spectrograms from short time Fourier transform (STFT) during training. These mel-spectrograms are used for loss computation in case of Tacotron 2 and as conditioning input to the network in case of WaveGlow.

The training loss is averaged over an entire training epoch, whereas the validation loss is averaged over the validation dataset. Performance is reported in total output mel-spectrograms per second for the Tacotron 2 model and in total output samples per second for the WaveGlow model. Both measures are recorded as train_iter_items/sec (after each iteration) and train_epoch_items/sec (averaged over epoch) in the output log file ./output/nvlog.json. The result is averaged over an entire training epoch and summed over all GPUs that were included in the training.

Even though the training script uses all available GPUs, you can change this behavior by setting the CUDA_VISIBLE_DEVICES variable in your environment or by setting the NV_GPU variable at the Docker container launch (see section "GPU isolation").

Inference process

You can run inference using the ./inference.py script. This script takes text as input and runs Tacotron 2 and then WaveGlow inference to produce an audio file. It requires pre-trained checkpoints from Tacotron 2 and WaveGlow models and input text as a text file, with one phrase per line.

To run inference, issue:

python inference.py --tacotron2 <Tacotron2_checkpoint> --waveglow <WaveGlow_checkpoint> --wn-channels 256 -o output/ --include-warmup -i phrases/phrase.txt --fp16

Here, Tacotron2_checkpoint and WaveGlow_checkpoint are pre-trained checkpoints for the respective models, and phrases/phrase.txt contains input phrases. The number of text lines determines the inference batch size. Audio will be saved in the output folder. The audio files audio_fp16 and audio_fp32 were generated using checkpoints from mixed precision and FP32 training, respectively.

You can find all the available options by calling python inference.py --help.

You can also run inference on CPU with TorchScript by adding flag --cpu:

export CUDA_VISIBLE_DEVICES=

python inference.py --tacotron2 <Tacotron2_checkpoint> --waveglow <WaveGlow_checkpoint> --wn-channels 256 --cpu -o output/ -i phrases/phrase.txt