TransformerTTS.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Esl31BIC1FWb"
      },
      "source": [
        "# Transformer TTS\n",
        "\n",
        "A implementation of [Transformer TTS](https://ojs.aaai.org/index.php/AAAI/article/view/4642) using PyTorch. A significant portion of the code was developed with reference to [Nvidia Tacotron2](https://github.com/NVIDIA/tacotron2) and [choiHkk Transformer TTS](https://github.com/choiHkk/Transformer-TTS).\n",
        "\n",
        "<br>\n",
        "\n",
        "**Paper reference**:  \n",
        "Li, N., Liu, S., Liu, Y., Zhao, S., & Liu, M. (2019, July). Neural speech synthesis with transformer network. In Proceedings of the AAAI conference on artificial intelligence (Vol. 33, No. 01, pp. 6706-6713)."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 86,
      "metadata": {
        "id": "mSVdfbqB065S"
      },
      "outputs": [],
      "source": [
        "# pacakages\n",
        "import os\n",
        "import math\n",
        "import librosa\n",
        "import numpy as np\n",
        "import matplotlib.pyplot as plt\n",
        "from tqdm import tqdm\n",
        "from scipy import signal\n",
        "\n",
        "# torch\n",
        "import torch\n",
        "import torch.nn as nn\n",
        "import torch.nn.functional as F\n",
        "\n",
        "# torchaudio\n",
        "import torchaudio\n",
        "import torchaudio.functional as af\n",
        "import torchaudio.transforms as T"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 87,
      "metadata": {
        "id": "zEqyad9AUrDX"
      },
      "outputs": [],
      "source": [
        "class Hparams:\n",
        "    # audio\n",
        "    preemphasis = 0.97\n",
        "    n_fft = 1024\n",
        "    hop_length = 256\n",
        "    win_length = 1024\n",
        "    n_mels = 80\n",
        "    f_min = 0.0\n",
        "    f_max = 8000.0\n",
        "    ref_level_db = 20.0\n",
        "    min_level_db = -100.0\n",
        "    max_abs_value = 4.0\n",
        "\n",
        "\n",
        "    # train\n",
        "    epochs = 10\n",
        "    batch_size = 16\n",
        "    vis_every = 50\n",
        "\n",
        "    # model\n",
        "    n_phonemes = 70\n",
        "    n_mels = 80\n",
        "    embedding_dim = 512\n",
        "    d_model = 256\n",
        "\n",
        "hp = Hparams()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 88,
      "metadata": {
        "id": "6aY1qUcfJ4ig"
      },
      "outputs": [],
      "source": [
        "# visualize mel spectrogram\n",
        "def visualize(melspec):\n",
        "    fig, ax = plt.subplots(1, 1, figsize=(12, 3))\n",
        "    im1 = ax.imshow(melspec, aspect=\"auto\", origin=\"lower\", interpolation='none')\n",
        "    plt.colorbar(im1, ax=ax)\n",
        "    plt.tight_layout()\n",
        "    plt.show()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "xa6xvxPSr6E7"
      },
      "source": [
        "## Dataset\n",
        "\n",
        "The dataset used in the experiment was LJSpeech, which can be easily downloaded from torchaudio. To convert text to phonemes and encode it as numbers, we used the text processor from `torchaudio.pipelines.TACOTRON2_WAVERNN_CHAR_LJSPEECH`. The preprocessing of wave files was referenced from Nvidia's Tacotron2 and Kyubong's Tacotron."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 89,
      "metadata": {
        "id": "Iy95XhPFrNA4"
      },
      "outputs": [],
      "source": [
        "class TransformerTTSDataset(torch.utils.data.Dataset):\n",
        "\n",
        "    \"\"\"\n",
        "    https://github.com/Kyubyong/tacotron/blob/master/utils.py#L21\n",
        "    \"\"\"\n",
        "\n",
        "    def __init__(self):\n",
        "        self.dataset = torchaudio.datasets.LJSPEECH('.', download=True)\n",
        "        self.text_preprocess = torchaudio.pipelines.TACOTRON2_WAVERNN_CHAR_LJSPEECH.get_text_processor()\n",
        "\n",
        "    def __len__(self):\n",
        "        return len(self.dataset)\n",
        "\n",
        "    def __getitem__(self, idx):\n",
        "\n",
        "        # get data\n",
        "        y, sr, _, text = self.dataset[idx]\n",
        "\n",
        "        # text preprocess\n",
        "        text, _ = self.text_preprocess(text)\n",
        "\n",
        "        # audio preprocess\n",
        "        mel = self.audio_preprocess(y, sr)\n",
        "\n",
        "        # (text_len), (n_mel, mel_len)\n",
        "        return text.squeeze(), mel\n",
        "\n",
        "    def audio_preprocess(self, y, sr):\n",
        "\n",
        "        # trimming\n",
        "        y, _ = librosa.effects.trim(y)\n",
        "\n",
        "        # Preemphasis\n",
        "        y = af.preemphasis(y, hp.preemphasis)\n",
        "\n",
        "        # stft\n",
        "        linear = librosa.stft(y=y.numpy().squeeze(),\n",
        "                            n_fft=hp.n_fft,\n",
        "                            hop_length=hp.hop_length,\n",
        "                            win_length=hp.win_length,\n",
        "                            pad_mode='constant')\n",
        "\n",
        "        # magnitude spectrogram\n",
        "        mag = np.abs(linear)\n",
        "\n",
        "        # mel spectrogram\n",
        "        mel_basis = librosa.filters.mel(sr=sr, n_fft=hp.n_fft, n_mels=hp.n_mels, fmin=hp.f_min, fmax=hp.f_max)\n",
        "        mel = np.dot(mel_basis, mag)\n",
        "\n",
        "        # to decibel\n",
        "        mel = 20 * np.log10(np.maximum(1e-5, mel))\n",
        "\n",
        "        mel = mel - hp.ref_level_db\n",
        "\n",
        "        # normalize\n",
        "        max_abs_value = hp.max_abs_value\n",
        "        min_level_db = hp.min_level_db\n",
        "        mel = np.clip((2 * max_abs_value) * ((mel - min_level_db) / (-min_level_db)) - max_abs_value,\n",
        "                      -max_abs_value, max_abs_value)\n",
        "\n",
        "        return torch.FloatTensor(mel)\n",
        "\n",
        "\n",
        "class TransformerTTSCollate():\n",
        "\n",
        "    def __init__(self):\n",
        "        ...\n",
        "\n",
        "    def __call__(self, batch):\n",
        "\n",
        "        # get decreasing order by text length within batch\n",
        "        text_lengths, ids_sorted_decreasing = torch.sort(\n",
        "            torch.LongTensor([len(text) for text, _ in batch]),\n",
        "            dim=0, descending=True\n",
        "        )\n",
        "\n",
        "        # all zero padded tensor\n",
        "        max_text_len = text_lengths[0]\n",
        "        text_padded = torch.LongTensor(len(batch), max_text_len)\n",
        "        text_padded.zero_()\n",
        "\n",
        "        # allocate text to zero padded tensor\n",
        "        for i in range(len(ids_sorted_decreasing)):\n",
        "            text = batch[ids_sorted_decreasing[i]][0]\n",
        "            text_padded[i, :text.size(0)] = text\n",
        "\n",
        "        # get maximum length of mel sequence within batch\n",
        "        num_mels = batch[0][1].size(0)\n",
        "        max_mel_len = max([mel.size(1) for _, mel in batch])\n",
        "\n",
        "        # all zero padded tensor\n",
        "        mel_padded = torch.FloatTensor(len(batch), num_mels, max_mel_len)\n",
        "        mel_padded.zero_()\n",
        "        gate_padded = torch.FloatTensor(len(batch), max_mel_len)\n",
        "        gate_padded.zero_()\n",
        "        mel_lengths = torch.LongTensor(len(batch))\n",
        "\n",
        "        for i in range(len(ids_sorted_decreasing)):\n",
        "            mel = batch[ids_sorted_decreasing[i]][1]\n",
        "            mel_padded[i, :, :mel.size(1)] = mel\n",
        "            gate_padded[i, mel.size(1) - 1:] = 1\n",
        "            mel_lengths[i] = mel.size(1)\n",
        "\n",
        "        return (\n",
        "            text_padded,\n",
        "            text_lengths,\n",
        "            mel_padded.transpose(1, 2),\n",
        "            gate_padded,\n",
        "            mel_lengths\n",
        "        )"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 90,
      "metadata": {
        "id": "tbsdkFdvZD5U"
      },
      "outputs": [],
      "source": [
        "dataset = TransformerTTSDataset()\n",
        "dataloader = torch.utils.data.DataLoader(dataset, batch_size=hp.batch_size, shuffle=True, collate_fn=TransformerTTSCollate())"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "nYk7UQDkr9t5"
      },
      "source": [
        "## Model"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "WOuJpFb3sBzI"
      },
      "source": [
        "### Scaled Positional Encoding\n",
        "\n",
        "**Page 3**\n",
        "\n",
        "$$ PE(pos, 2i) = \\sin (\\frac{pos}{10000^{\\frac{2i}{d_{model}}}}) \\tag{6} $$\n",
        "\n",
        "$$ PE(pos, 2i + 1) = \\cos (\\frac{pos}{10000^{\\frac{2i}{d_{model}}}}) \\tag{7} $$\n",
        "\n",
        "<br>\n",
        "\n",
        "In NMT, the embeddings for both source and target language are from language spaces, so the scales of these embeddings are similar. This condition doesn’t hold in the TTS scenarioe, since the source domain is of texts while the target domain is of mel spectrograms, hence using fixed positional embeddings may impose heavy constraints on both the encoder and decoder pre-nets (which will be described in Sec. 3.3 and 3.4)."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 91,
      "metadata": {
        "id": "9OVZpT-1rRsT"
      },
      "outputs": [],
      "source": [
        "class ScaledPositionalEncoding(nn.Module):\n",
        "\n",
        "    \"\"\"\n",
        "    https://pytorch.org/tutorials/beginner/transformer_tutorial\n",
        "    \"\"\"\n",
        "\n",
        "    def __init__(self, d_model, max_len=5000):\n",
        "\n",
        "        super(ScaledPositionalEncoding, self).__init__()\n",
        "\n",
        "        # terms for positional embedding\n",
        "        position = torch.arange(0, max_len).unsqueeze(1).float()\n",
        "        div_term = torch.exp(torch.arange(0, d_model, 2).float() * -(np.log(10000.0) / d_model))\n",
        "\n",
        "        # fixed positional embedding\n",
        "        pe = torch.zeros(1, max_len, d_model)\n",
        "        pe[0, :, 0::2] = torch.sin(position * div_term)\n",
        "        pe[0, :, 1::2] = torch.cos(position * div_term)\n",
        "        self.register_buffer('pe', pe)\n",
        "\n",
        "        # trainable weight\n",
        "        self.alpha = nn.Parameter(torch.ones(1))\n",
        "\n",
        "\n",
        "    def forward(self, x):\n",
        "        return x + self.alpha * self.pe[:, :x.size(1)]"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "i13aAJnnr_6b"
      },
      "source": [
        "### Encoder Prenet\n",
        "\n",
        "**page 3**\n",
        "\n",
        "The output of each convolution layer has 512 channels, followed by a batch normalization and ReLU activation, and a dropout layer as well. In addition, we add a linear projection after the final ReLU activation, since the output range of ReLU is $[0, +∞)$, while each dimension of these triangle positional embeddings is in $[−1, 1]$."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 92,
      "metadata": {
        "id": "VX0qDfz5rQOO"
      },
      "outputs": [],
      "source": [
        "class EncoderPrenet(nn.Module):\n",
        "\n",
        "    \"\"\"\n",
        "    https://github.com/NVIDIA/tacotron2/blob/master/model.py#L89\n",
        "    \"\"\"\n",
        "\n",
        "    def __init__(self, in_dim=512, sizes=[512, 512, 512], encoder_kernel_size=3):\n",
        "        super(EncoderPrenet, self).__init__()\n",
        "\n",
        "        # convolution layer\n",
        "        in_sizes = [in_dim] + sizes[:-1]\n",
        "        self.layers = nn.ModuleList(\n",
        "            [nn.Sequential(\n",
        "                ConvNorm(in_size * 2, out_size * 2, kernel_size=encoder_kernel_size,\n",
        "                      padding=int((encoder_kernel_size - 1) / 2),\n",
        "                      dilation=1, w_init_gain='relu'),\n",
        "                nn.BatchNorm1d(out_size * 2)\n",
        "            ) for (in_size, out_size) in zip(in_sizes, sizes)]\n",
        "        )\n",
        "        self.projection = LinearNorm(sizes[-1] * 2, sizes[-1])\n",
        "\n",
        "\n",
        "    def forward(self, x):\n",
        "        x = x.transpose(1, 2)\n",
        "        for linear in self.layers:\n",
        "            x = F.dropout(F.relu(linear(x)), p=0.5, training=True)\n",
        "\n",
        "        x = self.projection(x.transpose(1, 2))\n",
        "        return x\n",
        "\n",
        "\n",
        "class LinearNorm(nn.Module):\n",
        "\n",
        "    \"\"\"\n",
        "    https://github.com/NVIDIA/tacotron2/blob/master/layers.py#L8\n",
        "    \"\"\"\n",
        "\n",
        "    def __init__(self, in_dim, out_dim, bias=True, w_init_gain='linear'):\n",
        "        super(LinearNorm, self).__init__()\n",
        "        self.linear_layer = nn.Linear(in_dim, out_dim, bias=bias)\n",
        "        nn.init.xavier_uniform_(\n",
        "            self.linear_layer.weight,\n",
        "            gain=nn.init.calculate_gain(w_init_gain)\n",
        "        )\n",
        "\n",
        "    def forward(self, x):\n",
        "        return self.linear_layer(x)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "zE7yvq8crXfp"
      },
      "source": [
        "### Decoder Prenet\n",
        "\n",
        "**page 3-4**\n",
        "\n",
        "The mel spectrogram is first consumed by a neural network composed of two fully connected layers(each has 256 hidden units) with ReLU activation, named \"decoder pre-net\", and it plays an important role in the TTS system."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 93,
      "metadata": {
        "id": "YIBuDnT7ro0A"
      },
      "outputs": [],
      "source": [
        "class DecoderPrenet(nn.Module):\n",
        "\n",
        "    \"\"\"\n",
        "    https://github.com/NVIDIA/tacotron2/blob/master/model.py#L89\n",
        "    \"\"\"\n",
        "\n",
        "    def __init__(self, in_dim=80, sizes=[256, 256]):\n",
        "        super(DecoderPrenet, self).__init__()\n",
        "\n",
        "        in_sizes = [in_dim] + sizes[:-1]\n",
        "        self.layers = nn.ModuleList(\n",
        "            [LinearNorm(in_size, out_size, bias=False)\n",
        "            for (in_size, out_size) in zip(in_sizes, sizes)]\n",
        "        )\n",
        "\n",
        "    def forward(self, x):\n",
        "        for linear in self.layers:\n",
        "            x = F.dropout(F.relu(linear(x)), p=0.5, training=True)\n",
        "\n",
        "        return x"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "mVJu9qE2sD0o"
      },
      "source": [
        "### Postnet\n",
        "\n",
        "**[Tacotron 2] page 2. Spectrogram Prediction Network**\n",
        "\n",
        "he predicted mel spectrogram is passed through a $5$-layer convolutional post-net which predicts a residual to add to the prediction to improve the overall reconstruction. Each post-net layer is  comprised of $512$ filters with shape $5 \\times 1$ with batch\n",
        "normalization, followed by tanh activations on all but the final layer.\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 94,
      "metadata": {
        "id": "u1_k6d47rTYl"
      },
      "outputs": [],
      "source": [
        "class PostNet(nn.Module):\n",
        "\n",
        "    \"\"\"\n",
        "    https://github.com/NVIDIA/tacotron2/blob/master/model.py#L103\n",
        "    \"\"\"\n",
        "\n",
        "    def __init__(self, n_mel_channels=80, postnet_embedding_dim=1024,\n",
        "                 postnet_kernel_size=5, postnet_n_convolutions=5):\n",
        "        super(PostNet, self).__init__()\n",
        "\n",
        "        in_channels = [n_mel_channels] + [postnet_embedding_dim] * (postnet_n_convolutions - 1)\n",
        "        out_channels = [postnet_embedding_dim] * (postnet_n_convolutions - 1) + [n_mel_channels]\n",
        "\n",
        "        # convolution layers\n",
        "        self.convolutions = nn.ModuleList(\n",
        "            [nn.Sequential(\n",
        "                ConvNorm(in_channel, out_channel,\n",
        "                         kernel_size=postnet_kernel_size, stride=1,\n",
        "                         padding=int((postnet_kernel_size - 1) / 2),\n",
        "                         dilation=1, w_init_gain='tanh'),\n",
        "                nn.BatchNorm1d(out_channel)\n",
        "            ) for (in_channel, out_channel) in zip(in_channels, out_channels)]\n",
        "        )\n",
        "\n",
        "\n",
        "    def forward(self, x):\n",
        "        for i in range(len(self.convolutions) - 1):\n",
        "            x = F.dropout(F.tanh(self.convolutions[i](x)), 0.5, self.training)\n",
        "        return F.dropout(self.convolutions[-1](x), 0.5, self.training)\n",
        "\n",
        "\n",
        "class ConvNorm(nn.Module):\n",
        "\n",
        "    \"\"\"\n",
        "    https://github.com/NVIDIA/tacotron2/blob/master/layers.py#L21\n",
        "    \"\"\"\n",
        "\n",
        "    def __init__(self, in_channels, out_channels, kernel_size=1, stride=1,\n",
        "                 padding=None, dilation=1, bias=True, w_init_gain='linear'):\n",
        "        super(ConvNorm, self).__init__()\n",
        "\n",
        "        # check padding is valid\n",
        "        if padding is None:\n",
        "            assert(kernel_size % 2 == 1)\n",
        "            padding = int(dilation * (kernel_size - 1) / 2)\n",
        "\n",
        "        # 1d convolution\n",
        "        self.conv = nn.Conv1d(in_channels, out_channels,\n",
        "                              kernel_size=kernel_size, stride=stride,\n",
        "                              padding=padding, dilation=dilation,\n",
        "                              bias=bias)\n",
        "\n",
        "        # weight normalization\n",
        "        nn.init.xavier_uniform_(\n",
        "            self.conv.weight,\n",
        "            gain=nn.init.calculate_gain(w_init_gain)\n",
        "        )\n",
        "\n",
        "\n",
        "    def forward(self, x):\n",
        "        return self.conv(x)\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "-hP4NY5_sFIb"
      },
      "source": [
        "### Transformer TTS"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 95,
      "metadata": {
        "id": "sARM3DawrU05"
      },
      "outputs": [],
      "source": [
        "class TransformerTTS(nn.Module):\n",
        "\n",
        "    \"\"\"\n",
        "    https://github.com/choiHkk/Transformer-TTS/blob/main/model.py#L20\n",
        "    \"\"\"\n",
        "\n",
        "    def __init__(self, hp):\n",
        "        super(TransformerTTS, self).__init__()\n",
        "\n",
        "        n_phonemes = hp.n_phonemes\n",
        "        n_mels = hp.n_mels\n",
        "        d_model = hp.d_model\n",
        "\n",
        "        # phoneme embedding\n",
        "        self.n_mels = hp.n_mels\n",
        "        self.phoneme_embedding = nn.Embedding(n_phonemes, d_model * 2)\n",
        "        std = math.sqrt(2.0 / (n_phonemes + d_model))\n",
        "        val = math.sqrt(3.0) * std\n",
        "        self.phoneme_embedding.weight.data.uniform_(-val, val)\n",
        "\n",
        "        # encoder pre-net\n",
        "        self.encoder_prenet = EncoderPrenet(d_model, [d_model] * 3, 3)\n",
        "\n",
        "        # decoder pre-net\n",
        "        self.decoder_prenet = DecoderPrenet(n_mels, [d_model, d_model])\n",
        "\n",
        "        # scaled positional encoding\n",
        "        self.scaled_positional_encoding = ScaledPositionalEncoding(d_model)\n",
        "\n",
        "        # transformer encoder layer\n",
        "        transformer_encoder_layer = nn.TransformerEncoderLayer(\n",
        "            d_model=d_model, nhead=4, dim_feedforward=1024, dropout=0.1,\n",
        "            activation=\"relu\", batch_first=True, norm_first=False, layer_norm_eps=1e-5)\n",
        "\n",
        "        # transformer decoder layer\n",
        "        transformer_decoder_layer = nn.TransformerDecoderLayer(\n",
        "            d_model=d_model, nhead=4, dim_feedforward=1024, dropout=0.1,\n",
        "            activation=\"relu\", batch_first=True, norm_first=False, layer_norm_eps=1e-5)\n",
        "\n",
        "        # transformer encoder\n",
        "        self.transformer_encoder = nn.TransformerEncoder(\n",
        "            encoder_layer=transformer_encoder_layer, num_layers=3, norm=None)\n",
        "\n",
        "        # transformer decoer\n",
        "        self.transformer_decoder = nn.TransformerDecoder(\n",
        "            decoder_layer=transformer_decoder_layer, num_layers=3, norm=None)\n",
        "\n",
        "        # mel linear projection\n",
        "        self.mel_linear_projection = LinearNorm(d_model, n_mels)\n",
        "\n",
        "        # stop linear projection\n",
        "        self.stop_linear_projection = LinearNorm(d_model, 1, w_init_gain='sigmoid')\n",
        "\n",
        "        # postnet\n",
        "        self.postnet = PostNet(n_mels, 1024, 5, 5)\n",
        "\n",
        "\n",
        "    def forward(self, text, text_lengths, mel, mel_lengths):\n",
        "\n",
        "        self.train()\n",
        "        self.initialize_masks(text_lengths, mel_lengths)\n",
        "\n",
        "        # phoneme embedding + encoder prenet\n",
        "        x = self.encoder_prenet(self.phoneme_embedding(text))\n",
        "\n",
        "        # positional encoding\n",
        "        x = self.scaled_positional_encoding(x)\n",
        "\n",
        "        # transformer encoder\n",
        "        memory = self.transformer_encoder(\n",
        "            src=x,\n",
        "            mask=self.src_mask,\n",
        "            src_key_padding_mask=self.src_key_padding_mask)\n",
        "\n",
        "        y = torch.cat([self.get_go_frame(memory).unsqueeze(1), mel[:,:-1,:]], dim=1)\n",
        "        y = self.scaled_positional_encoding(self.decoder_prenet(y))\n",
        "\n",
        "        # transformer decoder\n",
        "        features = self.transformer_decoder(\n",
        "            tgt=y,\n",
        "            memory=memory,\n",
        "            tgt_mask=self.tgt_mask,\n",
        "            memory_mask=self.memory_mask,\n",
        "            tgt_key_padding_mask=self.tgt_key_padding_mask,\n",
        "            memory_key_padding_mask=self.src_key_padding_mask\n",
        "        )\n",
        "\n",
        "        # stop linear, mel linear, postnet\n",
        "        mel = self.mel_linear_projection(features)\n",
        "        stop = self.stop_linear_projection(features)\n",
        "        post_mel = self.postnet(mel.transpose(1, 2)).transpose(1, 2) + mel\n",
        "\n",
        "        # return padding masked output\n",
        "        return self.parse_output(post_mel, mel, stop)\n",
        "\n",
        "\n",
        "    def get_go_frame(self, memory):\n",
        "\n",
        "        # transformer decoder initial input\n",
        "        batch = memory.size(0)\n",
        "        decoder_input = torch.autograd.Variable(memory.data.new(batch, self.n_mels).zero_())\n",
        "\n",
        "        return decoder_input\n",
        "\n",
        "\n",
        "    def initialize_masks(self, text_lengths, mel_lengths):\n",
        "\n",
        "        # transformer encoder layer mask\n",
        "        S = text_lengths.max().item()\n",
        "        self.src_mask = self.generate_square_subsequent_mask(S, S).to(\n",
        "            device=text_lengths.device)\n",
        "        self.src_key_padding_mask = self.generate_padding_mask(text_lengths).to(\n",
        "            device=text_lengths.device)\n",
        "\n",
        "        # transformer decoder layer mask\n",
        "        T = mel_lengths.max().item()\n",
        "        self.tgt_mask = self.generate_square_subsequent_mask(T, T).to(\n",
        "            device=mel_lengths.device)\n",
        "        self.tgt_key_padding_mask = self.generate_padding_mask(mel_lengths).to(\n",
        "            device=mel_lengths.device)\n",
        "\n",
        "        # transformer decoder input mask\n",
        "        self.memory_mask = self.generate_square_subsequent_mask(T, S).to(\n",
        "            device=mel_lengths.device)\n",
        "\n",
        "\n",
        "    def parse_output(self, post_mel, mel, stop):\n",
        "        post_mel.data.masked_fill_(\n",
        "            self.tgt_key_padding_mask.unsqueeze(-1).repeat(1, 1, post_mel.size(-1)), 0.0)\n",
        "        mel.data.masked_fill_(\n",
        "            self.tgt_key_padding_mask.unsqueeze(-1).repeat(1, 1, mel.size(-1)), 0.0)\n",
        "        stop.data.masked_fill_(\n",
        "            self.tgt_key_padding_mask.unsqueeze(-1), 1e3)\n",
        "\n",
        "        return post_mel, mel, stop\n",
        "\n",
        "\n",
        "    def generate_square_subsequent_mask(self, lsz, rsz):\n",
        "\n",
        "        \"\"\"\n",
        "        https://pytorch.org/docs/stable/_modules/torch/nn/modules/transformer.html#Transformer\n",
        "        \"\"\"\n",
        "\n",
        "        return torch.triu(torch.ones(lsz, rsz) * float('-inf'), diagonal=1)\n",
        "\n",
        "\n",
        "    def generate_padding_mask(self, lengths, max_len=None):\n",
        "\n",
        "        \"\"\"\n",
        "        https://github.com/ming024/FastSpeech2/blob/master/utils/tools.py\n",
        "        \"\"\"\n",
        "\n",
        "        batch_size = lengths.size(0)\n",
        "        if max_len is None:\n",
        "            max_len = lengths.max().item()\n",
        "        ids = torch.arange(0, max_len).unsqueeze(0).expand(batch_size, -1).to(\n",
        "            dtype=lengths.dtype, device=lengths.device)\n",
        "        return ids >= lengths.unsqueeze(1).expand(-1, max_len)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "o2NT3VP9sH8G"
      },
      "source": [
        "## Loss function\n",
        "\n",
        "**Page 4. Mel Linear, Stop Linear and Post-net**\n",
        "\n",
        "It's worth mentioning that, for the stop linear, there is only one positive sample in the end of each sequence which means ”stop”, while hundreds of negative samples for other frames. This imbalance may result in unstoppable inference. We impose a positive weight ($5.0 ~ 8.0$) on the tail positive stop token when calculating binary cross entropy loss, and this problem was efficiently solved.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 96,
      "metadata": {
        "id": "YBBYbrQVrWCH"
      },
      "outputs": [],
      "source": [
        "class TransformerTTSLoss():\n",
        "\n",
        "    \"\"\"\n",
        "    https://github.com/choiHkk/Transformer-TTS/blob/main/loss_function.py\n",
        "    \"\"\"\n",
        "\n",
        "    def __init__(self):\n",
        "        self.mel_loss = nn.MSELoss()\n",
        "        self.stop_loss = nn.BCEWithLogitsLoss()\n",
        "        self.alpha = 5.0\n",
        "\n",
        "    def __call__(self, post_mel, mel_pred, mel_target, stop_pred, stop_target):\n",
        "\n",
        "        # match stop dimension\n",
        "        stop_pred = stop_pred.view(-1, 1)\n",
        "        stop_target = stop_target.view(-1, 1)\n",
        "\n",
        "        # calculate loss\n",
        "        mel_loss = self.mel_loss(mel_pred, mel_target) + self.mel_loss(post_mel, mel_target)\n",
        "        stop_loss = self.stop_loss(stop_pred, stop_target)\n",
        "        return mel_loss + stop_loss * self.alpha"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "d7CpM3ffsJnD"
      },
      "source": [
        "## Train"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 97,
      "metadata": {
        "id": "vBxbEF1HN54O"
      },
      "outputs": [],
      "source": [
        "def visualize_training_mel(exp_dir, mel_pred, mel, epoch, step):\n",
        "\n",
        "    fig, (ax1, ax2) = plt.subplots(2, 1, figsize=(12, 6))\n",
        "\n",
        "    # predicted mel spectrogram\n",
        "    im1 = ax1.imshow(mel_pred, aspect=\"auto\", origin=\"lower\", interpolation='none')\n",
        "    plt.colorbar(im1, ax=ax1)\n",
        "    ax1.set_title(\"Prediction\")\n",
        "    ax1.set_xlabel(\"Frames\")\n",
        "    ax1.set_ylabel(\"Channels\")\n",
        "\n",
        "    # target mel spectrogram\n",
        "    im2 = ax2.imshow(mel, aspect=\"auto\", origin=\"lower\", interpolation='none')\n",
        "    plt.colorbar(im2, ax=ax2)\n",
        "    ax2.set_title(\"Target\")\n",
        "    ax2.set_xlabel(\"Frames\")\n",
        "    ax2.set_ylabel(\"Channels\")\n",
        "\n",
        "    plt.tight_layout()\n",
        "    exp_dir = os.path.join(exp_dir, f\"epoch-{epoch}-step-{step}.png\")\n",
        "    plt.savefig(exp_dir)\n",
        "    plt.close()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 98,
      "metadata": {
        "id": "A6F61KG5YEJK"
      },
      "outputs": [],
      "source": [
        "def get_experiment_dir(base_dir: str=\"./exp\"):\n",
        "\n",
        "    # make base directory\n",
        "    base_name = \"experiment\"\n",
        "    if not os.path.exists(base_dir):\n",
        "        os.makedirs(base_dir)\n",
        "\n",
        "    # experiment dir lists\n",
        "    existing_dirs = [d for d in os.listdir(base_dir) if os.path.isdir(os.path.join(base_dir, d))]\n",
        "    experiment_dirs = [d for d in existing_dirs if d.startswith(base_name)]\n",
        "\n",
        "    # get maximum experiment number\n",
        "    max_num = 0\n",
        "    for dir_name in experiment_dirs:\n",
        "        try:\n",
        "            num = int(dir_name[len(base_name):])\n",
        "            if num > max_num:\n",
        "                max_num = num\n",
        "        except ValueError:\n",
        "            continue\n",
        "\n",
        "    # new directory name / path\n",
        "    new_dir_name = f\"{base_name}{max_num + 1}\"\n",
        "    new_dir_path = os.path.join(base_dir, new_dir_name)\n",
        "    os.makedirs(new_dir_path)\n",
        "\n",
        "    return new_dir_path"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 99,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "TyidJRMWrXOp",
        "outputId": "e082d89f-c559-46f4-fe6b-a7d906da4f94"
      },
      "outputs": [
        {
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "100%|██████████| 819/819 [16:16<00:00,  1.19s/it, loss=0.562]\n",
            "100%|██████████| 819/819 [15:57<00:00,  1.17s/it, loss=0.469]\n",
            "100%|██████████| 819/819 [15:16<00:00,  1.12s/it, loss=0.465]\n",
            "100%|██████████| 819/819 [15:17<00:00,  1.12s/it, loss=0.309]\n",
            "100%|██████████| 819/819 [15:38<00:00,  1.15s/it, loss=0.385]\n",
            "100%|██████████| 819/819 [15:23<00:00,  1.13s/it, loss=0.379]\n",
            "100%|██████████| 819/819 [15:11<00:00,  1.11s/it, loss=0.33]\n",
            "100%|██████████| 819/819 [14:59<00:00,  1.10s/it, loss=0.292]\n",
            "100%|██████████| 819/819 [15:02<00:00,  1.10s/it, loss=0.271]\n",
            "100%|██████████| 819/819 [15:00<00:00,  1.10s/it, loss=0.306]\n"
          ]
        }
      ],
      "source": [
        "device = 'cuda' if torch.cuda.is_available() else 'cpu'\n",
        "epochs = hp.epochs\n",
        "\n",
        "# model\n",
        "model = TransformerTTS(hp).to(device)\n",
        "\n",
        "# loss\n",
        "criterion = TransformerTTSLoss()\n",
        "\n",
        "# optimizer\n",
        "optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)\n",
        "\n",
        "# make experiment directory\n",
        "os.makedirs(\"./exp\", exist_ok=True)\n",
        "experiment_dir = get_experiment_dir(\"./exp\")\n",
        "\n",
        "\n",
        "for epoch in range(epochs):\n",
        "\n",
        "    step = 0\n",
        "    train_tqdm_bar = tqdm(enumerate(dataloader), total=len(dataloader))\n",
        "\n",
        "    for i, data in train_tqdm_bar:\n",
        "\n",
        "        # get data\n",
        "        text, text_len, mel, stop, mel_len = data\n",
        "\n",
        "        # move to appropriate device\n",
        "        text = text.to(device)\n",
        "        text_len = text_len.to(device)\n",
        "        mel = mel.to(device)\n",
        "        stop = stop.to(device)\n",
        "        mel_len = mel_len.to(device)\n",
        "\n",
        "        # model prediction + loss calculation\n",
        "        post_mel, mel_pred, stop_pred = model(text, text_len, mel, mel_len)\n",
        "        loss = criterion(post_mel, mel_pred, mel, stop_pred, stop)\n",
        "        train_tqdm_bar.set_postfix(loss=loss.item())\n",
        "\n",
        "        # back propgation\n",
        "        optimizer.zero_grad()\n",
        "        loss.backward()\n",
        "        optimizer.step()\n",
        "\n",
        "        # predicted mel, target mel visualization\n",
        "        if i % hp.vis_every == 0:\n",
        "\n",
        "            # first data\n",
        "            _mel_pred = mel_pred[0].transpose(1, 0)\n",
        "            _mel_targ = mel[0].transpose(1, 0)\n",
        "\n",
        "            # detach from gpu\n",
        "            _mel_pred = _mel_pred.detach().cpu().numpy()\n",
        "            _mel_targ = _mel_targ.detach().cpu().numpy()\n",
        "\n",
        "            # save fig\n",
        "            visualize_training_mel(experiment_dir, _mel_pred, _mel_targ,\n",
        "                epoch, step)\n",
        "\n",
        "        # increment step\n",
        "        step = step + 1\n",
        "\n"
      ]
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "gpuType": "T4",
      "provenance": []
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}