tutorials/tts/Tacotron2_Training.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "htbJiaJjYQAD"
   },
   "source": [
    "# Tacotron 2 Training\n",
    "\n",
    "This notebook is designed to provide a guide on how to train Tacotron2 as part of the TTS pipeline. It contains the following sections\n",
    "\n",
    "  1. Tacotron2 and NeMo - An introduction to the Tacotron2 model\n",
    "  2. LJSpeech - How to train Tacotron2 on LJSpeech\n",
    "  3. Custom Datasets - How to collect audio data to train Tacotron2 for difference voices and languages"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "wqPMTEXXYUP4"
   },
   "source": [
    "# License\n",
    "\n",
    "> Copyright 2020 NVIDIA. All Rights Reserved.\n",
    "> \n",
    "> Licensed under the Apache License, Version 2.0 (the \"License\");\n",
    "> you may not use this file except in compliance with the License.\n",
    "> You may obtain a copy of the License at\n",
    "> \n",
    ">     http://www.apache.org/licenses/LICENSE-2.0\n",
    "> \n",
    "> Unless required by applicable law or agreed to in writing, software\n",
    "> distributed under the License is distributed on an \"AS IS\" BASIS,\n",
    "> WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
    "> See the License for the specific language governing permissions and\n",
    "> limitations under the License."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "SUkq9HAvYU7T"
   },
   "outputs": [],
   "source": [
    "\"\"\"\n",
    "You can either run this notebook locally (if you have all the dependencies and a GPU) or on Google Colab.\n",
    "Instructions for setting up Colab are as follows:\n",
    "1. Open a new Python 3 notebook.\n",
    "2. Import this notebook from GitHub (File -> Upload Notebook -> \"GITHUB\" tab -> copy/paste GitHub URL)\n",
    "3. Connect to an instance with a GPU (Runtime -> Change runtime type -> select \"GPU\" for hardware accelerator)\n",
    "4. Run this cell to set up dependencies# .\n",
    "\"\"\"\n",
    "BRANCH = 'r1.17.0'\n",
    "# # If you're using Colab and not running locally, uncomment and run this cell.\n",
    "# !apt-get install sox libsndfile1 ffmpeg\n",
    "# !pip install wget text-unidecode\n",
    "# !python -m pip install git+https://github.com/NVIDIA/NeMo.git@$BRANCH#egg=nemo_toolkit[all]\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "ZivXzmq0YYLj"
   },
   "source": [
    "# Tacotron2 and NeMo\n",
    "\n",
    "Tacotron2 is a neural network that converts text characters into a mel spectrogram. For more details on the model, please refer to Nvidia's [Tacotron2 Model Card](https://ngc.nvidia.com/catalog/models/nvidia:nemo:tts_en_tacotron2), or the original [paper](https://arxiv.org/abs/1712.05884).\n",
    "\n",
    "Tacotron2 like most NeMo models are defined as a LightningModule, allowing for easy training via PyTorch Lightning, and parameterized by a configuration, currently defined via a yaml file and loading using Hydra.\n",
    "\n",
    "Let's take a look using NeMo's pretrained model and how to use it to generate spectrograms."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "HEvdSU5WYZbj"
   },
   "outputs": [],
   "source": [
    "# Load the Tacotron2Model\n",
    "from nemo.collections.tts.models import Tacotron2Model\n",
    "from nemo.collections.tts.models.base import SpectrogramGenerator\n",
    "\n",
    "# Let's see what pretrained models are available\n",
    "print(Tacotron2Model.list_available_models())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "3W8unatgYbUp"
   },
   "outputs": [],
   "source": [
    "# We can load the pre-trained model as follows\n",
    "model = Tacotron2Model.from_pretrained(\"tts_en_tacotron2\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "xsyBa9tIdHp4"
   },
   "outputs": [],
   "source": [
    "# Tacotron2 is a SpectrogramGenerator\n",
    "assert isinstance(model, SpectrogramGenerator)\n",
    "\n",
    "# SpectrogramGenerators in NeMo have two helper functions:\n",
    "#   1. parse(self, text: str, normalize=True) which takes an English string and produces a token tensor\n",
    "#   2. generate_spectrogram(self, *, tokens) which takes the token tensor and generates a spectrogram\n",
    "# Let's try it out\n",
    "tokens = model.parse(text = \"Hey, this produces speech!\")\n",
    "spectrogram = model.generate_spectrogram(tokens = tokens)\n",
    "\n",
    "# Now we can visualize the generated spectrogram\n",
    "# If we want to generate speech, we have to use a vocoder in conjunction to a spectrogram generator.\n",
    "# Refer to the TTS Inference notebook on how to convert spectrograms to speech.\n",
    "from matplotlib.pyplot import imshow\n",
    "from matplotlib import pyplot as plt\n",
    "%matplotlib inline\n",
    "imshow(spectrogram.cpu().detach().numpy()[0,...], origin=\"lower\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "zZ90eCfdrNIf"
   },
   "source": [
    "# Training\n",
    "\n",
    "Now that we looked at the Tacotron2 model, let's see how to train a Tacotron2 Model\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "7rHG-LERrPRY"
   },
   "outputs": [],
   "source": [
    "# NeMo's training scripts are stored inside the examples/ folder. Let's grab the tacotron2.py file\n",
    "# as well as the tacotron2.yaml file\n",
    "!wget https://raw.githubusercontent.com/NVIDIA/NeMo/$BRANCH/examples/tts/tacotron2.py\n",
    "!(mkdir -p conf \\\n",
    "  && cd conf \\\n",
    "  && wget https://raw.githubusercontent.com/NVIDIA/NeMo/$BRANCH/examples/tts/conf/tacotron2.yaml \\\n",
    "  && cd ..)\n",
    "\n",
    "# We will also need a few extra files for handling text.\n",
    "!(mkdir -p scripts/tts_dataset_files \\\n",
    "  && cd scripts/tts_dataset_files \\\n",
    "  && wget https://raw.githubusercontent.com/NVIDIA/NeMo/$BRANCH/scripts/tts_dataset_files/cmudict-0.7b_nv22.10 \\\n",
    "  && wget https://raw.githubusercontent.com/NVIDIA/NeMo/$BRANCH/scripts/tts_dataset_files/heteronyms-052722 \\\n",
    "  && cd ..)\n",
    "        \n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "Upv_LxBIsC51"
   },
   "source": [
    "Let's take a look at the tacotron2.py file\n",
    "\n",
    "```python\n",
    "import pytorch_lightning as pl\n",
    "\n",
    "from nemo.collections.common.callbacks import LogEpochTimeCallback\n",
    "from nemo.collections.tts.models import Tacotron2Model\n",
    "from nemo.core.config import hydra_runner\n",
    "from nemo.utils.exp_manager import exp_manager\n",
    "\n",
    "\n",
    "# hydra_runner is a thin NeMo wrapper around Hydra\n",
    "# It looks for a config named tacotron2.yaml inside the conf folder\n",
    "# Hydra parses the yaml and returns it as a Omegaconf DictConfig\n",
    "@hydra_runner(config_path=\"conf\", config_name=\"tacotron2\")\n",
    "def main(cfg):\n",
    "    # Define the Lightning trainer\n",
    "    trainer = pl.Trainer(**cfg.trainer)\n",
    "    # exp_manager is a NeMo construct that helps with logging and checkpointing\n",
    "    exp_manager(trainer, cfg.get(\"exp_manager\", None))\n",
    "    # Define the Tacotron 2 model, this will construct the model as well as\n",
    "    # define the training and validation dataloaders\n",
    "    model = Tacotron2Model(cfg=cfg.model, trainer=trainer)\n",
    "    # Let's add a few more callbacks\n",
    "    lr_logger = pl.callbacks.LearningRateMonitor()\n",
    "    epoch_time_logger = LogEpochTimeCallback()\n",
    "    trainer.callbacks.extend([lr_logger, epoch_time_logger])\n",
    "    # Call lightning trainer's fit() to train the model\n",
    "    trainer.fit(model)\n",
    "\n",
    "\n",
    "if __name__ == '__main__':\n",
    "    main()  # noqa pylint: disable=no-value-for-parameter\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "6nM-fZO-s75u"
   },
   "source": [
    "Let's take a look at the yaml config\n",
    "\n",
    "```yaml\n",
    "name: &name Tacotron2\n",
    "\n",
    "train_dataset: ???\n",
    "validation_datasets: ???\n",
    "sup_data_path: null\n",
    "sup_data_types: null\n",
    "\n",
    "phoneme_dict_path: \"scripts/tts_dataset_files/cmudict-0.7b_nv22.10\"\n",
    "heteronyms_path: \"scripts/tts_dataset_files/heteronyms-052722\"\n",
    "```\n",
    "\n",
    "The first part of the yaml defines dataset parameters used by Tacotron. Then in the head of 'model' section there are processing - related parameters. You can see\n",
    "that the sample rate is set to 22050 for LJSpeech. \n",
    "\n",
    "Looking at the yaml, there is `train_dataset: ???` and `validation_datasets: ???`. The ??? indicates to hydra that these values must be passed via the command line or the script will fail.\n",
    "\n",
    "Looking further down the yaml, we get to the pytorch lightning trainer parameters.\n",
    "\n",
    "```yaml\n",
    "trainer:\n",
    "  devices: 1 # number of gpus\n",
    "  accelerator: 'gpu' \n",
    "  max_epochs: ???\n",
    "  num_nodes: 1\n",
    "  accelerator: 'gpu'\n",
    "  strategy: 'ddp'\n",
    "  accumulate_grad_batches: 1\n",
    "  enable_checkpointing: False  # Provided by exp_manager\n",
    "  logger: False  # Provided by exp_manager\n",
    "  gradient_clip_val: 1.0\n",
    "  log_every_n_steps: 200\n",
    "  check_val_every_n_epoch: 25\n",
    "```\n",
    "\n",
    "These values can be changed either by editing the yaml or through the command line.\n",
    "\n",
    "Let's grab some simple audio data and test Tacotron2."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "GnEzODcorugt"
   },
   "outputs": [],
   "source": [
    "!wget https://github.com/NVIDIA/NeMo/releases/download/v0.11.0/test_data.tar.gz \\\n",
    "&& mkdir -p tests/data \\\n",
    "&& tar xzf test_data.tar.gz -C tests/data\n",
    "\n",
    "# Just like ASR, the Tacotron2 require .json files to define the training and validation data.\n",
    "!cat tests/data/asr/an4_val.json"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now that we have some sample data, we can try training Tacotron 2!\n",
    "\n",
    "Note that the sample data is not enough data to fully train a Tacotron 2 model. The following code uses a toy dataset to illustrate how the pipeline for training would work."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!(python tacotron2.py \\\n",
    "  model.sample_rate=16000 \\\n",
    "  train_dataset=tests/data/asr/an4_train.json \\\n",
    "  validation_datasets=tests/data/asr/an4_val.json \\\n",
    "  trainer.max_epochs=3 \\\n",
    "  trainer.accelerator=null \\\n",
    "  trainer.check_val_every_n_epoch=1 \\\n",
    " +trainer.gpus=1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "9erGDGZJ1H_p"
   },
   "source": [
    "# Training Data\n",
    "\n",
    "In order to train Tacotron2, it is highly recommended to obtain high quality speech data with the following properties:\n",
    "  - Sampling rate of 22050Hz or higher\n",
    "  - Single speaker\n",
    "  - Speech should contain a variety of speech phonemes\n",
    "  - Audio split into segments of 1-10 seconds\n",
    "  - Audio segments should not have silence at the beginning and end\n",
    "  - Audio segments should not contain long silences inside\n",
    "\n",
    "After obtaining the speech data and splitting into training, validation, and test sections, it is required to construct .json files to tell NeMo where to find these audio files.\n",
    "\n",
    "The .json files should adhere to the format required by the `nemo.collections.tts.data.tts_dataset.TTSDataset` class. For example, here is a sample .json file\n",
    "\n",
    "```json\n",
    "{\"audio_filepath\": \"/path/to/audio1.wav\", \"text\": \"the transcription\", \"duration\": 0.82}\n",
    "{\"audio_filepath\": \"/path/to/audio2.wav\", \"text\": \"the other transcription\", \"duration\": 2.1}\n",
    "...\n",
    "```\n",
    "Please note that the duration is in seconds.\n",
    "\n",
    "\n",
    "Then you are ready to run your training script:\n",
    "```bash\n",
    "python tacotron2.py train_dataset=YOUR_TRAIN.json validation_datasets=YOUR_VAL.json trainer.devices=-1\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": []
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "colab": {
   "collapsed_sections": [],
   "name": "Taco2.ipynb",
   "provenance": []
  },
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}