module_test_m.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import tensorflow as tf\n",
    "tf.reset_default_graph()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Input Data: .wav -> Pitch contour (f0s), Harmonic spectral envelope (sps), Aperiodic spectral envelope (aps)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import os\n",
    "import time\n",
    "import argparse\n",
    "import librosa\n",
    "import matplotlib.pyplot as plt\n",
    "import matplotlib as mpl\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pyworld\n",
    "import CycleGAN\n",
    "import readData"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "from utils import *\n",
    "from ops import *"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "import librosa.display\n",
    "from IPython.display import Audio\n",
    "import matplotlib.pyplot as plt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "%matplotlib inline\n",
    "plt.rcParams['figure.figsize'] = (16, 4)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "random_seed = 0\n",
    "np.random.seed(random_seed)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Autoencoder: Style_Encoder, Content_Encoder, MLP, Decoder, Discriminator"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "def Style_Encoder(inputs, style_dim=16, reuse=False, scope='style_encoder'):                                                            # [1, 24, 128] = [batch_size, feature_channel, time]\n",
    "\n",
    "    inputs = tf.transpose(inputs, perm=[0, 2, 1], name='input_transpose')                                                               # [1, 128, 24] = [batch_size, time, feature_channel]\n",
    "\n",
    "    with tf.variable_scope(scope, reuse=reuse):\n",
    "\n",
    "        h1 = conv1d_layer(inputs=inputs, filters=128, kernel_size=15, strides=1, name='h1_conv')                                        # [1, 128, 128]\n",
    "        h1_gates = conv1d_layer(inputs=inputs, filters=128, kernel_size=15, strides=1, name='h1_conv_gates')\n",
    "        h1_glu = gated_linear_layer(inputs=h1, gates=h1_gates, name='h1_glu')\n",
    "\n",
    "        # Downsample\n",
    "        d1 = downsample1d_block_withoutIN(inputs=h1_glu, filters=256, kernel_size=5, strides=2, name_prefix='downsample1d_block1')      # [1, 64, 256]\n",
    "        d2 = downsample1d_block_withoutIN(inputs=d1, filters=512, kernel_size=5, strides=2, name_prefix='downsample1d_block2')          # [1, 32, 512]\n",
    "\n",
    "        d3 = downsample1d_block_withoutIN(inputs=d2, filters=512, kernel_size=3, strides=2, name_prefix='downsample1d_block3')          # [1, 16, 512]\n",
    "        d4 = downsample1d_block_withoutIN(inputs=d3, filters=512, kernel_size=3, strides=2, name_prefix='downsample1d_block4')          # [1, 8, 512]\n",
    "\n",
    "        # Global Average Pooling\n",
    "        p1 = adaptive_avg_pooling(d4)                                                                                                   # [1, 1, 512]\n",
    "        style = conv1d_layer(inputs=p1, filters=style_dim, kernel_size=1, strides=1, name='SE_logit')                                   # [1, 1, 16]\n",
    "\n",
    "        return style                                                                                                                    # [1, 1, 16]\n",
    "\n",
    "\n",
    "def Content_Encoder(inputs, reuse=False, scope='content_encoder'):\n",
    "    # IN removes the original feature mean and variance that represent important style information\n",
    "    inputs = tf.transpose(inputs, perm=[0, 2, 1], name='input_transpose')                                                               # [1, 24, 128] = [batch_size, time, feature_channel]\n",
    "\n",
    "    with tf.variable_scope(scope, reuse=reuse):\n",
    "\n",
    "        h1 = conv1d_layer(inputs=inputs, filters=128, kernel_size=15, strides=1, name='h1_conv')                                        # [1, 128, 128]\n",
    "        h1_norm = instance_norm_layer(inputs=h1, name='h1_norm')\n",
    "        h1_gates = conv1d_layer(inputs=inputs, filters=128, kernel_size=15, strides=1, name='h1_gates')\n",
    "        h1_norm_gates = instance_norm_layer(inputs=h1_gates, name='h1_norm_gates')\n",
    "        h1_glu = gated_linear_layer(inputs=h1_norm, gates=h1_norm_gates, name='h1_glu')\n",
    "\n",
    "        # downsample\n",
    "        d1 = downsample1d_block(inputs=h1_glu, filters=256, kernel_size=5, strides=2, name_prefix='downsample1d_block1')                # [1, 64, 256]\n",
    "        d2 = downsample1d_block(inputs=d1, filters=512, kernel_size=5, strides=2, name_prefix='downsample1d_block2')                    # [1, 32, 512]\n",
    "               \n",
    "        # Residual blocks\n",
    "        r1 = residual1d_block(inputs=d2, filters = 512, kernel_size=3, strides=1, name_prefix='residual1d_block1')                      # [1, 32, 512]\n",
    "        r2 = residual1d_block(inputs=r1, filters = 512, kernel_size=3, strides=1, name_prefix='residual1d_block2')\n",
    "        r3 = residual1d_block(inputs=r2, filters = 512, kernel_size=3, strides=1, name_prefix='residual1d_block3')\n",
    "        content = residual1d_block(inputs=r3, filters=512, kernel_size=3, strides=1, name_prefix='residual1d_block4')\n",
    "\n",
    "        return content                                                                                                                  # [1, 32, 512]\n",
    "\n",
    "\n",
    "def MLP(style, reuse=False, scope='MLP'):                                                                                               # [1, 1, 16]\n",
    "\n",
    "    with tf.variable_scope(scope, reuse=reuse):\n",
    "\n",
    "        x1 = linear(style, 512, scope='linear_1')                                                                                       # [1, 1, 512]\n",
    "        x1_gates = linear(x1, 512, scope='linear_1_gates')\n",
    "        x1_glu = gated_linear_layer(inputs=x1, gates=x1_gates, name='x1_glu')\n",
    "\n",
    "        x2 = linear(x1_glu, 512, scope='linear_2')\n",
    "        x2_gates = linear(x2, 512, scope='linear_2_gates')\n",
    "        x2_glu = gated_linear_layer(inputs=x2, gates=x2_gates, name='x2_glu')\n",
    "\n",
    "        mu = linear(x2_glu, 512, scope='mu')\n",
    "        sigma = linear(x2_glu, 512, scope='sigma')\n",
    "\n",
    "        mu = tf.reshape(mu, shape=[-1, 1, 512])                                                                                         # [1, 1, 512]\n",
    "        sigma = tf.reshape(sigma, shape=[-1, 1, 512])                                                                                   # [1, 1, 512]\n",
    "\n",
    "        return mu, sigma                                                                                                                # [1, 1, 512]\n",
    "\n",
    "\n",
    "def Decoder(content, style, reuse=False, scope=\"decoder\"):\n",
    "\n",
    "    with tf.variable_scope(scope, reuse=reuse):\n",
    "\n",
    "        mu, sigma = MLP(style, reuse)                                                                                                   # [1, 1, 512]\n",
    "        x = content                                                                                                                     # [1, 32, 512]\n",
    "\n",
    "        # Adaptive Residual blocks\n",
    "        r1 = residual1d_block_adaptive(inputs=x, filters=512, mu=mu, sigma=sigma, kernel_size=3, strides=1, name_prefix='residual1d_block1')        # [1, 32, 512]\n",
    "        r2 = residual1d_block_adaptive(inputs=r1, filters=512, mu=mu, sigma=sigma, kernel_size=3, strides=1, name_prefix='residual1d_block2')\n",
    "        r3 = residual1d_block_adaptive(inputs=r2, filters=512, mu=mu, sigma=sigma, kernel_size=3, strides=1, name_prefix='residual1d_block3')\n",
    "\n",
    "        # Upsample\n",
    "        u1 = upsample1d_block(inputs=r3, filters=512, kernel_size=5, strides=1, shuffle_size=2, name_prefix='upsample1d_block1')        # [1, 64, 512]\n",
    "        u2 = upsample1d_block(inputs=u1, filters=256, kernel_size=5, strides=1, shuffle_size=2, name_prefix='upsample1d_block2')        # [1, 128, 256]\n",
    "\n",
    "        # Output\n",
    "        o1 = conv1d_layer(inputs=u2, filters=24, kernel_size=15, strides=1, name='o1_conv')                                             # [1, 128, 24]\n",
    "        o2 = tf.transpose(o1, perm=[0, 2, 1], name='output_transpose')                                                                  # [1, 24, 128]\n",
    "\n",
    "        return o2                                                                                                                       # [1, 24, 128] = [batch_size, feature_channel, time]\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "def Discriminator(inputs, reuse=False, scope='discriminator'):\n",
    "\n",
    "    # inputs = [batch_size, num_features, time]\n",
    "    # add channel for 2D convolution [batch_size, num_features, time, 1]\n",
    "    inputs = tf.expand_dims(inputs, -1)                                                                                                 # [1, 24, 128, 1]\n",
    "\n",
    "    with tf.variable_scope(scope, reuse=reuse):\n",
    "\n",
    "        h1 = conv2d_layer(inputs=inputs, filters=128, kernel_size=[3, 3], strides=[1, 2], name='h1_conv')                               # [1, 24, 64, 128]\n",
    "        h1_gates = conv2d_layer(inputs=inputs, filters=128, kernel_size=[3, 3], strides=[1, 2], name='h1_conv_gates')\n",
    "        h1_glu = gated_linear_layer(inputs=h1, gates=h1_gates, name='h1_glu')\n",
    "\n",
    "        # Downsample\n",
    "        d1 = downsample2d_block(inputs=h1_glu, filters=256, kernel_size=[3, 3], strides=[2, 2], name_prefix='downsample2d_block1')      # [1, 12, 32, 256]\n",
    "        d2 = downsample2d_block(inputs=d1, filters=512, kernel_size=[3, 3], strides=[2, 2], name_prefix='downsample2d_block2')          # [1, 6, 16, 512]\n",
    "        d3 = downsample2d_block(inputs=d2, filters=1024, kernel_size=[6, 3], strides=[1, 2], name_prefix='downsample2d_block3')         # [1, 6, 8, 1024]\n",
    "\n",
    "        # Output\n",
    "        o1 = tf.layers.dense(inputs=d3, units=1, activation=tf.nn.sigmoid)\n",
    "\n",
    "        return [o1]                                                                                                                       # [1, 6, 8, 1]\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Autoencoder Model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "##################################################################################\n",
    "# Model\n",
    "##################################################################################\n",
    "\n",
    "def Encoder_A(x_A, reuse=False):\n",
    "    style_A = Style_Encoder(x_A, reuse=reuse, scope='style_encoder_A')\n",
    "    content_A = Content_Encoder(x_A, reuse=reuse, scope='content_encoder_A')\n",
    "\n",
    "    return content_A, style_A\n",
    "\n",
    "def Encoder_B(x_B, reuse=False):\n",
    "    style_B = Style_Encoder(x_B, reuse=reuse, scope='style_encoder_B')\n",
    "    content_B = Content_Encoder(x_B, reuse=reuse, scope='content_encoder_B')\n",
    "\n",
    "    return content_B, style_B\n",
    "\n",
    "def Decoder_A(content_B, style_A, reuse=False):\n",
    "    x_ba = Decoder(content=content_B, style=style_A, reuse=reuse, scope='decoder_A')\n",
    "\n",
    "    return x_ba\n",
    "\n",
    "def Decoder_B(content_A, style_B, reuse=False):\n",
    "    x_ab = Decoder(content=content_A, style=style_B, reuse=reuse, scope='decoder_B')\n",
    "\n",
    "    return x_ab\n",
    "\n",
    "def discriminate_real(x_A, x_B):\n",
    "    real_A_logit = Discriminator(x_A, scope=\"discriminator_A\")\n",
    "    real_B_logit = Discriminator(x_B, scope=\"discriminator_B\")\n",
    "\n",
    "    return real_A_logit, real_B_logit\n",
    "\n",
    "def discriminate_fake(x_ba, x_ab):\n",
    "    fake_A_logit = Discriminator(x_ba, reuse=True, scope=\"discriminator_A\")\n",
    "    fake_B_logit = Discriminator(x_ab, reuse=True, scope=\"discriminator_B\")\n",
    "\n",
    "    return fake_A_logit, fake_B_logit"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Module: EmoMUNIT"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "class EmoMUNIT(object):\n",
    "    def __init__(self, sess):\n",
    "        \n",
    "        # self.train_A_dir = 'C:/Users/Gayane/Desktop/EmotionTask/F0_results/ang_neu/ang'\n",
    "        # self.train_B_dir = './../../../Database/Emotion/ang_neu/neu'\n",
    "        # self.validation_A_dir = './../../../Database/Emotion/ang_neu/val_ang'\n",
    "        # self.validation_B_dir = './../../../Database/Emotion/ang_neu/val_neu'\n",
    "\n",
    "        dir = \"C:/Users/Gayane/Desktop/EmotionTask/Data/IEMOCAP_full_release/\"\n",
    "        dirAudio = \"Session1/sentences/wav/\"\n",
    "        dirText = \"Session1/dialog/EmoEvaluation/\"\n",
    "        obj = readData.preProcessing(dir, dirAudio, dirText)\n",
    "        wavFoldersPaths =  obj.wavFilePath()\n",
    "        self.train_A_dir = wavFoldersPaths[0][:10]\n",
    "        self.train_B_dir = wavFoldersPaths[0][10:20]\n",
    "        self.validation_A_dir = wavFoldersPaths[1][:10]\n",
    "        self.validation_B_dir = wavFoldersPaths[2][:10]\n",
    "#         self.max_samples = 1000\n",
    "        \n",
    "        self.batch_size = 1\n",
    "        self.style_dim = 16\n",
    "        \n",
    "        self.Encoder_A = Encoder_A\n",
    "        self.Encoder_B = Encoder_B\n",
    "        self.Decoder_A = Decoder_A\n",
    "        self.Decoder_B = Decoder_B\n",
    "        self.discriminate_real = discriminate_real\n",
    "        self.discriminate_fake = discriminate_fake\n",
    "        \n",
    "        self.gan_type = 'lsgan'\n",
    "        \n",
    "        self.gan_w = 10.0\n",
    "        self.recon_x_w = 1.0\n",
    "        self.recon_s_w = 1.0\n",
    "        self.recon_c_w = 1.0\n",
    "        self.recon_x_cyc_w = 10.0\n",
    "               \n",
    "        self.audio_len = 128    # = n_frames, time_length\n",
    "        self.audio_ch = 24      # = num_mcep, num_features\n",
    "        \n",
    "        self.direction = 'A2B'\n",
    "        \n",
    "        self.model_name = 'EmoMUNIT'\n",
    "        self.gan_type = 'lsgan'\n",
    "        self.dataset_name = 'ang2neu'\n",
    "        self.log_dir = 'logs'\n",
    "        self.sample_dir = 'samples'\n",
    "        self.checkpoint_dir = 'checkpoint'\n",
    "        self.result_dir = 'results'\n",
    "        \n",
    "        self.sess = sess\n",
    "        self.epoch = 3# 1000\n",
    "        self.iteration = 3 # 1000\n",
    "        self.init_lr_D = 0.0001\n",
    "        self.init_lr_G = 0.0002\n",
    "        \n",
    "        self.sample_freq = 1000\n",
    "        self.save_freq = 1000\n",
    "        \n",
    "        self.sampling_rate = 16000\n",
    "        self.frame_period = 5.0\n",
    "        self.num_mcep = 24\n",
    "        \n",
    "    \n",
    "    def build_model(self):\n",
    "        self.lr_D = tf.placeholder(tf.float32, name='learning_rate_D')\n",
    "        self.lr_G = tf.placeholder(tf.float32, name='learning_rate_G')\n",
    "        \n",
    "        # Iterate from train_data_A and train_data_A\n",
    "        self.domain_A = tf.placeholder(tf.float32, shape=[self.batch_size, self.audio_ch, self.audio_len], name='domain_a')\n",
    "        self.domain_B = tf.placeholder(tf.float32, shape=[self.batch_size, self.audio_ch, self.audio_len], name='domain_b')\n",
    "    \n",
    "        self.style_a = tf.placeholder(tf.float32, shape=[self.batch_size, 1, self.style_dim], name='style_a')\n",
    "        self.style_b = tf.placeholder(tf.float32, shape=[self.batch_size, 1, self.style_dim], name='style_b')  \n",
    "    \n",
    "        # encode\n",
    "        content_a, style_a_prime = self.Encoder_A(self.domain_A)\n",
    "        content_b, style_b_prime = self.Encoder_B(self.domain_B)\n",
    "\n",
    "        # decode (within domain)\n",
    "        x_aa = self.Decoder_A(content_B=content_a, style_A=style_a_prime)\n",
    "        x_bb = self.Decoder_B(content_A=content_b, style_B=style_b_prime)\n",
    "    \n",
    "        # decode (cross domain)\n",
    "        x_ba = self.Decoder_A(content_B=content_b, style_A=self.style_a, reuse=True)\n",
    "        x_ab = self.Decoder_B(content_A=content_a, style_B=self.style_b, reuse=True)   \n",
    "    \n",
    "        # encode again\n",
    "        content_b_, style_a_ = self.Encoder_A(x_ba, reuse=True)\n",
    "        content_a_, style_b_ = self.Encoder_B(x_ab, reuse=True)    \n",
    "    \n",
    "        # decode again (if needed)\n",
    "        if self.recon_x_cyc_w > 0 :\n",
    "            x_aba = self.Decoder_A(content_B=content_a_, style_A=style_a_prime, reuse=True)\n",
    "            x_bab = self.Decoder_B(content_A=content_b_, style_B=style_b_prime, reuse=True)\n",
    "\n",
    "            cyc_recon_A = L1_loss(x_aba, self.domain_A)\n",
    "            cyc_recon_B = L1_loss(x_bab, self.domain_B)\n",
    "\n",
    "        else :\n",
    "            cyc_recon_A = 0.0\n",
    "            cyc_recon_B = 0.0    \n",
    "      \n",
    "        real_A_logit, real_B_logit = self.discriminate_real(self.domain_A, self.domain_B)\n",
    "        fake_A_logit, fake_B_logit = self.discriminate_fake(x_ba, x_ab)   \n",
    "    \n",
    "    \n",
    "        \"\"\" Define Loss \"\"\"\n",
    "        # Adversarial Loss\n",
    "        G_ad_loss_a = generator_loss(self.gan_type, fake_A_logit)\n",
    "        G_ad_loss_b = generator_loss(self.gan_type, fake_B_logit)\n",
    "    \n",
    "        # Discrimination Loss (real/fake)\n",
    "        D_ad_loss_a = discriminator_loss(self.gan_type, real_A_logit, fake_A_logit)\n",
    "        D_ad_loss_b = discriminator_loss(self.gan_type, real_B_logit, fake_B_logit)\n",
    "    \n",
    "        # Reconstruction Loss\n",
    "        recon_A = L1_loss(x_aa, self.domain_A) # reconstruction\n",
    "        recon_B = L1_loss(x_bb, self.domain_B) # reconstruction   \n",
    "    \n",
    "        # Semi-CycleGAN Loss\n",
    "        # For style, encourages diverse outputs given different style codes\n",
    "        recon_style_A = L1_loss(style_a_, self.style_a)\n",
    "        recon_style_B = L1_loss(style_b_, self.style_b)\n",
    "    \n",
    "        # For content, encourages the translated image to preserve semantic content of the input image\n",
    "        recon_content_A = L1_loss(content_a_, content_a)\n",
    "        recon_content_B = L1_loss(content_b_, content_b)   \n",
    "    \n",
    "        # Attacker Loss\n",
    "        Generator_A_loss = self.gan_w * G_ad_loss_a + \\\n",
    "                                   self.recon_x_w * recon_A + \\\n",
    "                                   self.recon_s_w * recon_style_A + \\\n",
    "                                   self.recon_c_w * recon_content_A + \\\n",
    "                                   self.recon_x_cyc_w * cyc_recon_A\n",
    "\n",
    "        Generator_B_loss = self.gan_w * G_ad_loss_b + \\\n",
    "                           self.recon_x_w * recon_B + \\\n",
    "                           self.recon_s_w * recon_style_B + \\\n",
    "                           self.recon_c_w * recon_content_B + \\\n",
    "                           self.recon_x_cyc_w * cyc_recon_B   \n",
    "    \n",
    "        # Defender Loss\n",
    "        Discriminator_A_loss = self.gan_w * D_ad_loss_a\n",
    "        Discriminator_B_loss = self.gan_w * D_ad_loss_b\n",
    "    \n",
    "        # Total Loss\n",
    "        self.Generator_loss = Generator_A_loss + Generator_B_loss\n",
    "        self.Discriminator_loss = Discriminator_A_loss + Discriminator_B_loss\n",
    "    \n",
    "    \n",
    "        \"\"\" Training Variables \"\"\"\n",
    "        t_vars = tf.trainable_variables()\n",
    "        G_vars = [var for var in t_vars if 'decoder' in var.name or 'encoder' in var.name]\n",
    "        D_vars = [var for var in t_vars if 'discriminator' in var.name]\n",
    "    \n",
    "        self.G_optim = tf.train.AdamOptimizer(self.lr_G, beta1=0.5, beta2=0.999).minimize(self.Generator_loss, var_list=G_vars)\n",
    "        self.D_optim = tf.train.AdamOptimizer(self.lr_D, beta1=0.5, beta2=0.999).minimize(self.Discriminator_loss, var_list=D_vars)\n",
    "    \n",
    "        \"\"\"\" Summary \"\"\"\n",
    "        self.all_G_loss = tf.summary.scalar(\"Generator_loss\", self.Generator_loss)\n",
    "        self.all_D_loss = tf.summary.scalar(\"Discriminator_loss\", self.Discriminator_loss)\n",
    "        \n",
    "        self.R_A_loss = tf.summary.scalar(\"Reconstruction_A_loss\", recon_A)\n",
    "        self.R_B_loss = tf.summary.scalar(\"Reconstruction_B_loss\", recon_B)\n",
    "        self.recon_style_A = tf.summary.scalar(\"recon_style_A\", recon_style_A)\n",
    "        self.recon_style_B = tf.summary.scalar(\"recon_style_B\", recon_style_B)        \n",
    "        self.recon_content_A = tf.summary.scalar(\"recon_content_A\", recon_content_A)\n",
    "        self.recon_content_B = tf.summary.scalar(\"recon_content_B\", recon_content_B)         \n",
    "        \n",
    "        self.G_A_loss = tf.summary.scalar(\"G_A_loss\", Generator_A_loss)\n",
    "        self.G_B_loss = tf.summary.scalar(\"G_B_loss\", Generator_B_loss)\n",
    "        self.G_ad_loss_a = tf.summary.scalar(\"G_ad_loss_a\", G_ad_loss_a)\n",
    "        self.G_ad_loss_b = tf.summary.scalar(\"G_ad_loss_b\", G_ad_loss_b)\n",
    "        \n",
    "        self.D_A_loss = tf.summary.scalar(\"D_A_loss\", Discriminator_A_loss)\n",
    "        self.D_B_loss = tf.summary.scalar(\"D_B_loss\", Discriminator_B_loss)\n",
    "\n",
    "        self.G_loss = tf.summary.merge([self.R_A_loss, self.R_B_loss, self.recon_style_A, self.recon_style_B, self.recon_content_A, self.recon_content_B, self.G_ad_loss_a, self.G_ad_loss_b, self.G_A_loss, self.G_B_loss, self.all_G_loss])\n",
    "        self.D_loss = tf.summary.merge([self.D_A_loss, self.D_B_loss, self.all_D_loss])\n",
    "    \n",
    "    \n",
    "        \"\"\" Speech: real and fake \"\"\"\n",
    "        self.real_A = self.domain_A\n",
    "        self.real_B = self.domain_B\n",
    "\n",
    "        self.fake_A = x_ba\n",
    "        self.fake_B = x_ab \n",
    "    \n",
    "        \"\"\" Test Variables \"\"\"\n",
    "        self.test_domain_A = tf.placeholder(tf.float32, [1, self.audio_ch, None], name='test_domain_a') # [1 24 None]\n",
    "        self.test_domain_B = tf.placeholder(tf.float32, [1, self.audio_ch, None], name='test_domain_b') # [1 24 None]\n",
    "        \n",
    "        self.test_style_a = tf.placeholder(tf.float32, [1, 1, self.style_dim], name='test_style_a')   # [1 1 16]\n",
    "        self.test_style_b = tf.placeholder(tf.float32, [1, 1, self.style_dim], name='test_style_b')   # [1 1 16]\n",
    "        \n",
    "        test_content_a, test_style_a = self.Encoder_A(self.test_domain_A, reuse=True)\n",
    "        test_content_b, test_style_b = self.Encoder_B(self.test_domain_B, reuse=True)\n",
    "\n",
    "        self.test_fake_A = self.Decoder_A(content_B=test_content_b, style_A=self.test_style_a, reuse=True)\n",
    "        self.test_fake_B = self.Decoder_B(content_A=test_content_a, style_B=self.test_style_b, reuse=True)\n",
    "\n",
    "        self.test_recon_A = self.Decoder_A(content_B=test_content_a, style_A=test_style_a, reuse=True)\n",
    "        self.test_recon_B = self.Decoder_B(content_A=test_content_b, style_B=test_style_b, reuse=True)\n",
    "        \n",
    "        \n",
    "        \"\"\" Guided Speech Translation \"\"\"\n",
    "        self.content_audio = tf.placeholder(tf.float32, [1, self.audio_ch, self.audio_len], name='content_audio')\n",
    "        self.style_audio = tf.placeholder(tf.float32, [1, self.audio_ch, self.audio_len], name='guide_style_audio_ch')\n",
    "\n",
    "        if self.direction == 'A2B' :\n",
    "            guide_content_A, guide_style_A = self.Encoder_A(self.content_audio, reuse=True)\n",
    "            guide_content_B, guide_style_B = self.Encoder_B(self.style_audio, reuse=True)\n",
    "\n",
    "        else :\n",
    "            guide_content_B, guide_style_B = self.Encoder_B(self.content_audio, reuse=True)\n",
    "            guide_content_A, guide_style_A = self.Encoder_A(self.style_audio, reuse=True)\n",
    "\n",
    "        self.guide_fake_A = self.Decoder_A(content_B=guide_content_B, style_A=guide_style_A, reuse=True)\n",
    "        self.guide_fake_B = self.Decoder_B(content_A=guide_content_A, style_B=guide_style_B, reuse=True)\n",
    "    \n",
    "    \n",
    "    def data_prepare(self, f0s_A, f0s_B, coded_sps_norm_A, coded_sps_norm_B):\n",
    "        \n",
    "        train_data_A = sample_train_data03(sps=list(coded_sps_norm_A), f0s=list(f0s_A), n_frames=self.audio_len)\n",
    "        train_data_B = sample_train_data03(sps=list(coded_sps_norm_B), f0s=list(f0s_B), n_frames=self.audio_len)\n",
    "\n",
    "        minlen = min(len(train_data_A), len(train_data_B))\n",
    "        np.random.shuffle(train_data_A)\n",
    "        np.random.shuffle(train_data_B)\n",
    "        train_data_A = np.array(train_data_A[0:minlen])\n",
    "        train_data_B = np.array(train_data_B[0:minlen])\n",
    "\n",
    "        return train_data_A, train_data_B\n",
    "    \n",
    "    \n",
    "    def train(self):\n",
    "        # initialize all variables\n",
    "        tf.global_variables_initializer().run()\n",
    "        \n",
    "        # saver to save model\n",
    "        self.saver = tf.train.Saver()\n",
    "        \n",
    "        # summary writer\n",
    "        self.writer = tf.summary.FileWriter(self.log_dir + '/' + self.model_dir, self.sess.graph)\n",
    "        \n",
    "        # restore check-point if it exits\n",
    "        could_load, checkpoint_counter = self.load(self.checkpoint_dir)\n",
    "        if could_load:\n",
    "            start_epoch = (int)(checkpoint_counter / self.iteration)\n",
    "            start_batch_id = checkpoint_counter - start_epoch * self.iteration\n",
    "            counter = checkpoint_counter\n",
    "            print(\" [*] Load SUCCESS\")\n",
    "        else:\n",
    "            start_epoch = 0\n",
    "            start_batch_id = 0\n",
    "            counter = 1\n",
    "            print(\" [!] Load FAILED...\")\n",
    "            \n",
    "        # check sample_dir    \n",
    "        check_folder(self.sample_dir)\n",
    "        \n",
    "        \n",
    "        '''Training loop for epoch'''\n",
    "        \n",
    "        # load data and extract features\n",
    "        f0s_A, coded_sps_norm_A, log_f0s_mean_A, log_f0s_std_A, coded_sps_mean_A, coded_sps_std_A = vocoder_extract(self.train_A_dir)\n",
    "        f0s_B, coded_sps_norm_B, log_f0s_mean_B, log_f0s_std_B, coded_sps_mean_B, coded_sps_std_B = vocoder_extract(self.train_B_dir)\n",
    "        \n",
    "        # load validation data\n",
    "        wavs_val_A = load_wavs(wav_dir=self.validation_A_dir, sr=self.sampling_rate)\n",
    "        wavs_val_B = load_wavs(wav_dir=self.validation_B_dir, sr=self.sampling_rate)\n",
    "\n",
    "        start_time = time.time()\n",
    "        for epoch in range(start_epoch, self.epoch):\n",
    "            \n",
    "            train_data_A, train_data_B = self.data_prepare(f0s_A, f0s_B, coded_sps_norm_A, coded_sps_norm_B)\n",
    "            print('Epoch[%d]: Input data sampled from %d A and %d B audio files: train_data_A' %(epoch, len(f0s_A), len(f0s_B)), np.shape(train_data_A), 'train_data_B', np.shape(train_data_B))\n",
    "\n",
    "            lr_D, lr_G = self.init_lr_D * pow(0.995, epoch), self.init_lr_G * pow(0.995, epoch)\n",
    "            for idx in range(start_batch_id, self.iteration):\n",
    "                style_a = np.random.normal(loc=0.0, scale=1.0, size=[self.batch_size, 1, self.style_dim])\n",
    "                style_b = np.random.normal(loc=0.0, scale=1.0, size=[self.batch_size, 1, self.style_dim])\n",
    "                \n",
    "                idx_A = idx%len(train_data_A)\n",
    "                idx_B = idx%len(train_data_B)\n",
    "                domain_A = train_data_A[idx_A:idx_A+1].astype('float32')\n",
    "                domain_B = train_data_B[idx_B:idx_B+1].astype('float32')\n",
    "                \n",
    "                train_feed_dict = {\n",
    "                    self.style_a : style_a,\n",
    "                    self.style_b : style_b,\n",
    "                    self.lr_D : lr_D,\n",
    "                    self.lr_G : lr_G,\n",
    "                    self.domain_A : domain_A,\n",
    "                    self.domain_B : domain_B\n",
    "                }\n",
    "                \n",
    "                # Update D\n",
    "                _, d_loss, summary_str = self.sess.run([self.D_optim, self.Discriminator_loss, self.D_loss], feed_dict = train_feed_dict)\n",
    "                self.writer.add_summary(summary_str, counter)\n",
    "            \n",
    "                # Update G\n",
    "                batch_A_audios, batch_B_audios, fake_A, fake_B, _, g_loss, summary_str = \\\n",
    "                self.sess.run([self.real_A, self.real_B, self.fake_A, self.fake_B, self.G_optim, \\\n",
    "                               self.Generator_loss, self.G_loss], feed_dict = train_feed_dict)\n",
    "                self.writer.add_summary(summary_str, counter)           \n",
    "            \n",
    "                # display training status\n",
    "                counter += 1\n",
    "                print(\"Epoch: [%2d] [%6d/%6d] time: %4.4f d_loss: %.8f, g_loss: %.8f\" \\\n",
    "                      % (epoch, idx, self.iteration, time.time() - start_time, d_loss, g_loss), end='\\r')\n",
    "            \n",
    "                # save generated samples\n",
    "                if np.mod(counter+1, self.sample_freq) == 0:\n",
    "                    # A2B\n",
    "                    idx_val_A = (counter//self.sample_freq)%len(wavs_val_A)\n",
    "                    wav = wavs_val_A[idx_val_A]\n",
    "                    wav = wav_padding(wav = wav, sr = self.sampling_rate, frame_period = self.frame_period, multiple = 4)\n",
    "                    # f0 conversion\n",
    "                    f0, timeaxis, sp, ap = world_decompose(wav = wav, fs = self.sampling_rate, frame_period = self.frame_period)\n",
    "                    f0_converted = pitch_conversion(f0 = f0, mean_log_src = log_f0s_mean_A, std_log_src = log_f0s_std_A, mean_log_target = log_f0s_mean_B, std_log_target = log_f0s_std_B)\n",
    "                    # sp normalization\n",
    "                    coded_sp = world_encode_spectral_envelop(sp = sp, fs = self.sampling_rate, dim = self.num_mcep)\n",
    "                    coded_sp_transposed = coded_sp.T\n",
    "                    coded_sp_norm = (coded_sp_transposed - coded_sps_mean_A) / coded_sps_std_A\n",
    "                    # random sampled style\n",
    "                    test_style_b = np.random.normal(loc=0.0, scale=1.0, size=[1, 1, self.style_dim])\n",
    "                    # sp conversion (A2B)\n",
    "                    coded_sp_converted_norm = self.sess.run(self.test_fake_B, feed_dict = {self.test_domain_A: np.array([coded_sp_norm]), self.test_style_b : test_style_b})\n",
    "                    coded_sp_converted_norm_recon = self.sess.run(self.test_recon_A, feed_dict = {self.test_domain_A: np.array([coded_sp_norm])})\n",
    "                    # [1,24,None]\n",
    "                    # de-normalization\n",
    "                    coded_sp_converted = coded_sp_converted_norm[0] * coded_sps_std_B + coded_sps_mean_B\n",
    "                    coded_sp_converted = coded_sp_converted.T\n",
    "                    coded_sp_converted = np.ascontiguousarray(coded_sp_converted)\n",
    "                    coded_sp_converted_recon = coded_sp_converted_norm_recon[0] * coded_sps_std_A + coded_sps_mean_A\n",
    "                    coded_sp_converted_recon = coded_sp_converted_recon.T\n",
    "                    coded_sp_converted_recon = np.ascontiguousarray(coded_sp_converted_recon)\n",
    "                    # combine converted f0, sp and ap\n",
    "                    decoded_sp_converted = world_decode_spectral_envelop(coded_sp = coded_sp_converted, fs = self.sampling_rate)\n",
    "                    decoded_sp_converted_recon = world_decode_spectral_envelop(coded_sp = coded_sp_converted_recon, fs = self.sampling_rate)                   \n",
    "                    wav_transformed = world_speech_synthesis(f0 = f0_converted, decoded_sp = decoded_sp_converted, ap = ap, fs = self.sampling_rate, frame_period = self.frame_period)\n",
    "                    wav_transformed_recon = world_speech_synthesis(f0 = f0, decoded_sp = decoded_sp_converted_recon, ap = ap, fs = self.sampling_rate, frame_period = self.frame_period)\n",
    "                    # write .wav file\n",
    "                    path_A2B = './{}/fake_A2B_id{:03d}_iter{:03d}K.wav'.format(self.sample_dir, idx_val_A, counter//1000)\n",
    "                    path_A2A = './{}/recon_A2A_id{:03d}_iter{:03d}K.wav'.format(self.sample_dir, idx_val_A, counter//1000)\n",
    "                    save_audio(wav=wav_transformed, path=path_A2B, sr=self.sampling_rate)\n",
    "                    save_audio(wav=wav_transformed_recon, path=path_A2A, sr=self.sampling_rate)\n",
    "                    \n",
    "                \n",
    "                # save checkpoints\n",
    "                if np.mod(counter+1, self.save_freq) == 0 :\n",
    "                    self.save(self.checkpoint_dir, counter)\n",
    "        \n",
    "            # After an epoch, start_batch_id reset to zero\n",
    "            # non-zero value is only for the first epoch after loading pre-trained model\n",
    "            start_batch_id = 0\n",
    "\n",
    "            # save model for final step\n",
    "            self.save(self.checkpoint_dir, counter)\n",
    "            \n",
    "        print(\" [*] Training finished!\")\n",
    "        \n",
    "    \n",
    "    def test(self):\n",
    "        # initialize all variables\n",
    "        tf.global_variables_initializer().run()\n",
    "               \n",
    "        # saver to save model\n",
    "        self.saver = tf.train.Saver()\n",
    "    \n",
    "        # load check-point if it exits\n",
    "        could_load, checkpoint_counter = self.load(self.checkpoint_dir)\n",
    "    \n",
    "        if could_load :\n",
    "            print(\" [*] Load SUCCESS\")\n",
    "        else :\n",
    "            print(\" [!] Load FAILED...\")\n",
    "    \n",
    "        # check result_dir\n",
    "        check_folder(self.result_dir)\n",
    "        \n",
    "        # write html for visual comparison\n",
    "        \n",
    "        \n",
    "        # Get statistic from train_A, train_B\n",
    "        _, _, log_f0s_mean_A, log_f0s_std_A, coded_sps_A_mean, coded_sps_A_std = vocoder_extract(self.train_A_dir)\n",
    "        _, _, log_f0s_mean_B, log_f0s_std_B, coded_sps_B_mean, coded_sps_B_std = vocoder_extract(self.train_B_dir) \n",
    "        print('std_log_src:', log_f0s_std_A, 'std_log_target', log_f0s_std_B)\n",
    "        \n",
    "        \n",
    "        # A2B\n",
    "        test_files_A = os.listdir(self.validation_A_dir)\n",
    "        for i in range(len(test_files_A)):\n",
    "            file = test_files_A[i]\n",
    "            # filepath = os.path.join(self.validation_A_dir, file)\n",
    "            wav, _ = librosa.load(file, sr = self.sampling_rate, mono = True)\n",
    "            wav = wav_padding(wav = wav, sr = self.sampling_rate, frame_period = self.frame_period, multiple = 4)\n",
    "            f0, timeaxis, sp, ap = world_decompose(wav = wav, fs = self.sampling_rate, frame_period = self.frame_period)\n",
    "            \n",
    "            # f0 conversion\n",
    "            f0_converted = pitch_conversion(f0 = f0, mean_log_src = log_f0s_mean_A, std_log_src = log_f0s_std_A, mean_log_target = log_f0s_mean_B, std_log_target = log_f0s_std_B)\n",
    "\n",
    "            # sp normalization\n",
    "            coded_sp = world_encode_spectral_envelop(sp = sp, fs = self.sampling_rate, dim = self.num_mcep)\n",
    "            coded_sp_transposed = coded_sp.T\n",
    "            coded_sp_norm = (coded_sp_transposed - coded_sps_A_mean) / coded_sps_A_std\n",
    "            \n",
    "            # random sampled style\n",
    "            test_style_b = np.random.normal(loc=0.0, scale=1.0, size=[1, 1, self.style_dim])\n",
    "            \n",
    "            # sp conversion (A2B)\n",
    "            coded_sp_converted_norm = self.sess.run(self.test_fake_B, feed_dict = {self.test_domain_A: np.array([coded_sp_norm]), self.test_style_b : test_style_b})\n",
    "            # [1,24,None]\n",
    "            \n",
    "            # print('coded_sp_converted_norm', np.shape(coded_sp_converted_norm[0]), 'coded_sps_B_mean', np.shape(coded_sps_B_mean), 'coded_sps_B_std:', np.shape(coded_sps_B_std))          \n",
    "            coded_sp_converted = coded_sp_converted_norm[0] * coded_sps_B_std + coded_sps_B_mean\n",
    "            coded_sp_converted = coded_sp_converted.T\n",
    "            coded_sp_converted = np.ascontiguousarray(coded_sp_converted)\n",
    "            decoded_sp_converted = world_decode_spectral_envelop(coded_sp = coded_sp_converted, fs = self.sampling_rate)\n",
    "            wav_transformed = world_speech_synthesis(f0 = f0_converted, decoded_sp = decoded_sp_converted, ap = ap, fs = self.sampling_rate, frame_period = self.frame_period)\n",
    "            librosa.output.write_wav(os.path.join(self.result_dir, os.path.basename(file)), wav_transformed, self.sampling_rate)\n",
    "            \n",
    "            print('converting test samples: [%d/%d]' %(i+1, len(test_files_A)), end='\\r')\n",
    "            \n",
    "        print(\" [*] Testing finished!\")\n",
    "        \n",
    "        \n",
    "        \n",
    "    @property\n",
    "    def model_dir(self):\n",
    "        return \"{}_{}_{}\".format(self.model_name, self.dataset_name, self.gan_type)\n",
    "    \n",
    "    \n",
    "    def load(self, checkpoint_dir):\n",
    "        import re\n",
    "        print(\" [*] Reading checkpoints...\")\n",
    "        checkpoint_dir = os.path.join(checkpoint_dir, self.model_dir)\n",
    "\n",
    "        ckpt = tf.train.get_checkpoint_state(checkpoint_dir)\n",
    "        if ckpt and ckpt.model_checkpoint_path:\n",
    "            ckpt_name = os.path.basename(ckpt.model_checkpoint_path)\n",
    "            self.saver.restore(self.sess, os.path.join(checkpoint_dir, ckpt_name))\n",
    "            counter = int(next(re.finditer(\"(\\d+)(?!.*\\d)\", ckpt_name)).group(0))\n",
    "            print(\" [*] Success to read {}\".format(ckpt_name))\n",
    "            return True, counter\n",
    "        else:\n",
    "            print(\" [*] Failed to find a checkpoint\")\n",
    "            return False, 0\n",
    "        \n",
    "        \n",
    "    def save(self, checkpoint_dir, step):\n",
    "        checkpoint_dir = os.path.join(checkpoint_dir, self.model_dir)\n",
    "\n",
    "        if not os.path.exists(checkpoint_dir):\n",
    "            os.makedirs(checkpoint_dir)\n",
    "\n",
    "        self.saver.save(self.sess, os.path.join(checkpoint_dir, self.model_name + '.model'), global_step=step)        \n",
    "        \n",
    "        \n",
    "        \n",
    "        "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Train"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      " [*] Reading checkpoints...\n",
      " [*] Failed to find a checkpoint\n",
      " [!] Load FAILED...\n",
      "Epoch[0]: Input data sampled from 10 A and 10 B audio files: train_data_A (34, 24, 128) train_data_B (34, 24, 128)\n",
      "saved file at ./samples/fake_A2B_id000_iter000K.wav shape: (149760,)\n",
      "saved file at ./samples/recon_A2A_id000_iter000K.wav shape: (149760,)\n",
      "Epoch[1]: Input data sampled from 10 A and 10 B audio files: train_data_A (34, 24, 128) train_data_B (34, 24, 128)\n",
      "saved file at ./samples/fake_A2B_id001_iter001K.wav shape: (47360,)\n",
      "saved file at ./samples/recon_A2A_id001_iter001K.wav shape: (47360,)\n",
      "Epoch[2]: Input data sampled from 10 A and 10 B audio files: train_data_A (36, 24, 128) train_data_B (36, 24, 128)\n"
     ]
    }
   ],
   "source": [
    "with tf.Session(config=tf.ConfigProto(allow_soft_placement=True)) as sess:\n",
    "    gan = EmoMUNIT(sess)\n",
    "    gan.build_model()\n",
    "    gan.train()\n",
    "    gan.test()\n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "interpreter": {
   "hash": "69191a637d571d0fc652f70933094b43ef8be802af5773b5a7842a5fe03dd230"
  },
  "kernelspec": {
   "display_name": "Python 3.5.6 64-bit ('finalenv': conda)",
   "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.5.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}