NVIDIA Deep Learning Examples
NVIDIA Deep Learning Examples
HiFi-GAN for PyTorch
Resource
NVIDIA Deep Learning Examples
NVIDIA Deep Learning Examples
HiFi-GAN for PyTorch

HiFi-GAN model implements a spectrogram inversion model that allows to synthesize speech waveforms from mel-spectrograms.

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

Parameters

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

FlagDescription
--epochsnumber of epochs (default: 1000)
--learning_ratelearning rate (default: 0.1)
--batch_sizeactual batch size for a single forward-backward step (default: 16)
--grad_accumulationnumber of forward-backward steps over which gradients are accumulated (default: 1)
--ampuse mixed precision training (default: disabled)

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 -h.

Getting the data

The ./scripts/download_dataset.sh script will automatically download and extract the dataset to the ./data/LJSpeech-1.1 directory. The ./scripts/prepare_dataset.sh script will preprocess the dataset by generating split filelists in ./data/filelists directory and extracting mel-spectrograms into the ./data/LJSpeech-1.1/mels directory. Data preparation for LJSpeech-1.1 takes around 3 hours on a CPU.

Dataset guidelines

The LJSpeech dataset has 13,100 clips that amount to about 24 hours of speech of a single, female speaker. Since the original dataset does not define a train/dev/test split of the data, we provide a split in the form of three file lists:

./data/filelists
|-- ljs_audio_train_v3.txt
|-- ljs_audio_test.txt
|-- ljs_audio_val.txt

These files are generated during ./scripts/prepare_dataset.sh script execution.

Multi-dataset

Follow these steps to use datasets different from the default LJSpeech dataset.

  1. Prepare a directory with .wav files.

    ./data/my_dataset
    |-- wavs
    
  2. Prepare filelists with paths to .wav files. They define training/validation split of the data (test is currently unused, but it's a good practice to create it for the final evaluation):

    ./data/filelists
    |-- my-dataset_audio_train.txt
    |-- my-dataset_audio_val.txt
    

    Those filelists should list a single wavefile per line as:

    path/to/file001.wav
    path/to/file002.wav
    ...
    

    Those paths should be relative to the path provided by the --dataset-path option of train.py.

  3. (Optional) Prepare file lists with paths to pre-calculated pitch when doing fine-tuning:

    ./data/filelists
    |-- my-dataset_audio_pitch_text_train.txt
    |-- my-dataset_audio_pitch_text_val.txt
    

In order to use the prepared dataset, pass the following to the train.py script:

--dataset-path ./data/my_dataset` \
--training-files ./data/filelists/my-dataset_audio_text_train.txt \
--validation files ./data/filelists/my-dataset_audio_text_val.txt

Training process

HiFi-GAN is trained to generate waveforms from input mel-spectrograms. During training and validation, the network processes small, random chunks of the input of fixed length.

The training can be started with scripts/train.sh script. Output models, DLLogger logs and TensorBoard logs will be saved in the output/ directory.

The following example output is printed when running the model:

DLL 2021-06-30 10:58:05.828323 - epoch    1 | iter   1/24 | d loss  7.966 | g loss 95.839 | mel loss 87.291 |  3092.31 frames/s | took 13.25 s | g lr 3.00e-04 | d lr 3.00e-04
DLL 2021-06-30 10:58:06.999175 - epoch    1 | iter   2/24 | d loss  7.957 | g loss 96.151 | mel loss 87.627 | 35109.29 frames/s | took 1.17 s | g lr 3.00e-04 | d lr 3.00e-04
DLL 2021-06-30 10:58:07.945764 - epoch    1 | iter   3/24 | d loss  7.956 | g loss 93.872 | mel loss 88.154 | 43443.33 frames/s | took 0.94 s | g lr 3.00e-04 | d lr 3.00e-04

Performance is reported in total input mel-spectrogram frames per second and recorded as train_frames/s (after each iteration) and avg_train_frames/s (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. The metrics are averaged in such a way, that gradient accumulation steps would be transparent to the user.

The scripts/train.sh script is configured for 8x GPU with at least 16GB of memory. In a single accumulated step, there are batch_size x grad_accumulation x GPUs = 16 x 1 x 8 = 128 examples being processed in parallel. With a smaller number of GPUs, increase gradient accumulation steps to keep the relation satisfied, e.g., through env variables

NUM_GPUS=1 GRAD_ACCUMULATION=8 BATCH_SIZE=16 bash scripts/train.sh

The script also enables automatic mixed precision training. To train with mixed precision, specify the AMP variable

AMP=true bash scripts/train.sh

Inference process

You can run inference using the ./inference.py script. This script takes mel-spectrograms as input and runs HiFi-GAN inference to produce audio files.

Pre-trained HiFi-GAN models are available for download on NGC. The latest model can be downloaded with:

scripts/download_model.sh hifigan

Having pre-trained models in place, extract validation mel-spectrograms from the LJSpeech-1.1 test-set, and run inference with:

bash scripts/inference_example.sh

Examine the inference_example.sh script to adjust paths to pre-trained models, and call python inference.py --help to learn all available options. By default, synthesized audio samples are saved in ./output/audio_* folders.