text-to-speech/4.seq2seq-bahdanau.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import tensorflow as tf\n",
    "import os\n",
    "from utils import *\n",
    "from tqdm import tqdm"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "def prenet(inputs, num_units = None, is_training = True, scope = 'prenet'):\n",
    "    if num_units is None:\n",
    "        num_units = [embed_size, embed_size // 2]\n",
    "    with tf.variable_scope(scope):\n",
    "        outputs = tf.layers.dense(\n",
    "            inputs,\n",
    "            units = num_units[0],\n",
    "            activation = tf.nn.relu,\n",
    "            name = 'dense1',\n",
    "        )\n",
    "        outputs = tf.layers.dropout(\n",
    "            outputs,\n",
    "            rate = dropout_rate,\n",
    "            training = is_training,\n",
    "            name = 'dropout1',\n",
    "        )\n",
    "        outputs = tf.layers.dense(\n",
    "            outputs,\n",
    "            units = num_units[1],\n",
    "            activation = tf.nn.relu,\n",
    "            name = 'dense2',\n",
    "        )\n",
    "        outputs = tf.layers.dropout(\n",
    "            outputs,\n",
    "            rate = dropout_rate,\n",
    "            training = is_training,\n",
    "            name = 'dropout2',\n",
    "        )\n",
    "    return outputs\n",
    "\n",
    "\n",
    "def highwaynet(inputs, num_units = None, scope = 'highwaynet'):\n",
    "    if not num_units:\n",
    "        num_units = inputs.get_shape()[-1]\n",
    "    with tf.variable_scope(scope):\n",
    "        H = tf.layers.dense(\n",
    "            inputs, units = num_units, activation = tf.nn.relu, name = 'dense1'\n",
    "        )\n",
    "        T = tf.layers.dense(\n",
    "            inputs,\n",
    "            units = num_units,\n",
    "            activation = tf.nn.sigmoid,\n",
    "            bias_initializer = tf.constant_initializer(-1.0),\n",
    "            name = 'dense2',\n",
    "        )\n",
    "        outputs = H * T + inputs * (1.0 - T)\n",
    "    return outputs\n",
    "\n",
    "\n",
    "def conv1d_banks(inputs, K = 16, is_training = True, scope = 'conv1d_banks'):\n",
    "    with tf.variable_scope(scope):\n",
    "        outputs = tf.layers.conv1d(inputs, embed_size // 2, 1, padding = 'SAME')\n",
    "        for k in range(2, K + 1):\n",
    "            with tf.variable_scope('num_{}'.format(k)):\n",
    "                output = tf.layers.conv1d(\n",
    "                    inputs, embed_size // 2, k, padding = 'SAME'\n",
    "                )\n",
    "                outputs = tf.concat((outputs, output), -1)\n",
    "        outputs = tf.nn.relu(\n",
    "            tf.layers.batch_normalization(outputs, training = is_training)\n",
    "        )\n",
    "    return outputs\n",
    "\n",
    "class Model:\n",
    "    def __init__(\n",
    "        self,\n",
    "        num_layers,\n",
    "        size_layers,\n",
    "        learning_rate = 1e-3,\n",
    "        dropout = 1.0,\n",
    "    ):\n",
    "        self.X = tf.placeholder(tf.int32, (None, None))\n",
    "        self.training = tf.placeholder(tf.bool, None)\n",
    "        lookup_table = tf.get_variable(\n",
    "            'lookup_table',\n",
    "            dtype = tf.float32,\n",
    "            shape = [len(vocab), size_layers],\n",
    "            initializer = tf.truncated_normal_initializer(\n",
    "                mean = 0.0, stddev = 0.01\n",
    "            ),\n",
    "        )\n",
    "        lookup_table = tf.concat(\n",
    "            (tf.zeros(shape = [1, size_layers]), lookup_table[1:, :]), 0\n",
    "        )\n",
    "        forward = tf.nn.embedding_lookup(lookup_table, self.X)\n",
    "        self.Y = tf.placeholder(tf.float32, (None, None, n_mels * resampled))\n",
    "        self.decoder_inputs = tf.concat(\n",
    "            (tf.zeros_like(self.Y[:, :1, :]), self.Y[:, :-1, :]), 1\n",
    "        )\n",
    "        self.decoder_inputs = self.decoder_inputs[:, :, -n_mels:]\n",
    "        self.Z = tf.placeholder(\n",
    "            tf.float32, (None, None, fourier_window_size // 2 + 1)\n",
    "        )\n",
    "        \n",
    "        batch_size = tf.shape(self.X)[0]\n",
    "        seq_lens = tf.count_nonzero(\n",
    "            tf.reduce_sum(self.decoder_inputs, -1), 1, dtype = tf.int32\n",
    "        ) + 1\n",
    "        \n",
    "        def cells(reuse = False):\n",
    "            return tf.contrib.rnn.DropoutWrapper(\n",
    "                tf.nn.rnn_cell.LSTMCell(\n",
    "                    size_layers,\n",
    "                    initializer = tf.orthogonal_initializer(),\n",
    "                    reuse = reuse,\n",
    "                ),\n",
    "                state_keep_prob = dropout,\n",
    "                output_keep_prob = dropout,\n",
    "            )\n",
    "        def attention(encoder_out, seq_len, reuse=False):\n",
    "            attention_mechanism = tf.contrib.seq2seq.BahdanauAttention(num_units = size_layers, \n",
    "                                                                    memory = encoder_out,\n",
    "                                                                    memory_sequence_length = seq_len)\n",
    "            return tf.contrib.seq2seq.AttentionWrapper(\n",
    "            cell = tf.nn.rnn_cell.MultiRNNCell([cells(reuse) for _ in range(num_layers)]), \n",
    "                attention_mechanism = attention_mechanism,\n",
    "                attention_layer_size = size_layers,alignment_history = True)\n",
    "        \n",
    "        encoder_cells = tf.nn.rnn_cell.MultiRNNCell([cells() for _ in range(num_layers)])\n",
    "        encoder_out, encoder_state = tf.nn.dynamic_rnn(cell = encoder_cells, \n",
    "                                                                 inputs = forward, \n",
    "                                                                 sequence_length = seq_lens,\n",
    "                                                                 dtype = tf.float32)\n",
    "        \n",
    "        encoder_state = tuple(encoder_state[-1] for _ in range(num_layers))\n",
    "        decoder_cell = attention(encoder_out, seq_lens)\n",
    "        dense_layer = tf.layers.Dense(n_mels * resampled)\n",
    "        \n",
    "        training_helper = tf.contrib.seq2seq.TrainingHelper(\n",
    "                inputs = self.decoder_inputs,\n",
    "                sequence_length = seq_lens,\n",
    "                time_major = False)\n",
    "        training_decoder = tf.contrib.seq2seq.BasicDecoder(\n",
    "                cell = decoder_cell,\n",
    "                helper = training_helper,\n",
    "                initial_state = decoder_cell.zero_state(batch_size, tf.float32).clone(cell_state=encoder_state),\n",
    "                output_layer = dense_layer)\n",
    "        training_decoder_output, _, _ = tf.contrib.seq2seq.dynamic_decode(\n",
    "                decoder = training_decoder,\n",
    "                impute_finished = True,\n",
    "                maximum_iterations = tf.reduce_max(seq_lens))\n",
    "        \n",
    "        self.Y_hat = training_decoder_output.rnn_output\n",
    "        out_decoder2 = tf.reshape(\n",
    "                self.Y_hat, [tf.shape(self.Y_hat)[0], -1, n_mels]\n",
    "        )\n",
    "        dec = conv1d_banks(\n",
    "                out_decoder2, K = decoder_num_banks, is_training = self.training\n",
    "        )\n",
    "        dec = tf.layers.max_pooling1d(\n",
    "                dec, pool_size = 2, strides = 1, padding = 'same'\n",
    "        )\n",
    "        dec = tf.layers.conv1d(\n",
    "                dec,\n",
    "                embed_size // 2,\n",
    "                3,\n",
    "                name = 'decoder-conv1-1',\n",
    "                padding = 'SAME',\n",
    "            )\n",
    "        dec = tf.nn.relu(\n",
    "                tf.layers.batch_normalization(dec, training = self.training)\n",
    "            )\n",
    "        dec = tf.layers.conv1d(\n",
    "                dec,\n",
    "                embed_size // 2,\n",
    "                3,\n",
    "                name = 'decoder-conv1-2',\n",
    "                padding = 'SAME',\n",
    "            )\n",
    "        dec = tf.layers.batch_normalization(dec, training = self.training)\n",
    "        dec = tf.layers.dense(dec, embed_size // 2)\n",
    "        for i in range(4):\n",
    "            dec = highwaynet(\n",
    "                    dec,\n",
    "                    num_units = embed_size // 2,\n",
    "                    scope = 'decoder-highwaynet-{}'.format(i),\n",
    "                )\n",
    "        with tf.variable_scope('decoder-gru', reuse = False):\n",
    "            cell = tf.contrib.rnn.GRUCell(embed_size // 2)\n",
    "            cell_bw = tf.contrib.rnn.GRUCell(embed_size // 2)\n",
    "            outputs, _ = tf.nn.bidirectional_dynamic_rnn(\n",
    "                    cell, cell_bw, dec, dtype = tf.float32\n",
    "                )\n",
    "            outputs = tf.concat(outputs, 2)\n",
    "        self.Z_hat = tf.layers.dense(outputs, 1 + fourier_window_size // 2)\n",
    "        self.loss1 = tf.reduce_mean(tf.abs(self.Y_hat - self.Y))\n",
    "        self.loss2 = tf.reduce_mean(tf.abs(self.Z_hat - self.Z))\n",
    "        self.loss = self.loss1 + self.loss2\n",
    "        self.optimizer = tf.train.AdamOptimizer(\n",
    "            learning_rate = learning_rate\n",
    "        ).minimize(self.loss)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "tf.reset_default_graph()\n",
    "sess = tf.InteractiveSession()\n",
    "\n",
    "size_layers = 128\n",
    "learning_rate = 1e-3\n",
    "num_layers = 2\n",
    "\n",
    "model = Model(num_layers, size_layers, learning_rate)\n",
    "sess.run(tf.global_variables_initializer())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "paths, lengths, texts, raw_texts = [], [], [], []\n",
    "text_files = [f for f in os.listdir('mel') if f.endswith('.npy')]\n",
    "for fpath in text_files:\n",
    "    with open('%s/%s' % (path, fpath.replace('npy', 'txt'))) as fopen:\n",
    "        text = fopen.read()\n",
    "    paths.append(fpath.replace('.npy', ''))\n",
    "    text = text_normalize(text)\n",
    "    raw_texts.append(text)\n",
    "    text = text + 'E'\n",
    "    texts.append(np.array([char2idx[char] for char in text], np.int32))\n",
    "    lengths.append(len(text))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "def dynamic_batching(paths):\n",
    "    files, max_y, max_z = [], 0, 0\n",
    "    for n in range(len(paths)):\n",
    "        files.append(get_cached(paths[n]))\n",
    "        if files[-1][0].shape[0] > max_y:\n",
    "            max_y = files[-1][0].shape[0]\n",
    "        if files[-1][1].shape[0] > max_z:\n",
    "            max_z = files[-1][1].shape[0]\n",
    "    return files, max_y, max_z"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "minibatch loop: 100%|██████████| 88/88 [00:36<00:00,  2.46it/s, cost=0.153]\n",
      "minibatch loop: 100%|██████████| 88/88 [00:34<00:00,  2.60it/s, cost=0.0978]\n",
      "minibatch loop: 100%|██████████| 88/88 [00:34<00:00,  2.63it/s, cost=0.0886]\n",
      "minibatch loop: 100%|██████████| 88/88 [00:34<00:00,  2.65it/s, cost=0.0795]\n",
      "minibatch loop: 100%|██████████| 88/88 [00:34<00:00,  2.64it/s, cost=0.0741]\n",
      "minibatch loop: 100%|██████████| 88/88 [00:34<00:00,  2.65it/s, cost=0.0716]\n",
      "minibatch loop: 100%|██████████| 88/88 [00:34<00:00,  2.64it/s, cost=0.0696]\n",
      "minibatch loop: 100%|██████████| 88/88 [00:34<00:00,  2.69it/s, cost=0.0666]\n",
      "minibatch loop: 100%|██████████| 88/88 [00:34<00:00,  2.68it/s, cost=0.0648]\n",
      "minibatch loop: 100%|██████████| 88/88 [00:34<00:00,  2.68it/s, cost=0.0636]\n",
      "minibatch loop: 100%|██████████| 88/88 [00:34<00:00,  2.67it/s, cost=0.0615]\n",
      "minibatch loop: 100%|██████████| 88/88 [00:34<00:00,  2.66it/s, cost=0.0618]\n",
      "minibatch loop: 100%|██████████| 88/88 [00:34<00:00,  2.68it/s, cost=0.0636]\n",
      "minibatch loop: 100%|██████████| 88/88 [00:34<00:00,  2.68it/s, cost=0.0607]\n",
      "minibatch loop: 100%|██████████| 88/88 [00:34<00:00,  2.68it/s, cost=0.0592]\n",
      "minibatch loop: 100%|██████████| 88/88 [00:34<00:00,  2.65it/s, cost=0.0586]\n",
      "minibatch loop: 100%|██████████| 88/88 [00:34<00:00,  2.68it/s, cost=0.0568]\n",
      "minibatch loop: 100%|██████████| 88/88 [00:34<00:00,  2.68it/s, cost=0.0561]\n",
      "minibatch loop: 100%|██████████| 88/88 [00:34<00:00,  2.68it/s, cost=0.056] \n",
      "minibatch loop: 100%|██████████| 88/88 [00:34<00:00,  2.67it/s, cost=0.056] \n",
      "minibatch loop: 100%|██████████| 88/88 [00:34<00:00,  2.69it/s, cost=0.0551]\n",
      "minibatch loop: 100%|██████████| 88/88 [00:34<00:00,  2.66it/s, cost=0.0552]\n",
      "minibatch loop: 100%|██████████| 88/88 [00:34<00:00,  2.70it/s, cost=0.0538]\n",
      "minibatch loop: 100%|██████████| 88/88 [00:33<00:00,  2.70it/s, cost=0.0539]\n",
      "minibatch loop: 100%|██████████| 88/88 [00:33<00:00,  2.71it/s, cost=0.0523]\n",
      "minibatch loop: 100%|██████████| 88/88 [00:33<00:00,  2.71it/s, cost=0.0524]\n",
      "minibatch loop: 100%|██████████| 88/88 [00:33<00:00,  2.70it/s, cost=0.0525]\n",
      "minibatch loop: 100%|██████████| 88/88 [00:33<00:00,  2.69it/s, cost=0.0524]\n",
      "minibatch loop: 100%|██████████| 88/88 [00:33<00:00,  2.70it/s, cost=0.0512]\n",
      "minibatch loop: 100%|██████████| 88/88 [00:33<00:00,  2.72it/s, cost=0.0511]\n"
     ]
    }
   ],
   "source": [
    "EPOCH = 30\n",
    "for i in range(EPOCH):\n",
    "    pbar = tqdm(range(0, len(paths), batch_size), desc = 'minibatch loop')\n",
    "    for k in pbar:\n",
    "        index = min(k + batch_size, len(paths))\n",
    "        files, max_y, max_z = dynamic_batching(paths[k:index])\n",
    "        max_x = max(lengths[k:index])\n",
    "        batch_x = np.zeros((batch_size, max_x))\n",
    "        batch_y = np.zeros((batch_size, max_y, n_mels * resampled))\n",
    "        batch_z = np.zeros((batch_size, max_z, fourier_window_size // 2 + 1))\n",
    "        for n in range(len(files)):\n",
    "            batch_x[n, :] = np.pad(\n",
    "                texts[k + n],\n",
    "                ((0, max_x - texts[k + n].shape[0])),\n",
    "                mode = 'constant',\n",
    "            )\n",
    "            batch_y[n, :, :] = np.pad(\n",
    "                files[n][0],\n",
    "                ((0, max_y - files[n][0].shape[0]), (0, 0)),\n",
    "                mode = 'constant',\n",
    "            )\n",
    "            batch_z[n, :, :] = np.pad(\n",
    "                files[n][1],\n",
    "                ((0, max_z - files[n][1].shape[0]), (0, 0)),\n",
    "                mode = 'constant',\n",
    "            )\n",
    "        _, cost = sess.run(\n",
    "            [model.optimizer, model.loss],\n",
    "            feed_dict = {model.X: batch_x, model.Y: batch_y, model.Z: batch_z,\n",
    "                        model.training: True},\n",
    "        )\n",
    "        pbar.set_postfix(cost = cost)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 50/50 [00:00<00:00, 37.10it/s] \n"
     ]
    }
   ],
   "source": [
    "y_hat = np.zeros((1, 50, n_mels * resampled), np.float32)\n",
    "for j in tqdm(range(50)):\n",
    "    _y_hat = sess.run(model.Y_hat, {model.X: [texts[0]], model.Y: y_hat})\n",
    "    y_hat[:, j, :] = _y_hat[:, j, :]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "mags = sess.run(model.Z_hat, {model.Y_hat: y_hat,model.training:False})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "audio = spectrogram2wav(mags[0])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "saving: say the word burn\n"
     ]
    }
   ],
   "source": [
    "from scipy.io.wavfile import write\n",
    "print('saving: %s'%(raw_texts[0]))\n",
    "write(os.path.join('test.wav'), sample_rate, audio)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "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.6.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}