Tacotron2_Training.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "Tacotron2_Training.ipynb",
      "provenance": [],
      "collapsed_sections": []
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "accelerator": "GPU"
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ODI651CWUdFs"
      },
      "source": [
        "# **Tacotron 2 Training Colab**\n",
        "\n",
        "\n",
        "### Put together by Justin John\n",
        "\n",
        "##### Tacotron 2: https://github.com/NVIDIA/tacotron2\n",
        "\n",
        "## Synthesis notebook: For testing your model\n",
        "### https://colab.research.google.com/github/justinjohn0306/Tacotron2-Colab/blob/main/Tacotron2_synthesis.ipynb\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "iQBUZTuMogui"
      },
      "source": [
        "# Optional cells (unhide by clicking the arrow to the left)"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "75SFQOP_Wnpz",
        "cellView": "form"
      },
      "source": [
        "#@markdown # Check GPU type\n",
        "#@markdown ### Factory reset runtime if you don't have the desired GPU.\n",
        "\n",
        "#@markdown ---\n",
        "\n",
        "#@markdown ## It is recommended to not use the K80 as it does not produce audio correctly.\n",
        "\n",
        "!nvidia-smi -L\n",
        "#@markdown P100 = Suitable\n",
        "\n",
        "#@markdown V100 = Suitable\n",
        "\n",
        "#@markdown T4 = Suitable\n",
        "\n",
        "#@markdown P4 = Suitable\n",
        "\n",
        "#@markdown K80 = Not Suitable, Doesn't produce audio.\n",
        "\n",
        "#@markdown ---"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "cellView": "form",
        "id": "VKYvFCrUhI4O"
      },
      "source": [
        "#@markdown # Anti-Disconnect for Google Colab\n",
        "#@markdown ## Run this to stop it from disconnecting automatically (will disconnect after 6 - 12 hours, though.)\n",
        "\n",
        "import IPython\n",
        "js_code = '''\n",
        "function ClickConnect(){\n",
        "console.log(\"Working\");\n",
        "document.querySelector(\"colab-toolbar-button#connect\").click()\n",
        "}\n",
        "setInterval(ClickConnect,60000)\n",
        "'''\n",
        "display(IPython.display.Javascript(js_code))"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "OhZSeuWHoj6g"
      },
      "source": [
        "# Training"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "z6aq6vRzqEV9",
        "cellView": "form"
      },
      "source": [
        "#@markdown ## Mount your Google Drive (optional but is recommended)\n",
        "\n",
        "#Google Drive Authentication Token\n",
        "from google.colab import drive\n",
        "drive.mount('drive')"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "8mC6GTuQO1Qe",
        "cellView": "form"
      },
      "source": [
        "#@markdown # After running this cell, insert your wav files (must be 22.05Khz) into the wavs folder inside tacotron/ and transcript file into filelists folder\n",
        "\n",
        "#@markdown #### This cell will install Tacotron2 and Dependancies\n",
        "%tensorflow_version 1.x\n",
        "import os\n",
        "!git clone -q https://github.com/NVIDIA/tacotron2\n",
        "os.chdir('tacotron2')\n",
        "!git submodule init\n",
        "!git submodule update\n",
        "!pip install -q unidecode tensorboardX\n",
        "\n",
        "data_path = 'wavs'\n",
        "!mkdir {data_path}"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "c6zQ-DWo4lcs",
        "cellView": "form"
      },
      "source": [
        "#@markdown # Warm up Tacotron 2 A.I \n",
        "\n",
        "%matplotlib inline\n",
        "import os\n",
        "if os.getcwd() != '/content/tacotron2':\n",
        "    os.chdir('tacotron2')\n",
        "import time\n",
        "import argparse\n",
        "import math\n",
        "from numpy import finfo\n",
        "\n",
        "import torch\n",
        "from distributed import apply_gradient_allreduce\n",
        "import torch.distributed as dist\n",
        "from torch.utils.data.distributed import DistributedSampler\n",
        "from torch.utils.data import DataLoader\n",
        "\n",
        "from model import Tacotron2\n",
        "from data_utils import TextMelLoader, TextMelCollate\n",
        "from loss_function import Tacotron2Loss\n",
        "from logger import Tacotron2Logger\n",
        "from hparams import create_hparams\n",
        " \n",
        "import random\n",
        "import numpy as np\n",
        "\n",
        "import layers\n",
        "from utils import load_wav_to_torch, load_filepaths_and_text\n",
        "from text import text_to_sequence\n",
        "from math import e\n",
        "#from tqdm import tqdm # Terminal\n",
        "#from tqdm import tqdm_notebook as tqdm # Legacy Notebook TQDM\n",
        "from tqdm.notebook import tqdm # Modern Notebook TQDM\n",
        "from distutils.dir_util import copy_tree\n",
        "import matplotlib.pylab as plt\n",
        "\n",
        "def download_from_google_drive(file_id, file_name):\n",
        "  # download a file from the Google Drive link\n",
        "  !rm -f ./cookie\n",
        "  !curl -c ./cookie -s -L \"https://drive.google.com/uc?export=download&id={file_id}\" > /dev/null\n",
        "  confirm_text = !awk '/download/ {print $NF}' ./cookie\n",
        "  confirm_text = confirm_text[0]\n",
        "  !curl -Lb ./cookie \"https://drive.google.com/uc?export=download&confirm={confirm_text}&id={file_id}\" -o {file_name}\n",
        "\n",
        "def create_mels():\n",
        "    print(\"Generating Mels\")\n",
        "    stft = layers.TacotronSTFT(\n",
        "                hparams.filter_length, hparams.hop_length, hparams.win_length,\n",
        "                hparams.n_mel_channels, hparams.sampling_rate, hparams.mel_fmin,\n",
        "                hparams.mel_fmax)\n",
        "    def save_mel(filename):\n",
        "        audio, sampling_rate = load_wav_to_torch(filename)\n",
        "        if sampling_rate != stft.sampling_rate:\n",
        "            raise ValueError(\"{} {} SR doesn't match target {} SR\".format(filename, \n",
        "                sampling_rate, stft.sampling_rate))\n",
        "        audio_norm = audio / hparams.max_wav_value\n",
        "        audio_norm = audio_norm.unsqueeze(0)\n",
        "        audio_norm = torch.autograd.Variable(audio_norm, requires_grad=False)\n",
        "        melspec = stft.mel_spectrogram(audio_norm)\n",
        "        melspec = torch.squeeze(melspec, 0).cpu().numpy()\n",
        "        np.save(filename.replace('.wav', ''), melspec)\n",
        "\n",
        "    import glob\n",
        "    wavs = glob.glob('wavs/*.wav')\n",
        "    for i in tqdm(wavs):\n",
        "        save_mel(i)\n",
        "\n",
        "\n",
        "def reduce_tensor(tensor, n_gpus):\n",
        "    rt = tensor.clone()\n",
        "    dist.all_reduce(rt, op=dist.reduce_op.SUM)\n",
        "    rt /= n_gpus\n",
        "    return rt\n",
        "\n",
        "\n",
        "def init_distributed(hparams, n_gpus, rank, group_name):\n",
        "    assert torch.cuda.is_available(), \"Distributed mode requires CUDA.\"\n",
        "    print(\"Initializing Distributed\")\n",
        "\n",
        "    # Set cuda device so everything is done on the right GPU.\n",
        "    torch.cuda.set_device(rank % torch.cuda.device_count())\n",
        "\n",
        "    # Initialize distributed communication\n",
        "    dist.init_process_group(\n",
        "        backend=hparams.dist_backend, init_method=hparams.dist_url,\n",
        "        world_size=n_gpus, rank=rank, group_name=group_name)\n",
        "\n",
        "    print(\"Done initializing distributed\")\n",
        "\n",
        "\n",
        "def prepare_dataloaders(hparams):\n",
        "    # Get data, data loaders and collate function ready\n",
        "    trainset = TextMelLoader(hparams.training_files, hparams)\n",
        "    valset = TextMelLoader(hparams.validation_files, hparams)\n",
        "    collate_fn = TextMelCollate(hparams.n_frames_per_step)\n",
        "\n",
        "    if hparams.distributed_run:\n",
        "        train_sampler = DistributedSampler(trainset)\n",
        "        shuffle = False\n",
        "    else:\n",
        "        train_sampler = None\n",
        "        shuffle = True\n",
        "\n",
        "    train_loader = DataLoader(trainset, num_workers=1, shuffle=shuffle,\n",
        "                              sampler=train_sampler,\n",
        "                              batch_size=hparams.batch_size, pin_memory=False,\n",
        "                              drop_last=True, collate_fn=collate_fn)\n",
        "    return train_loader, valset, collate_fn\n",
        "\n",
        "\n",
        "def prepare_directories_and_logger(output_directory, log_directory, rank):\n",
        "    if rank == 0:\n",
        "        if not os.path.isdir(output_directory):\n",
        "            os.makedirs(output_directory)\n",
        "            os.chmod(output_directory, 0o775)\n",
        "        logger = Tacotron2Logger(os.path.join(output_directory, log_directory))\n",
        "    else:\n",
        "        logger = None\n",
        "    return logger\n",
        "\n",
        "\n",
        "def load_model(hparams):\n",
        "    model = Tacotron2(hparams).cuda()\n",
        "    if hparams.fp16_run:\n",
        "        model.decoder.attention_layer.score_mask_value = finfo('float16').min\n",
        "\n",
        "    if hparams.distributed_run:\n",
        "        model = apply_gradient_allreduce(model)\n",
        "\n",
        "    return model\n",
        "\n",
        "\n",
        "def warm_start_model(checkpoint_path, model, ignore_layers):\n",
        "    assert os.path.isfile(checkpoint_path)\n",
        "    print(\"Warm starting model from checkpoint '{}'\".format(checkpoint_path))\n",
        "    checkpoint_dict = torch.load(checkpoint_path, map_location='cpu')\n",
        "    model_dict = checkpoint_dict['state_dict']\n",
        "    if len(ignore_layers) > 0:\n",
        "        model_dict = {k: v for k, v in model_dict.items()\n",
        "                      if k not in ignore_layers}\n",
        "        dummy_dict = model.state_dict()\n",
        "        dummy_dict.update(model_dict)\n",
        "        model_dict = dummy_dict\n",
        "    model.load_state_dict(model_dict)\n",
        "    return model\n",
        "\n",
        "\n",
        "def load_checkpoint(checkpoint_path, model, optimizer):\n",
        "    assert os.path.isfile(checkpoint_path)\n",
        "    print(\"Loading checkpoint '{}'\".format(checkpoint_path))\n",
        "    checkpoint_dict = torch.load(checkpoint_path, map_location='cpu')\n",
        "    model.load_state_dict(checkpoint_dict['state_dict'])\n",
        "    optimizer.load_state_dict(checkpoint_dict['optimizer'])\n",
        "    learning_rate = checkpoint_dict['learning_rate']\n",
        "    iteration = checkpoint_dict['iteration']\n",
        "    print(\"Loaded checkpoint '{}' from iteration {}\" .format(\n",
        "        checkpoint_path, iteration))\n",
        "    return model, optimizer, learning_rate, iteration\n",
        "\n",
        "\n",
        "def save_checkpoint(model, optimizer, learning_rate, iteration, filepath):\n",
        "    print(\"Saving model and optimizer state at iteration {} to {}\".format(\n",
        "        iteration, filepath))\n",
        "    try:\n",
        "        torch.save({'iteration': iteration,\n",
        "                'state_dict': model.state_dict(),\n",
        "                'optimizer': optimizer.state_dict(),\n",
        "                'learning_rate': learning_rate}, filepath)\n",
        "    except KeyboardInterrupt:\n",
        "        print(\"interrupt received while saving, waiting for save to complete.\")\n",
        "        torch.save({'iteration': iteration,'state_dict': model.state_dict(),'optimizer': optimizer.state_dict(),'learning_rate': learning_rate}, filepath)\n",
        "    print(\"Model Saved\")\n",
        "\n",
        "def plot_alignment(alignment, info=None):\n",
        "    %matplotlib inline\n",
        "    fig, ax = plt.subplots(figsize=(int(alignment_graph_width/100), int(alignment_graph_height/100)))\n",
        "    im = ax.imshow(alignment, cmap='inferno', aspect='auto', origin='lower',\n",
        "                   interpolation='none')\n",
        "    ax.autoscale(enable=True, axis=\"y\", tight=True)\n",
        "    fig.colorbar(im, ax=ax)\n",
        "    xlabel = 'Decoder timestep'\n",
        "    if info is not None:\n",
        "        xlabel += '\\n\\n' + info\n",
        "    plt.xlabel(xlabel)\n",
        "    plt.ylabel('Encoder timestep')\n",
        "    plt.tight_layout()\n",
        "    fig.canvas.draw()\n",
        "    plt.show()\n",
        "\n",
        "def validate(model, criterion, valset, iteration, batch_size, n_gpus,\n",
        "             collate_fn, logger, distributed_run, rank, epoch, start_eposh, learning_rate):\n",
        "    \"\"\"Handles all the validation scoring and printing\"\"\"\n",
        "    model.eval()\n",
        "    with torch.no_grad():\n",
        "        val_sampler = DistributedSampler(valset) if distributed_run else None\n",
        "        val_loader = DataLoader(valset, sampler=val_sampler, num_workers=1,\n",
        "                                shuffle=False, batch_size=batch_size,\n",
        "                                pin_memory=False, collate_fn=collate_fn)\n",
        "\n",
        "        val_loss = 0.0\n",
        "        for i, batch in enumerate(val_loader):\n",
        "            x, y = model.parse_batch(batch)\n",
        "            y_pred = model(x)\n",
        "            loss = criterion(y_pred, y)\n",
        "            if distributed_run:\n",
        "                reduced_val_loss = reduce_tensor(loss.data, n_gpus).item()\n",
        "            else:\n",
        "                reduced_val_loss = loss.item()\n",
        "            val_loss += reduced_val_loss\n",
        "        val_loss = val_loss / (i + 1)\n",
        "\n",
        "    model.train()\n",
        "    if rank == 0:\n",
        "        print(\"Epoch: {} Validation loss {}: {:9f}  Time: {:.1f}m LR: {:.6f}\".format(epoch, iteration, val_loss,(time.perf_counter()-start_eposh)/60, learning_rate))\n",
        "        logger.log_validation(val_loss, model, y, y_pred, iteration)\n",
        "        if hparams.show_alignments:\n",
        "            %matplotlib inline\n",
        "            _, mel_outputs, gate_outputs, alignments = y_pred\n",
        "            idx = random.randint(0, alignments.size(0) - 1)\n",
        "            plot_alignment(alignments[idx].data.cpu().numpy().T)\n",
        "\n",
        "def train(output_directory, log_directory, checkpoint_path, warm_start, n_gpus,\n",
        "          rank, group_name, hparams, log_directory2):\n",
        "    \"\"\"Training and validation logging results to tensorboard and stdout\n",
        "\n",
        "    Params\n",
        "    ------\n",
        "    output_directory (string): directory to save checkpoints\n",
        "    log_directory (string) directory to save tensorboard logs\n",
        "    checkpoint_path(string): checkpoint path\n",
        "    n_gpus (int): number of gpus\n",
        "    rank (int): rank of current gpu\n",
        "    hparams (object): comma separated list of \"name=value\" pairs.\n",
        "    \"\"\"\n",
        "    if hparams.distributed_run:\n",
        "        init_distributed(hparams, n_gpus, rank, group_name)\n",
        "\n",
        "    torch.manual_seed(hparams.seed)\n",
        "    torch.cuda.manual_seed(hparams.seed)\n",
        "\n",
        "    model = load_model(hparams)\n",
        "    learning_rate = hparams.learning_rate\n",
        "    optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate,\n",
        "                                 weight_decay=hparams.weight_decay)\n",
        "\n",
        "    if hparams.fp16_run:\n",
        "        from apex import amp\n",
        "        model, optimizer = amp.initialize(\n",
        "            model, optimizer, opt_level='O2')\n",
        "\n",
        "    if hparams.distributed_run:\n",
        "        model = apply_gradient_allreduce(model)\n",
        "\n",
        "    criterion = Tacotron2Loss()\n",
        "\n",
        "    logger = prepare_directories_and_logger(\n",
        "        output_directory, log_directory, rank)\n",
        "\n",
        "    train_loader, valset, collate_fn = prepare_dataloaders(hparams)\n",
        "\n",
        "    # Load checkpoint if one exists\n",
        "    iteration = 0\n",
        "    epoch_offset = 0\n",
        "    if checkpoint_path is not None and os.path.isfile(checkpoint_path):\n",
        "        if warm_start:\n",
        "            model = warm_start_model(\n",
        "                checkpoint_path, model, hparams.ignore_layers)\n",
        "        else:\n",
        "            model, optimizer, _learning_rate, iteration = load_checkpoint(\n",
        "                checkpoint_path, model, optimizer)\n",
        "            if hparams.use_saved_learning_rate:\n",
        "                learning_rate = _learning_rate\n",
        "            iteration += 1  # next iteration is iteration + 1\n",
        "            epoch_offset = max(0, int(iteration / len(train_loader)))\n",
        "    else:\n",
        "      os.path.isfile(\"pretrained_model\")\n",
        "      download_from_google_drive(\"1c5ZTuT7J08wLUoVZ2KkUs_VdZuJ86ZqA\",\"pretrained_model\")\n",
        "      model = warm_start_model(\"pretrained_model\", model, hparams.ignore_layers)\n",
        "      # download LJSpeech pretrained model if no checkpoint already exists\n",
        "    \n",
        "    start_eposh = time.perf_counter()\n",
        "    learning_rate = 0.0\n",
        "    model.train()\n",
        "    is_overflow = False\n",
        "    # ================ MAIN TRAINNIG LOOP! ===================\n",
        "    for epoch in tqdm(range(epoch_offset, hparams.epochs)):\n",
        "        print(\"\\nStarting Epoch: {} Iteration: {}\".format(epoch, iteration))\n",
        "        start_eposh = time.perf_counter() # eposh is russian, not a typo\n",
        "        for i, batch in tqdm(enumerate(train_loader), total=len(train_loader)):\n",
        "            start = time.perf_counter()\n",
        "            if iteration < hparams.decay_start: learning_rate = hparams.A_\n",
        "            else: iteration_adjusted = iteration - hparams.decay_start; learning_rate = (hparams.A_*(e**(-iteration_adjusted/hparams.B_))) + hparams.C_\n",
        "            learning_rate = max(hparams.min_learning_rate, learning_rate) # output the largest number\n",
        "            for param_group in optimizer.param_groups:\n",
        "                param_group['lr'] = learning_rate\n",
        "\n",
        "            model.zero_grad()\n",
        "            x, y = model.parse_batch(batch)\n",
        "            y_pred = model(x)\n",
        "\n",
        "            loss = criterion(y_pred, y)\n",
        "            if hparams.distributed_run:\n",
        "                reduced_loss = reduce_tensor(loss.data, n_gpus).item()\n",
        "            else:\n",
        "                reduced_loss = loss.item()\n",
        "            if hparams.fp16_run:\n",
        "                with amp.scale_loss(loss, optimizer) as scaled_loss:\n",
        "                    scaled_loss.backward()\n",
        "            else:\n",
        "                loss.backward()\n",
        "\n",
        "            if hparams.fp16_run:\n",
        "                grad_norm = torch.nn.utils.clip_grad_norm_(\n",
        "                    amp.master_params(optimizer), hparams.grad_clip_thresh)\n",
        "                is_overflow = math.isnan(grad_norm)\n",
        "            else:\n",
        "                grad_norm = torch.nn.utils.clip_grad_norm_(\n",
        "                    model.parameters(), hparams.grad_clip_thresh)\n",
        "\n",
        "            optimizer.step()\n",
        "\n",
        "            if not is_overflow and rank == 0:\n",
        "                duration = time.perf_counter() - start\n",
        "                logger.log_training(\n",
        "                    reduced_loss, grad_norm, learning_rate, duration, iteration)\n",
        "                #print(\"Batch {} loss {:.6f} Grad Norm {:.6f} Time {:.6f}\".format(iteration, reduced_loss, grad_norm, duration), end='\\r', flush=True)\n",
        "\n",
        "            iteration += 1\n",
        "        validate(model, criterion, valset, iteration,\n",
        "                 hparams.batch_size, n_gpus, collate_fn, logger,\n",
        "                 hparams.distributed_run, rank, epoch, start_eposh, learning_rate)\n",
        "        save_checkpoint(model, optimizer, learning_rate, iteration, checkpoint_path)\n",
        "        if log_directory2 != None:\n",
        "            copy_tree(log_directory, log_directory2)\n",
        "def check_dataset(hparams):\n",
        "    from utils import load_wav_to_torch, load_filepaths_and_text\n",
        "    import os\n",
        "    import numpy as np\n",
        "    def check_arr(filelist_arr):\n",
        "        for i, file in enumerate(filelist_arr):\n",
        "            if len(file) > 2:\n",
        "                print(\"|\".join(file), \"\\nhas multiple '|', this may not be an error.\")\n",
        "            if hparams.load_mel_from_disk and '.wav' in file[0]:\n",
        "                print(\"[WARNING]\", file[0], \" in filelist while expecting .npy .\")\n",
        "            else:\n",
        "                if not hparams.load_mel_from_disk and '.npy' in file[0]:\n",
        "                    print(\"[WARNING]\", file[0], \" in filelist while expecting .wav .\")\n",
        "            if (not os.path.exists(file[0])):\n",
        "                print(\"|\".join(file), \"\\n[WARNING] does not exist.\")\n",
        "            if len(file[1]) < 3:\n",
        "                print(\"|\".join(file), \"\\n[info] has no/very little text.\")\n",
        "            if not ((file[1].strip())[-1] in r\"!?,.;:\"):\n",
        "                print(\"|\".join(file), \"\\n[info] has no ending punctuation.\")\n",
        "            mel_length = 1\n",
        "            if hparams.load_mel_from_disk and '.npy' in file[0]:\n",
        "                melspec = torch.from_numpy(np.load(file[0], allow_pickle=True))\n",
        "                mel_length = melspec.shape[1]\n",
        "            if mel_length == 0:\n",
        "                print(\"|\".join(file), \"\\n[WARNING] has 0 duration.\")\n",
        "    print(\"Checking Training Files\")\n",
        "    audiopaths_and_text = load_filepaths_and_text(hparams.training_files) # get split lines from training_files text file.\n",
        "    check_arr(audiopaths_and_text)\n",
        "    print(\"Checking Validation Files\")\n",
        "    audiopaths_and_text = load_filepaths_and_text(hparams.validation_files) # get split lines from validation_files text file.\n",
        "    check_arr(audiopaths_and_text)\n",
        "    print(\"Finished Checking\")\n",
        "\n",
        "warm_start=False#sorry bout that\n",
        "n_gpus=1\n",
        "rank=0\n",
        "group_name=None\n",
        "\n",
        "# ---- DEFAULT PARAMETERS DEFINED HERE ----\n",
        "hparams = create_hparams()\n",
        "model_filename = 'current_model'\n",
        "hparams.training_files = \"filelists/clipper_train_filelist.txt\"\n",
        "hparams.validation_files = \"filelists/clipper_val_filelist.txt\"\n",
        "#hparams.use_mmi=True,          # not used in this notebook\n",
        "#hparams.use_gaf=True,          # not used in this notebook\n",
        "#hparams.max_gaf=0.5,           # not used in this notebook\n",
        "#hparams.drop_frame_rate = 0.2  # not used in this notebook\n",
        "hparams.p_attention_dropout=0.1\n",
        "hparams.p_decoder_dropout=0.1\n",
        "hparams.decay_start = 15000\n",
        "hparams.A_ = 5e-4\n",
        "hparams.B_ = 8000\n",
        "hparams.C_ = 0\n",
        "hparams.min_learning_rate = 1e-5\n",
        "generate_mels = True\n",
        "hparams.show_alignments = True\n",
        "alignment_graph_height = 600\n",
        "alignment_graph_width = 1000\n",
        "hparams.batch_size = 32\n",
        "hparams.load_mel_from_disk = True\n",
        "hparams.ignore_layers = []\n",
        "hparams.epochs = 10000\n",
        "torch.backends.cudnn.enabled = hparams.cudnn_enabled\n",
        "torch.backends.cudnn.benchmark = hparams.cudnn_benchmark\n",
        "output_directory = '/content/drive/My Drive/colab/outdir' # Location to save Checkpoints\n",
        "log_directory = '/content/tacotron2/logs' # Location to save Log files locally\n",
        "log_directory2 = '/content/drive/My Drive/colab/logs' # Location to copy log files (done at the end of each epoch to cut down on I/O)\n",
        "checkpoint_path = output_directory+(r'/')+model_filename\n",
        "\n",
        "# ---- Replace .wav with .npy in filelists ----\n",
        "!sed -i -- 's,.wav|,.npy|,g' filelists/*.txt\n",
        "!sed -i -- 's,.wav|,.npy|,g' {hparams.training_files}\n",
        "!sed -i -- 's,.wav|,.npy|,g' {hparams.validation_files}\n",
        "# ---- Replace .wav with .npy in filelists ----"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "UP1trdpN_jV6",
        "cellView": "form"
      },
      "source": [
        "#@markdown ## Setting model parameters\n",
        "#@markdown ---\n",
        "#@markdown #### Your desired model name\n",
        "model_filename = 'current_model' #@param {type: \"string\"}\n",
        "\n",
        "#@markdown #### Upload your transcription / text to tacotron/filelists and right click -> copy path\n",
        "Training_file = \"filelists/list.txt\" #@param {type: \"string\"}\n",
        "hparams.training_files = Training_file\n",
        "hparams.validation_files = Training_file\n",
        "\n",
        "# hparams to Tune\n",
        "#hparams.use_mmi=True,          # not used in this notebook\n",
        "#hparams.use_gaf=True,          # not used in this notebook\n",
        "#hparams.max_gaf=0.5,           # not used in this notebook\n",
        "#hparams.drop_frame_rate = 0.2  # not used in this notebook\n",
        "hparams.p_attention_dropout=0.1\n",
        "hparams.p_decoder_dropout=0.1\n",
        "\n",
        "# Learning Rate             # https://www.desmos.com/calculator/ptgcz4vzsw / http://boards.4channel.org/mlp/thread/34778298#p34789030\n",
        "hparams.decay_start = 15000         # wait till decay_start to start decaying learning rate\n",
        "hparams.A_ = 5e-4                   # Start/Max Learning Rate\n",
        "hparams.B_ = 8000                   # Decay Rate\n",
        "hparams.C_ = 0                      # Shift learning rate equation by this value\n",
        "hparams.min_learning_rate = 1e-5    # Min Learning Rate\n",
        "\n",
        "# Quality of Life\n",
        "generate_mels = True\n",
        "hparams.show_alignments = True\n",
        "alignment_graph_height = 600\n",
        "alignment_graph_width = 1000\n",
        "\n",
        "#@markdown #### Your batch size, lower if you don't have enough ram.\n",
        "hparams.batch_size = 18 #@param {type: \"integer\"}\n",
        "hparams.load_mel_from_disk = True\n",
        "hparams.ignore_layers = [] # Layers to reset (None by default, other than foreign languages this param can be ignored)\n",
        "\n",
        "##@markdown #### Amount of epochs before stopping, preferably a very high amount to not stop.\n",
        "hparams.epochs = 10000\n",
        "\n",
        "torch.backends.cudnn.enabled = hparams.cudnn_enabled\n",
        "torch.backends.cudnn.benchmark = hparams.cudnn_benchmark\n",
        "\n",
        "#@markdown #### Where to save your model when training.\n",
        "output_directory = '/content/drive/My Drive/colab/outdir' #@param {type: \"string\"}\n",
        "log_directory = '/content/tacotron2/logs' # Location to save Log files locally\n",
        "log_directory2 = '/content/drive/My Drive/colab/logs' # Location to copy log files (done at the end of each epoch to cut down on I/O)\n",
        "checkpoint_path = output_directory+(r'/')+model_filename\n",
        "\n",
        "#@markdown ---"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "A2-xc6EcACWc",
        "cellView": "form"
      },
      "source": [
        "#@markdown ## Convert WAV files into Mel-Spectrograms (Run only once)\n",
        "\n",
        "if generate_mels:\n",
        "    create_mels()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "j_ATUV9uWHwx",
        "cellView": "form"
      },
      "source": [
        "#@markdown ## Check for missing files\n",
        "\n",
        "# ---- Replace .wav with .npy in filelists ----\n",
        "!sed -i -- 's,.wav|,.npy|,g' {hparams.training_files}; sed -i -- 's,.wav|,.npy|,g' {hparams.validation_files}\n",
        "\n",
        "check_dataset(hparams)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "iqIk1dIuCqhx",
        "cellView": "form"
      },
      "source": [
        "#@markdown # Begin training\n",
        "\n",
        "print('FP16 Run:', hparams.fp16_run)\n",
        "print('Dynamic Loss Scaling:', hparams.dynamic_loss_scaling)\n",
        "print('Distributed Run:', hparams.distributed_run)\n",
        "print('cuDNN Enabled:', hparams.cudnn_enabled)\n",
        "print('cuDNN Benchmark:', hparams.cudnn_benchmark)\n",
        "train(output_directory, log_directory, checkpoint_path,\n",
        "      warm_start, n_gpus, rank, group_name, hparams, log_directory2)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "h4ql3OWWPKZv"
      },
      "source": [
        "### **Validation Loss under 0.1 can generally be considered good enough to try out**\n",
        "\n",
        "JS to prevent idle timeout:\n",
        "\n",
        "Press F12 OR CTRL + SHIFT + I OR right click on this website -> inspect.\n",
        "Then click on the console tab and paste in the following code.\n",
        "\n",
        "```javascript\n",
        "function ClickConnect(){\n",
        "console.log(\"Working\");\n",
        "document.querySelector(\"colab-toolbar-button#connect\").click()\n",
        "}\n",
        "setInterval(ClickConnect,60000)\n",
        "```"
      ]
    }
  ]
}