hw4/vae.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/usr/local/lib/python3.5/dist-packages/h5py/__init__.py:36: FutureWarning: Conversion of the second argument of issubdtype from `float` to `np.floating` is deprecated. In future, it will be treated as `np.float64 == np.dtype(float).type`.\n",
      "  from ._conv import register_converters as _register_converters\n"
     ]
    }
   ],
   "source": [
    "import os\n",
    "os.environ['CUDA_VISIBLE_DEVICES'] = '1'\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "import tensorflow as tf\n",
    "\n",
    "import matplotlib.pyplot as plt\n",
    "import matplotlib.gridspec as gridspec\n",
    "%matplotlib inline"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import time\n",
    "import numpy as np\n",
    "import tensorflow as tf\n",
    "\n",
    "class VAE():\n",
    "    def __init__(self):\n",
    "        super().__init__()\n",
    "\n",
    "    def load(self, npy_path):\n",
    "        self.data_dict = np.load(npy_path, encoding='latin1').item()\n",
    "        print(\"Load %s as self.data_dict\" % npy_path)\n",
    "\n",
    "    def build(self, n_dim=512, lambda_kl=1e-5, shape=(64,64,3)):\n",
    "        \"\"\"\n",
    "        load pre-trained weights from path\n",
    "        :param vgg16_npy_path: file path of vgg16 pre-trained weights\n",
    "        \"\"\"\n",
    "        # input information\n",
    "        self.H, self.W, self.C = shape\n",
    "        self.n_dim = n_dim\n",
    "        self.lambda_kl = lambda_kl\n",
    "        \n",
    "        # parameter dictionary\n",
    "        self.para_dict = dict()\n",
    "        self.data_dict = dict()\n",
    "        self.net_shape = dict()\n",
    "\n",
    "        # input placeholder\n",
    "        self.x = tf.placeholder(tf.float32, [None, self.H, self.W, self.C])\n",
    "        self.is_train = tf.placeholder(tf.bool)\n",
    "        self.random_sample = tf.placeholder(tf.float32, [None, self.n_dim])\n",
    "        \n",
    "        # normalize inputs\n",
    "        # self.x = self.x/255.0\n",
    "        assert self.x.get_shape().as_list()[1:] == [self.H, self.W, self.C]\n",
    "        \n",
    "        with tf.variable_scope(\"VAE\",reuse=tf.AUTO_REUSE):\n",
    "            self.mean, self.logvar = self.encoder(self.x)\n",
    "            self.sample = self.sampler(self.mean, self.logvar)\n",
    "            self.output_image = self.decoder(self.sample)\n",
    "            \n",
    "            self.recon_loss = tf.reduce_mean(tf.square(tf.subtract(self.x, self.output_image)), [1,2,3])\n",
    "            self.kl_loss = 0.5*tf.reduce_mean(tf.subtract(tf.add(tf.square(self.mean), tf.exp(self.logvar)), tf.add(1.0, self.logvar)),1)\n",
    "            self.vae_loss = self.recon_loss + self.lambda_kl*self.kl_loss\n",
    "            \n",
    "            # Sampling from random z\n",
    "            self.random_sample_images = self.decoder(self.random_sample)\n",
    "    \n",
    "    def encoder(self, input_image):\n",
    "        # conv\n",
    "        conv1 = self.conv_bn_layer(input_image, shape=(4,4,3,32), stride=2, name=\"conv1\")\n",
    "        conv2 = self.conv_bn_layer(conv1 , shape=(4,4,32,64), stride=2, name=\"conv2\")\n",
    "        conv3 = self.conv_bn_layer(conv2 , shape=(4,4,64,128), stride=2, name=\"conv3\")\n",
    "        conv4 = self.conv_bn_layer(conv3 , shape=(4,4,128,256), stride=2, name=\"conv4\")\n",
    "        flatten = self.flatten_layer(conv4, name='flatten')\n",
    "\n",
    "        # mean and logvar\n",
    "        mean = self.dense_layer(flatten, n_hidden=self.n_dim, name='mean')\n",
    "        logvar = self.dense_layer(flatten, n_hidden=self.n_dim, name='logvar')\n",
    "        return mean, logvar\n",
    "\n",
    "    def sampler(self, mean, logvar):\n",
    "        eps = tf.random_normal(shape=tf.shape(mean))\n",
    "        return mean + tf.exp(logvar / 2) * eps\n",
    "        \n",
    "    def decoder(self, sample_input):\n",
    "        deconv_fc1 = self.dense_layer(sample_input, n_hidden=self.net_shape['flatten'][1], name='deconv_fc1')\n",
    "        deconv_input = tf.reshape(deconv_fc1, shape=[-1, 4, 4, 256])\n",
    "        \n",
    "        batch_size = tf.shape(sample_input)[0]\n",
    "        \n",
    "        deconv1 = self.trans_conv_layer(bottom=deconv_input, shape=(4,4,128,256),\n",
    "                                        output_shape=[batch_size, 8, 8, 128], stride=2, name='deconv1')\n",
    "        deconv2 = self.trans_conv_layer(bottom=deconv1, shape=(4,4,64,128),\n",
    "                                        output_shape=[batch_size, 16, 16, 64], stride=2, name='deconv2')\n",
    "        deconv3 = self.trans_conv_layer(bottom=deconv2, shape=(4,4,32,64),\n",
    "                                        output_shape=[batch_size, 32, 32, 32], stride=2, name='deconv3')\n",
    "        output = self.trans_conv_layer(bottom=deconv3, shape=(4,4,3,32),\n",
    "                                        output_shape=[batch_size, self.H, self.W, self.C], activation='tanh', stride=2, name='deconv_output')\n",
    "        return (output/2) + 0.5\n",
    "\n",
    "    def dense_layer(self, bottom, n_hidden=None, name=None):\n",
    "        bottom_shape = bottom.get_shape().as_list()\n",
    "        if n_hidden is not None:\n",
    "            W = self.get_weights(shape=(bottom_shape[1], n_hidden), name=name)\n",
    "            b = self.get_bias(shape=n_hidden, name=name)\n",
    "        elif name in self.data_dict.keys():\n",
    "            W = self.get_weights(name=name)\n",
    "            b = self.get_bias(name=name)\n",
    "        else:\n",
    "            print(\"Neither give a shape nor lack a pre-trained layer called %s\" % name)\n",
    "        self.para_dict[name] = [W, b]\n",
    "        fc = tf.nn.bias_add(tf.matmul(bottom, W), b)\n",
    "        self.net_shape[name] = fc.get_shape().as_list()\n",
    "        return fc\n",
    "\n",
    "    def flatten_layer(self, bottom, name):\n",
    "        shape = bottom.get_shape().as_list()\n",
    "        dim = 1\n",
    "        for d in shape[1:]:\n",
    "            dim *= d\n",
    "        flatten = tf.reshape(bottom, [-1, dim])\n",
    "        self.net_shape[name] = flatten.get_shape().as_list()\n",
    "        return flatten\n",
    "\n",
    "    def avg_pool_layer(self, bottom, name):\n",
    "        pool = tf.nn.avg_pool(bottom, ksize=[1, 2, 2, 1], strides=[1, 2, 2, 1], padding='SAME', name=name)\n",
    "        self.net_shape[name] = pool.get_shape().as_list()\n",
    "        return pool\n",
    "\n",
    "    def max_pool_layer(self, bottom, name):\n",
    "        pool = tf.nn.max_pool(bottom, ksize=[1, 2, 2, 1], strides=[1, 2, 2, 1], padding='SAME', name=name)\n",
    "        self.net_shape[name] = pool.get_shape().as_list()\n",
    "        return pool\n",
    "    \n",
    "    def dropout(self, bottom, keep_prob):\n",
    "        if self.is_train == True:\n",
    "            return tf.nn.dropout(bottom, keep_prob=keep_prob)\n",
    "        else:\n",
    "            return bottom\n",
    "\n",
    "    def trans_conv_layer(self, bottom, output_shape, stride, activation='relu', name=None, shape=None):\n",
    "        if shape is not None:\n",
    "            conv_filter, gamma, beta, bn_mean, bn_variance = self.get_conv_filter(shape=shape, name=name)\n",
    "            conv_bias = self.get_bias(shape=shape[2], name=name)\n",
    "        elif name in self.data_dict.keys():\n",
    "            conv_filter, gamma, beta, bn_mean, bn_variance = self.get_conv_filter(name=name)\n",
    "            conv_bias = self.get_bias(name=name)\n",
    "        else:\n",
    "            print(\"Neither give a shape nor lack a pre-trained layer called %s\" % name)\n",
    "\n",
    "        self.para_dict[name] = [conv_filter, conv_bias]\n",
    "        self.para_dict[name+\"_gamma\"] = gamma\n",
    "        self.para_dict[name+\"_beta\"] = beta\n",
    "        self.para_dict[name+\"_bn_mean\"] = bn_mean\n",
    "        self.para_dict[name+\"_bn_variance\"] = bn_variance\n",
    "\n",
    "        conv = tf.nn.conv2d_transpose(bottom, conv_filter, output_shape, strides=[1, stride, stride, 1], padding=\"SAME\")\n",
    "        conv = tf.nn.bias_add(conv, conv_bias)\n",
    "        \n",
    "        from tensorflow.python.training.moving_averages import assign_moving_average\n",
    "        def mean_var_with_update():\n",
    "            mean, variance = tf.nn.moments(conv, [0,1,2], name='moments')\n",
    "            with tf.control_dependencies([assign_moving_average(bn_mean, mean, 0.99),\n",
    "                                            assign_moving_average(bn_variance, variance, 0.99)]):\n",
    "                return tf.identity(mean), tf.identity(variance)\n",
    "\n",
    "        mean, variance = tf.cond(self.is_train, mean_var_with_update, lambda:(bn_mean, bn_variance))\n",
    "        conv = tf.nn.batch_normalization(conv, mean, variance, beta, gamma, 1e-05)\n",
    "        self.net_shape[name] = conv.get_shape().as_list()\n",
    "\n",
    "        if activation=='tanh':\n",
    "            print('tanh')\n",
    "            tanh = tf.nn.tanh(conv)\n",
    "            return tanh\n",
    "        else:\n",
    "            relu = tf.nn.leaky_relu(conv)\n",
    "            return relu\n",
    "\n",
    "    def conv_bn_layer(self, bottom, stride=1, activation='lrelu', name=None, shape=None):\n",
    "        if shape is not None:\n",
    "            conv_filter, gamma, beta, bn_mean, bn_variance = self.get_conv_filter(shape=shape, name=name)\n",
    "            conv_bias = self.get_bias(shape=shape[3], name=name)\n",
    "        elif name in self.data_dict.keys():\n",
    "            conv_filter, gamma, beta, bn_mean, bn_variance = self.get_conv_filter(name=name)\n",
    "            conv_bias = self.get_bias(name=name)\n",
    "        else:\n",
    "            print(\"Neither give a shape nor lack a pre-trained layer called %s\" % name)\n",
    "\n",
    "        self.para_dict[name] = [conv_filter, conv_bias]\n",
    "        self.para_dict[name+\"_gamma\"] = gamma\n",
    "        self.para_dict[name+\"_beta\"] = beta\n",
    "        self.para_dict[name+\"_bn_mean\"] = bn_mean\n",
    "        self.para_dict[name+\"_bn_variance\"] = bn_variance\n",
    "\n",
    "        conv = tf.nn.conv2d(bottom, conv_filter, [1, stride, stride, 1], padding='SAME')\n",
    "        conv = tf.nn.bias_add(conv, conv_bias)\n",
    "\n",
    "        from tensorflow.python.training.moving_averages import assign_moving_average\n",
    "        def mean_var_with_update():\n",
    "            mean, variance = tf.nn.moments(conv, [0,1,2], name='moments')\n",
    "            with tf.control_dependencies([assign_moving_average(bn_mean, mean, 0.99),\n",
    "                                            assign_moving_average(bn_variance, variance, 0.99)]):\n",
    "                return tf.identity(mean), tf.identity(variance)\n",
    "\n",
    "        mean, variance = tf.cond(self.is_train, mean_var_with_update, lambda:(bn_mean, bn_variance))\n",
    "\n",
    "        conv = tf.nn.batch_normalization(conv, mean, variance, beta, gamma, 1e-05)\n",
    "        self.net_shape[name] = conv.get_shape().as_list()\n",
    "\n",
    "        if activation=='tanh':\n",
    "            tanh = tf.nn.tanh(conv)\n",
    "            return tanh\n",
    "        else:\n",
    "            relu = tf.nn.leaky_relu(conv)\n",
    "            return relu\n",
    "\n",
    "    def get_conv_filter(self, shape=None, name=None, with_bn=True):\n",
    "        if shape is not None:\n",
    "            conv_filter = tf.get_variable(shape=shape, initializer=tf.truncated_normal_initializer(mean=0, stddev=0.1), name=name+\"_W\", dtype=tf.float32)\n",
    "        elif name in self.data_dict.keys():\n",
    "            conv_filter = tf.get_variable(initializer=self.data_dict[name][0], name=name+\"_W\")\n",
    "        else:\n",
    "            print(\"Neither give a shape nor lack a pre-trained layer called %s\" % name)\n",
    "            return None\n",
    "\n",
    "        if with_bn:\n",
    "            if 'deconv' in name:\n",
    "                H,W,O,C = conv_filter.get_shape().as_list()\n",
    "            else:\n",
    "                H,W,C,O = conv_filter.get_shape().as_list()\n",
    "\n",
    "            if name+\"_gamma\" in self.data_dict.keys(): \n",
    "                gamma = tf.get_variable(initializer=self.data_dict[name+\"_gamma\"], name=name+\"_gamma\")\n",
    "            else:\n",
    "                gamma = tf.get_variable(shape=(O,), initializer=tf.ones_initializer(), name=name+\"_gamma\")\n",
    "\n",
    "            if name+\"_beta\" in self.data_dict.keys(): \n",
    "                beta = tf.get_variable(initializer=self.data_dict[name+\"_beta\"], name=name+\"_beta\")\n",
    "            else:\n",
    "                beta = tf.get_variable(shape=(O,), initializer=tf.zeros_initializer(), name=name+'_beta')\n",
    "\n",
    "            if name+\"_bn_mean\" in self.data_dict.keys(): \n",
    "                bn_mean = tf.get_variable(initializer=self.data_dict[name+\"_bn_mean\"], name=name+\"_bn_mean\")\n",
    "            else:\n",
    "                bn_mean = tf.get_variable(shape=(O,), initializer=tf.zeros_initializer(), name=name+'_bn_mean')\n",
    "\n",
    "            if name+\"_bn_variance\" in self.data_dict.keys(): \n",
    "                bn_variance = tf.get_variable(initializer=self.data_dict[name+\"_bn_variance\"], name=name+\"_bn_variance\")\n",
    "            else:\n",
    "                bn_variance = tf.get_variable(shape=(O,), initializer=tf.ones_initializer(), name=name+'_bn_variance')\n",
    "            return conv_filter, gamma, beta, bn_mean, bn_variance\n",
    "        else:\n",
    "            return conv_filter\n",
    "    \n",
    "    def get_weights(self, shape=None, name=None):\n",
    "        if shape is not None:\n",
    "            return tf.get_variable(shape=shape, initializer=tf.truncated_normal_initializer(mean=0, stddev=0.1), name=name+\"_W\", dtype=tf.float32)\n",
    "        elif name in self.data_dict.keys(): \n",
    "            return tf.get_variable(initializer=self.data_dict[name][0], name=name+\"_W\")\n",
    "        else:\n",
    "            print(\"(get_weight) neither give a shape nor lack a pre-trained layer called %s\" % name)\n",
    "            return None\n",
    "            \n",
    "    def get_bias(self, shape=None, name=None):\n",
    "        if shape is not None:\n",
    "            return tf.get_variable(shape=shape, initializer=tf.truncated_normal_initializer(mean=0, stddev=0.1), name=name+\"_b\", dtype=tf.float32)\n",
    "        elif name in self.data_dict.keys(): \n",
    "            return tf.get_variable(initializer=self.data_dict[name][1], name=name+\"_b\")\n",
    "        else:\n",
    "            print(\"(get_bias) neither give a shape nor lack a pre-trained layer called %s\" % name)\n",
    "            return None\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Creating save_vae/\n",
      "Creating save_vae/recons\n",
      "Creating save_vae/samples\n",
      "Reading dataset...\n",
      "tanh\n",
      "tanh\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "9cb763e2649946d6aa4cf4c0c4b04502",
       "version_major": 2,
       "version_minor": 0
      },
      "text/html": [
       "<p>Failed to display Jupyter Widget of type <code>IntProgress</code>.</p>\n",
       "<p>\n",
       "  If you're reading this message in Jupyter Notebook or JupyterLab, it may mean\n",
       "  that the widgets JavaScript is still loading. If this message persists, it\n",
       "  likely means that the widgets JavaScript library is either not installed or\n",
       "  not enabled. See the <a href=\"https://ipywidgets.readthedocs.io/en/stable/user_install.html\">Jupyter\n",
       "  Widgets Documentation</a> for setup instructions.\n",
       "</p>\n",
       "<p>\n",
       "  If you're reading this message in another notebook frontend (for example, a static\n",
       "  rendering on GitHub or <a href=\"https://nbviewer.jupyter.org/\">NBViewer</a>),\n",
       "  it may mean that your frontend doesn't currently support widgets.\n",
       "</p>\n"
      ],
      "text/plain": [
       "IntProgress(value=0)"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "6d2b70710f0c4c24b2ea982e10551e1f",
       "version_major": 2,
       "version_minor": 0
      },
      "text/html": [
       "<p>Failed to display Jupyter Widget of type <code>IntProgress</code>.</p>\n",
       "<p>\n",
       "  If you're reading this message in Jupyter Notebook or JupyterLab, it may mean\n",
       "  that the widgets JavaScript is still loading. If this message persists, it\n",
       "  likely means that the widgets JavaScript library is either not installed or\n",
       "  not enabled. See the <a href=\"https://ipywidgets.readthedocs.io/en/stable/user_install.html\">Jupyter\n",
       "  Widgets Documentation</a> for setup instructions.\n",
       "</p>\n",
       "<p>\n",
       "  If you're reading this message in another notebook frontend (for example, a static\n",
       "  rendering on GitHub or <a href=\"https://nbviewer.jupyter.org/\">NBViewer</a>),\n",
       "  it may mean that your frontend doesn't currently support widgets.\n",
       "</p>\n"
      ],
      "text/plain": [
       "IntProgress(value=0)"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 1 (0), 29.33 sec >> train loss: 0.0823, train recon loss: 0.0206, train kl loss: 6.1661, val loss: 0.0312, val recon loss: 0.0133, val kl loss: 1.7939\n",
      "Epoch 2 (0), 25.41 sec >> train loss: 0.0251, train recon loss: 0.0124, train kl loss: 1.2638, val loss: 0.0217, val recon loss: 0.0112, val kl loss: 1.0471\n",
      "Epoch 3 (0), 25.29 sec >> train loss: 0.0192, train recon loss: 0.0111, train kl loss: 0.8133, val loss: 0.0173, val recon loss: 0.0101, val kl loss: 0.7207\n",
      "Epoch 4 (0), 24.93 sec >> train loss: 0.0163, train recon loss: 0.0103, train kl loss: 0.6036, val loss: 0.0153, val recon loss: 0.0096, val kl loss: 0.5714\n",
      "INFO:tensorflow:save_vae/model.ckpt-4 is not in all_model_checkpoint_paths. Manually adding it.\n",
      "Epoch 5 (0), 26.06 sec >> train loss: 0.0147, train recon loss: 0.0097, train kl loss: 0.4982, val loss: 0.0140, val recon loss: 0.0094, val kl loss: 0.4639\n",
      "Epoch 6 (0), 25.0 sec >> train loss: 0.0136, train recon loss: 0.0093, train kl loss: 0.4361, val loss: 0.0131, val recon loss: 0.0089, val kl loss: 0.4195\n",
      "Epoch 7 (0), 25.42 sec >> train loss: 0.0129, train recon loss: 0.0089, train kl loss: 0.3991, val loss: 0.0123, val recon loss: 0.0084, val kl loss: 0.3961\n",
      "Epoch 8 (0), 25.55 sec >> train loss: 0.0122, train recon loss: 0.0085, train kl loss: 0.3698, val loss: 0.0118, val recon loss: 0.0083, val kl loss: 0.3529\n",
      "INFO:tensorflow:save_vae/model.ckpt-8 is not in all_model_checkpoint_paths. Manually adding it.\n",
      "Epoch 9 (0), 25.22 sec >> train loss: 0.0117, train recon loss: 0.0082, train kl loss: 0.3458, val loss: 0.0111, val recon loss: 0.0077, val kl loss: 0.3388\n",
      "Epoch 10 (0), 25.2 sec >> train loss: 0.0111, train recon loss: 0.0079, train kl loss: 0.3212, val loss: 0.0105, val recon loss: 0.0073, val kl loss: 0.3154\n",
      "Epoch 11 (0), 28.45 sec >> train loss: 0.0106, train recon loss: 0.0076, train kl loss: 0.2997, val loss: 0.0101, val recon loss: 0.0072, val kl loss: 0.2874\n",
      "Epoch 12 (0), 25.68 sec >> train loss: 0.0101, train recon loss: 0.0073, train kl loss: 0.2810, val loss: 0.0098, val recon loss: 0.0071, val kl loss: 0.2700\n",
      "INFO:tensorflow:save_vae/model.ckpt-12 is not in all_model_checkpoint_paths. Manually adding it.\n",
      "Epoch 13 (0), 25.51 sec >> train loss: 0.0097, train recon loss: 0.0071, train kl loss: 0.2686, val loss: 0.0095, val recon loss: 0.0068, val kl loss: 0.2672\n",
      "Epoch 14 (0), 25.51 sec >> train loss: 0.0095, train recon loss: 0.0069, train kl loss: 0.2623, val loss: 0.0092, val recon loss: 0.0067, val kl loss: 0.2583\n",
      "Epoch 15 (0), 24.96 sec >> train loss: 0.0093, train recon loss: 0.0067, train kl loss: 0.2590, val loss: 0.0089, val recon loss: 0.0062, val kl loss: 0.2616\n",
      "Epoch 16 (0), 25.49 sec >> train loss: 0.0091, train recon loss: 0.0065, train kl loss: 0.2571, val loss: 0.0090, val recon loss: 0.0064, val kl loss: 0.2632\n",
      "INFO:tensorflow:save_vae/model.ckpt-16 is not in all_model_checkpoint_paths. Manually adding it.\n",
      "Epoch 17 (0), 25.18 sec >> train loss: 0.0087, train recon loss: 0.0062, train kl loss: 0.2530, val loss: 0.0084, val recon loss: 0.0059, val kl loss: 0.2517\n",
      "Epoch 18 (0), 25.57 sec >> train loss: 0.0087, train recon loss: 0.0061, train kl loss: 0.2539, val loss: 0.0083, val recon loss: 0.0058, val kl loss: 0.2559\n",
      "Epoch 19 (0), 25.14 sec >> train loss: 0.0086, train recon loss: 0.0061, train kl loss: 0.2544, val loss: 0.0083, val recon loss: 0.0057, val kl loss: 0.2538\n",
      "Epoch 20 (0), 25.26 sec >> train loss: 0.0085, train recon loss: 0.0060, train kl loss: 0.2546, val loss: 0.0082, val recon loss: 0.0057, val kl loss: 0.2575\n",
      "INFO:tensorflow:save_vae/model.ckpt-20 is not in all_model_checkpoint_paths. Manually adding it.\n",
      "Epoch 21 (0), 27.29 sec >> train loss: 0.0085, train recon loss: 0.0059, train kl loss: 0.2551, val loss: 0.0083, val recon loss: 0.0058, val kl loss: 0.2545\n",
      "Epoch 22 (0), 25.67 sec >> train loss: 0.0085, train recon loss: 0.0059, train kl loss: 0.2555, val loss: 0.0082, val recon loss: 0.0057, val kl loss: 0.2538\n",
      "Epoch 23 (0), 25.18 sec >> train loss: 0.0084, train recon loss: 0.0059, train kl loss: 0.2558, val loss: 0.0083, val recon loss: 0.0057, val kl loss: 0.2598\n",
      "Epoch 24 (0), 25.29 sec >> train loss: 0.0084, train recon loss: 0.0058, train kl loss: 0.2559, val loss: 0.0082, val recon loss: 0.0057, val kl loss: 0.2548\n",
      "INFO:tensorflow:save_vae/model.ckpt-24 is not in all_model_checkpoint_paths. Manually adding it.\n",
      "Epoch 25 (0), 25.52 sec >> train loss: 0.0083, train recon loss: 0.0058, train kl loss: 0.2561, val loss: 0.0082, val recon loss: 0.0056, val kl loss: 0.2586\n",
      "Epoch 26 (0), 25.42 sec >> train loss: 0.0083, train recon loss: 0.0058, train kl loss: 0.2564, val loss: 0.0081, val recon loss: 0.0055, val kl loss: 0.2593\n",
      "Epoch 27 (0), 25.2 sec >> train loss: 0.0083, train recon loss: 0.0057, train kl loss: 0.2566, val loss: 0.0081, val recon loss: 0.0055, val kl loss: 0.2557\n",
      "Epoch 28 (0), 24.97 sec >> train loss: 0.0083, train recon loss: 0.0057, train kl loss: 0.2568, val loss: 0.0081, val recon loss: 0.0055, val kl loss: 0.2547\n",
      "INFO:tensorflow:save_vae/model.ckpt-28 is not in all_model_checkpoint_paths. Manually adding it.\n",
      "Epoch 29 (0), 25.07 sec >> train loss: 0.0082, train recon loss: 0.0057, train kl loss: 0.2571, val loss: 0.0081, val recon loss: 0.0056, val kl loss: 0.2588\n",
      "Epoch 30 (0), 25.71 sec >> train loss: 0.0082, train recon loss: 0.0056, train kl loss: 0.2571, val loss: 0.0080, val recon loss: 0.0054, val kl loss: 0.2593\n",
      "Epoch 31 (0), 28.39 sec >> train loss: 0.0082, train recon loss: 0.0056, train kl loss: 0.2573, val loss: 0.0080, val recon loss: 0.0054, val kl loss: 0.2598\n",
      "Epoch 32 (0), 26.01 sec >> train loss: 0.0082, train recon loss: 0.0056, train kl loss: 0.2573, val loss: 0.0081, val recon loss: 0.0054, val kl loss: 0.2648\n",
      "INFO:tensorflow:save_vae/model.ckpt-32 is not in all_model_checkpoint_paths. Manually adding it.\n",
      "Epoch 33 (0), 25.76 sec >> train loss: 0.0080, train recon loss: 0.0055, train kl loss: 0.2567, val loss: 0.0079, val recon loss: 0.0053, val kl loss: 0.2548\n",
      "Epoch 34 (0), 25.92 sec >> train loss: 0.0080, train recon loss: 0.0055, train kl loss: 0.2568, val loss: 0.0079, val recon loss: 0.0054, val kl loss: 0.2546\n",
      "Epoch 35 (0), 26.17 sec >> train loss: 0.0080, train recon loss: 0.0054, train kl loss: 0.2569, val loss: 0.0078, val recon loss: 0.0052, val kl loss: 0.2572\n",
      "Epoch 36 (0), 25.81 sec >> train loss: 0.0080, train recon loss: 0.0054, train kl loss: 0.2572, val loss: 0.0079, val recon loss: 0.0053, val kl loss: 0.2589\n",
      "INFO:tensorflow:save_vae/model.ckpt-36 is not in all_model_checkpoint_paths. Manually adding it.\n",
      "Epoch 37 (0), 26.13 sec >> train loss: 0.0080, train recon loss: 0.0054, train kl loss: 0.2573, val loss: 0.0078, val recon loss: 0.0052, val kl loss: 0.2581\n",
      "Epoch 38 (0), 26.17 sec >> train loss: 0.0080, train recon loss: 0.0054, train kl loss: 0.2574, val loss: 0.0078, val recon loss: 0.0053, val kl loss: 0.2561\n",
      "Epoch 39 (0), 25.86 sec >> train loss: 0.0080, train recon loss: 0.0054, train kl loss: 0.2574, val loss: 0.0078, val recon loss: 0.0052, val kl loss: 0.2612\n",
      "Epoch 40 (0), 25.92 sec >> train loss: 0.0080, train recon loss: 0.0054, train kl loss: 0.2573, val loss: 0.0078, val recon loss: 0.0052, val kl loss: 0.2578\n",
      "INFO:tensorflow:save_vae/model.ckpt-40 is not in all_model_checkpoint_paths. Manually adding it.\n",
      "Epoch 41 (0), 28.62 sec >> train loss: 0.0080, train recon loss: 0.0054, train kl loss: 0.2575, val loss: 0.0079, val recon loss: 0.0053, val kl loss: 0.2571\n",
      "Epoch 42 (0), 26.2 sec >> train loss: 0.0079, train recon loss: 0.0054, train kl loss: 0.2575, val loss: 0.0078, val recon loss: 0.0052, val kl loss: 0.2591\n",
      "Epoch 43 (0), 25.68 sec >> train loss: 0.0079, train recon loss: 0.0054, train kl loss: 0.2575, val loss: 0.0078, val recon loss: 0.0052, val kl loss: 0.2567\n",
      "Epoch 44 (0), 25.75 sec >> train loss: 0.0079, train recon loss: 0.0053, train kl loss: 0.2576, val loss: 0.0078, val recon loss: 0.0052, val kl loss: 0.2594\n",
      "INFO:tensorflow:save_vae/model.ckpt-44 is not in all_model_checkpoint_paths. Manually adding it.\n",
      "Epoch 45 (0), 25.96 sec >> train loss: 0.0079, train recon loss: 0.0053, train kl loss: 0.2578, val loss: 0.0078, val recon loss: 0.0053, val kl loss: 0.2506\n",
      "Epoch 46 (0), 25.47 sec >> train loss: 0.0079, train recon loss: 0.0053, train kl loss: 0.2579, val loss: 0.0078, val recon loss: 0.0052, val kl loss: 0.2575\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 47 (0), 26.15 sec >> train loss: 0.0079, train recon loss: 0.0053, train kl loss: 0.2581, val loss: 0.0078, val recon loss: 0.0052, val kl loss: 0.2596\n",
      "Epoch 48 (0), 26.23 sec >> train loss: 0.0079, train recon loss: 0.0053, train kl loss: 0.2579, val loss: 0.0078, val recon loss: 0.0052, val kl loss: 0.2579\n",
      "INFO:tensorflow:save_vae/model.ckpt-48 is not in all_model_checkpoint_paths. Manually adding it.\n",
      "Epoch 49 (0), 25.7 sec >> train loss: 0.0078, train recon loss: 0.0052, train kl loss: 0.2576, val loss: 0.0077, val recon loss: 0.0051, val kl loss: 0.2562\n",
      "Epoch 50 (0), 25.14 sec >> train loss: 0.0078, train recon loss: 0.0053, train kl loss: 0.2578, val loss: 0.0077, val recon loss: 0.0051, val kl loss: 0.2591\n",
      "Epoch 51 (0), 27.28 sec >> train loss: 0.0078, train recon loss: 0.0052, train kl loss: 0.2581, val loss: 0.0077, val recon loss: 0.0051, val kl loss: 0.2585\n",
      "Epoch 52 (0), 25.63 sec >> train loss: 0.0078, train recon loss: 0.0052, train kl loss: 0.2581, val loss: 0.0077, val recon loss: 0.0051, val kl loss: 0.2598\n",
      "INFO:tensorflow:save_vae/model.ckpt-52 is not in all_model_checkpoint_paths. Manually adding it.\n",
      "Epoch 53 (0), 25.03 sec >> train loss: 0.0078, train recon loss: 0.0052, train kl loss: 0.2580, val loss: 0.0077, val recon loss: 0.0051, val kl loss: 0.2570\n",
      "Epoch 54 (0), 25.27 sec >> train loss: 0.0078, train recon loss: 0.0052, train kl loss: 0.2580, val loss: 0.0077, val recon loss: 0.0051, val kl loss: 0.2565\n",
      "Epoch 55 (0), 25.24 sec >> train loss: 0.0078, train recon loss: 0.0052, train kl loss: 0.2580, val loss: 0.0077, val recon loss: 0.0051, val kl loss: 0.2572\n",
      "Epoch 56 (0), 25.76 sec >> train loss: 0.0078, train recon loss: 0.0052, train kl loss: 0.2578, val loss: 0.0077, val recon loss: 0.0051, val kl loss: 0.2587\n",
      "INFO:tensorflow:save_vae/model.ckpt-56 is not in all_model_checkpoint_paths. Manually adding it.\n",
      "Epoch 57 (0), 25.39 sec >> train loss: 0.0078, train recon loss: 0.0052, train kl loss: 0.2579, val loss: 0.0077, val recon loss: 0.0051, val kl loss: 0.2576\n",
      "Epoch 58 (0), 25.98 sec >> train loss: 0.0078, train recon loss: 0.0052, train kl loss: 0.2581, val loss: 0.0077, val recon loss: 0.0051, val kl loss: 0.2553\n",
      "Epoch 59 (0), 25.86 sec >> train loss: 0.0078, train recon loss: 0.0052, train kl loss: 0.2581, val loss: 0.0077, val recon loss: 0.0051, val kl loss: 0.2573\n",
      "Epoch 60 (0), 25.78 sec >> train loss: 0.0078, train recon loss: 0.0052, train kl loss: 0.2579, val loss: 0.0077, val recon loss: 0.0051, val kl loss: 0.2579\n",
      "INFO:tensorflow:save_vae/model.ckpt-60 is not in all_model_checkpoint_paths. Manually adding it.\n",
      "Epoch 61 (0), 27.88 sec >> train loss: 0.0078, train recon loss: 0.0052, train kl loss: 0.2582, val loss: 0.0077, val recon loss: 0.0051, val kl loss: 0.2571\n",
      "Epoch 62 (0), 25.44 sec >> train loss: 0.0078, train recon loss: 0.0052, train kl loss: 0.2582, val loss: 0.0077, val recon loss: 0.0051, val kl loss: 0.2573\n",
      "Epoch 63 (0), 24.79 sec >> train loss: 0.0078, train recon loss: 0.0052, train kl loss: 0.2581, val loss: 0.0077, val recon loss: 0.0051, val kl loss: 0.2602\n",
      "Epoch 64 (0), 25.05 sec >> train loss: 0.0078, train recon loss: 0.0052, train kl loss: 0.2583, val loss: 0.0077, val recon loss: 0.0051, val kl loss: 0.2570\n",
      "INFO:tensorflow:save_vae/model.ckpt-64 is not in all_model_checkpoint_paths. Manually adding it.\n",
      "Epoch 65 (0), 25.46 sec >> train loss: 0.0077, train recon loss: 0.0052, train kl loss: 0.2580, val loss: 0.0076, val recon loss: 0.0051, val kl loss: 0.2561\n",
      "Epoch 66 (0), 25.28 sec >> train loss: 0.0077, train recon loss: 0.0051, train kl loss: 0.2582, val loss: 0.0076, val recon loss: 0.0050, val kl loss: 0.2587\n",
      "Epoch 67 (0), 25.11 sec >> train loss: 0.0077, train recon loss: 0.0051, train kl loss: 0.2582, val loss: 0.0076, val recon loss: 0.0051, val kl loss: 0.2565\n",
      "Epoch 68 (0), 25.62 sec >> train loss: 0.0077, train recon loss: 0.0051, train kl loss: 0.2580, val loss: 0.0076, val recon loss: 0.0051, val kl loss: 0.2565\n",
      "INFO:tensorflow:save_vae/model.ckpt-68 is not in all_model_checkpoint_paths. Manually adding it.\n",
      "Epoch 69 (0), 25.65 sec >> train loss: 0.0077, train recon loss: 0.0051, train kl loss: 0.2583, val loss: 0.0076, val recon loss: 0.0050, val kl loss: 0.2581\n",
      "Epoch 70 (0), 25.42 sec >> train loss: 0.0077, train recon loss: 0.0051, train kl loss: 0.2584, val loss: 0.0076, val recon loss: 0.0050, val kl loss: 0.2579\n",
      "Epoch 71 (0), 27.57 sec >> train loss: 0.0077, train recon loss: 0.0051, train kl loss: 0.2581, val loss: 0.0076, val recon loss: 0.0050, val kl loss: 0.2587\n",
      "Epoch 72 (0), 25.37 sec >> train loss: 0.0077, train recon loss: 0.0051, train kl loss: 0.2584, val loss: 0.0076, val recon loss: 0.0051, val kl loss: 0.2566\n",
      "INFO:tensorflow:save_vae/model.ckpt-72 is not in all_model_checkpoint_paths. Manually adding it.\n",
      "Epoch 73 (0), 25.97 sec >> train loss: 0.0077, train recon loss: 0.0051, train kl loss: 0.2582, val loss: 0.0076, val recon loss: 0.0050, val kl loss: 0.2573\n",
      "Epoch 74 (0), 25.46 sec >> train loss: 0.0077, train recon loss: 0.0051, train kl loss: 0.2582, val loss: 0.0076, val recon loss: 0.0051, val kl loss: 0.2577\n",
      "Epoch 75 (0), 25.92 sec >> train loss: 0.0077, train recon loss: 0.0051, train kl loss: 0.2584, val loss: 0.0076, val recon loss: 0.0050, val kl loss: 0.2584\n",
      "Epoch 76 (0), 25.46 sec >> train loss: 0.0077, train recon loss: 0.0051, train kl loss: 0.2583, val loss: 0.0076, val recon loss: 0.0050, val kl loss: 0.2572\n",
      "INFO:tensorflow:save_vae/model.ckpt-76 is not in all_model_checkpoint_paths. Manually adding it.\n",
      "Epoch 77 (0), 25.7 sec >> train loss: 0.0077, train recon loss: 0.0051, train kl loss: 0.2585, val loss: 0.0076, val recon loss: 0.0050, val kl loss: 0.2580\n",
      "Epoch 78 (0), 25.85 sec >> train loss: 0.0077, train recon loss: 0.0051, train kl loss: 0.2582, val loss: 0.0076, val recon loss: 0.0050, val kl loss: 0.2582\n",
      "Epoch 79 (0), 25.47 sec >> train loss: 0.0077, train recon loss: 0.0051, train kl loss: 0.2583, val loss: 0.0076, val recon loss: 0.0051, val kl loss: 0.2562\n",
      "Epoch 80 (0), 24.84 sec >> train loss: 0.0077, train recon loss: 0.0051, train kl loss: 0.2583, val loss: 0.0076, val recon loss: 0.0051, val kl loss: 0.2571\n",
      "INFO:tensorflow:save_vae/model.ckpt-80 is not in all_model_checkpoint_paths. Manually adding it.\n",
      "Epoch 81 (0), 27.28 sec >> train loss: 0.0077, train recon loss: 0.0051, train kl loss: 0.2583, val loss: 0.0076, val recon loss: 0.0050, val kl loss: 0.2581\n",
      "Epoch 82 (0), 25.25 sec >> train loss: 0.0077, train recon loss: 0.0051, train kl loss: 0.2582, val loss: 0.0076, val recon loss: 0.0050, val kl loss: 0.2584\n",
      "Epoch 83 (0), 25.68 sec >> train loss: 0.0077, train recon loss: 0.0051, train kl loss: 0.2584, val loss: 0.0076, val recon loss: 0.0050, val kl loss: 0.2579\n",
      "Epoch 84 (0), 24.86 sec >> train loss: 0.0077, train recon loss: 0.0051, train kl loss: 0.2584, val loss: 0.0076, val recon loss: 0.0050, val kl loss: 0.2580\n",
      "INFO:tensorflow:save_vae/model.ckpt-84 is not in all_model_checkpoint_paths. Manually adding it.\n",
      "Epoch 85 (0), 25.09 sec >> train loss: 0.0077, train recon loss: 0.0051, train kl loss: 0.2583, val loss: 0.0076, val recon loss: 0.0051, val kl loss: 0.2559\n",
      "Epoch 86 (0), 25.34 sec >> train loss: 0.0077, train recon loss: 0.0051, train kl loss: 0.2584, val loss: 0.0076, val recon loss: 0.0050, val kl loss: 0.2574\n",
      "Epoch 87 (0), 25.12 sec >> train loss: 0.0077, train recon loss: 0.0051, train kl loss: 0.2581, val loss: 0.0076, val recon loss: 0.0050, val kl loss: 0.2577\n",
      "Epoch 88 (0), 26.2 sec >> train loss: 0.0077, train recon loss: 0.0051, train kl loss: 0.2583, val loss: 0.0076, val recon loss: 0.0050, val kl loss: 0.2569\n",
      "INFO:tensorflow:save_vae/model.ckpt-88 is not in all_model_checkpoint_paths. Manually adding it.\n",
      "Epoch 89 (0), 25.62 sec >> train loss: 0.0077, train recon loss: 0.0051, train kl loss: 0.2585, val loss: 0.0076, val recon loss: 0.0050, val kl loss: 0.2572\n",
      "Epoch 90 (0), 25.77 sec >> train loss: 0.0077, train recon loss: 0.0051, train kl loss: 0.2583, val loss: 0.0076, val recon loss: 0.0050, val kl loss: 0.2582\n",
      "Epoch 91 (0), 28.58 sec >> train loss: 0.0077, train recon loss: 0.0051, train kl loss: 0.2585, val loss: 0.0076, val recon loss: 0.0050, val kl loss: 0.2580\n",
      "Epoch 92 (0), 26.5 sec >> train loss: 0.0077, train recon loss: 0.0051, train kl loss: 0.2583, val loss: 0.0076, val recon loss: 0.0050, val kl loss: 0.2574\n",
      "INFO:tensorflow:save_vae/model.ckpt-92 is not in all_model_checkpoint_paths. Manually adding it.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 93 (0), 25.38 sec >> train loss: 0.0077, train recon loss: 0.0051, train kl loss: 0.2584, val loss: 0.0076, val recon loss: 0.0050, val kl loss: 0.2578\n",
      "Epoch 94 (0), 25.43 sec >> train loss: 0.0077, train recon loss: 0.0051, train kl loss: 0.2583, val loss: 0.0076, val recon loss: 0.0050, val kl loss: 0.2577\n",
      "Epoch 95 (0), 26.13 sec >> train loss: 0.0077, train recon loss: 0.0051, train kl loss: 0.2585, val loss: 0.0076, val recon loss: 0.0050, val kl loss: 0.2578\n",
      "Epoch 96 (0), 26.13 sec >> train loss: 0.0077, train recon loss: 0.0051, train kl loss: 0.2584, val loss: 0.0076, val recon loss: 0.0050, val kl loss: 0.2587\n",
      "INFO:tensorflow:save_vae/model.ckpt-96 is not in all_model_checkpoint_paths. Manually adding it.\n",
      "Epoch 97 (0), 26.11 sec >> train loss: 0.0077, train recon loss: 0.0051, train kl loss: 0.2585, val loss: 0.0076, val recon loss: 0.0050, val kl loss: 0.2577\n",
      "Epoch 98 (0), 25.33 sec >> train loss: 0.0077, train recon loss: 0.0051, train kl loss: 0.2584, val loss: 0.0076, val recon loss: 0.0050, val kl loss: 0.2590\n",
      "Epoch 99 (0), 25.56 sec >> train loss: 0.0077, train recon loss: 0.0051, train kl loss: 0.2584, val loss: 0.0076, val recon loss: 0.0050, val kl loss: 0.2579\n",
      "Epoch 100 (0), 25.45 sec >> train loss: 0.0077, train recon loss: 0.0051, train kl loss: 0.2583, val loss: 0.0076, val recon loss: 0.0050, val kl loss: 0.2573\n",
      "INFO:tensorflow:save_vae/model.ckpt-100 is not in all_model_checkpoint_paths. Manually adding it.\n"
     ]
    }
   ],
   "source": [
    "# %load train.py\n",
    "import os\n",
    "import time\n",
    "import argparse\n",
    "import numpy as np\n",
    "import tensorflow as tf\n",
    "\n",
    "from progress.bar import Bar\n",
    "from ipywidgets import IntProgress\n",
    "from IPython.display import display\n",
    "import skimage.transform\n",
    "import imageio\n",
    "\n",
    "from utils import read_dataset\n",
    "\n",
    "TRAIN_CSV = \"hw4_data/train.csv\"\n",
    "TRAIN_DIR = \"hw4_data/train/\"\n",
    "TEST_CSV = \"hw4_data/test.csv\"\n",
    "TEST_DIR = \"hw4_data/test/\"\n",
    "\n",
    "FLAG_lr = 1e-3\n",
    "FLAG_save_dir = 'save_vae/'\n",
    "FLAG_lambda_kl = 1e-2\n",
    "FLAG_batch_size = 32\n",
    "FLAG_n_dim = 512\n",
    "\n",
    "if not os.path.exists(FLAG_save_dir):\n",
    "    print(\"Creating %s\" % FLAG_save_dir)\n",
    "    os.makedirs(FLAG_save_dir)\n",
    "if not os.path.exists(os.path.join(FLAG_save_dir, \"recons\")):\n",
    "    print(\"Creating %s\" % os.path.join(FLAG_save_dir, \"recons\"))\n",
    "    os.makedirs(os.path.join(FLAG_save_dir, \"recons\"))\n",
    "if not os.path.exists(os.path.join(FLAG_save_dir, \"samples\")):\n",
    "    print(\"Creating %s\" % os.path.join(FLAG_save_dir, \"samples\"))\n",
    "    os.makedirs(os.path.join(FLAG_save_dir, \"samples\"))\n",
    "\n",
    "print(\"Reading dataset...\")\n",
    "# load data\n",
    "Xtrain, df_train = read_dataset(TRAIN_CSV, TRAIN_DIR)\n",
    "Xtest , df_test  = read_dataset(TEST_CSV , TEST_DIR)\n",
    "\n",
    "vae = VAE()\n",
    "vae.build(lambda_kl=FLAG_lambda_kl,n_dim=FLAG_n_dim, shape=Xtrain.shape[1:])\n",
    "\n",
    "saver = tf.train.Saver(tf.global_variables(), max_to_keep=None)\n",
    "checkpoint_path = os.path.join(FLAG_save_dir, 'model.ckpt')\n",
    "\n",
    "def initialize_uninitialized(sess):\n",
    "    global_vars = tf.global_variables()\n",
    "    is_not_initialized = sess.run([tf.is_variable_initialized(var) for var in global_vars])\n",
    "    not_initialized_vars = [v for (v,f) in zip(global_vars, is_not_initialized) if not f]\n",
    "    if len(not_initialized_vars): \n",
    "            sess.run(tf.variables_initializer(not_initialized_vars))\n",
    "\n",
    "def res_plot(samples, n_row, n_col):     \n",
    "    fig = plt.figure(figsize=(n_col*2, n_row*2))\n",
    "    gs = gridspec.GridSpec(n_row, n_col)\n",
    "    gs.update(wspace=0.05, hspace=0.05)\n",
    "    for i, sample in enumerate(samples):\n",
    "        ax = plt.subplot(gs[i])\n",
    "        plt.axis('off')\n",
    "        ax.set_xticklabels([])\n",
    "        ax.set_yticklabels([])\n",
    "        ax.set_aspect('equal')\n",
    "        plt.imshow(sample.reshape(64, 64, 3))\n",
    "    return fig\n",
    "\n",
    "with tf.Session() as sess:\n",
    "    sess.run(tf.global_variables_initializer())\n",
    "\n",
    "    # hyper parameters\n",
    "    batch_size = FLAG_batch_size\n",
    "    epoch = 100\n",
    "    early_stop_patience = 10\n",
    "    min_delta = 0.0001\n",
    "\n",
    "    # recorder\n",
    "    epoch_counter = 0\n",
    "\n",
    "    # optimizer\n",
    "    global_step = tf.Variable(0, trainable=False)\n",
    "\n",
    "    # Passing global_step to minimize() will increment it at each step.\n",
    "    start_learning_rate = FLAG_lr\n",
    "    half_cycle = 20000\n",
    "    learning_rate = tf.train.exponential_decay(start_learning_rate, global_step, half_cycle, 0.5, staircase=True)\n",
    "    opt = tf.train.AdamOptimizer(learning_rate=learning_rate)\n",
    "\n",
    "    obj = vae.vae_loss\n",
    "    train_op = opt.minimize(obj, global_step=global_step)\n",
    "\n",
    "    # progress bar\n",
    "    ptrain = IntProgress()\n",
    "    pval = IntProgress()\n",
    "    display(ptrain)\n",
    "    display(pval)\n",
    "    ptrain.max = int(Xtrain.shape[0]/batch_size)\n",
    "    pval.max = int(Xtest.shape[0]/batch_size)\n",
    "\n",
    "    # re-initialize\n",
    "    initialize_uninitialized(sess)\n",
    "\n",
    "    # reset due to adding a new task\n",
    "    patience_counter = 0\n",
    "    current_best_val_loss = np.float('Inf')\n",
    "    \n",
    "    loss_dict = dict({'train_kl':[], \n",
    "                      'train_recon':[],\n",
    "                      'val_kl':[],\n",
    "                      'val_recon':[]\n",
    "                     })\n",
    "    \n",
    "    # optimize when the aggregated obj\n",
    "    while(patience_counter < early_stop_patience and epoch_counter < epoch):\n",
    "\n",
    "        # start training\n",
    "        stime = time.time()\n",
    "        bar_train = Bar('Training', max=int(Xtrain.shape[0]/batch_size), suffix='%(index)d/%(max)d - %(percent).1f%% - %(eta)ds')\n",
    "        bar_val =  Bar('Validation', max=int(Xtest.shape[0]/batch_size), suffix='%(index)d/%(max)d - %(percent).1f%% - %(eta)ds')\n",
    "\n",
    "        train_loss = 0.0\n",
    "        train_reconstruction_loss = 0.0\n",
    "        train_kl_loss = 0.0\n",
    "        for i in range(int(Xtrain.shape[0]/batch_size)):\n",
    "            st = i*batch_size\n",
    "            ed = (i+1)*batch_size\n",
    "            loss, reconstruction_loss, kl_loss ,_ = sess.run([vae.vae_loss, \n",
    "                                                              vae.recon_loss, \n",
    "                                                              vae.kl_loss, \n",
    "                                                              train_op],\n",
    "                                                             feed_dict={vae.x: Xtrain[st:ed,:],\n",
    "                                                                        vae.is_train: True})\n",
    "            train_loss += np.mean(loss)\n",
    "            train_reconstruction_loss += np.mean(reconstruction_loss)\n",
    "            train_kl_loss += np.mean(kl_loss)\n",
    "            ptrain.value +=1\n",
    "            ptrain.description = \"Training %s/%s\" % (ptrain.value, ptrain.max)\n",
    "\n",
    "        train_loss = train_loss/ptrain.value\n",
    "        train_reconstruction_loss = train_reconstruction_loss/ptrain.value\n",
    "        train_kl_loss = train_kl_loss/ptrain.value\n",
    "        \n",
    "        loss_dict['train_kl'].append(train_kl_loss)\n",
    "        loss_dict['train_recon'].append(train_reconstruction_loss)\n",
    "\n",
    "        # validation\n",
    "        val_loss = 0\n",
    "        val_reconstruction_loss = 0.0\n",
    "        val_kl_loss = 0.0\n",
    "        for i in range(int(Xtest.shape[0]/batch_size)):\n",
    "            st = i*batch_size\n",
    "            ed = (i+1)*batch_size\n",
    "            loss, reconstruction_loss, kl_loss = sess.run([vae.vae_loss, \n",
    "                                                           vae.recon_loss, \n",
    "                                                           vae.kl_loss],\n",
    "                                                          feed_dict={vae.x: Xtest[st:ed,:],\n",
    "                                                                     vae.is_train: False})\n",
    "            val_loss += np.mean(loss)\n",
    "            val_reconstruction_loss += np.mean(reconstruction_loss)\n",
    "            val_kl_loss += np.mean(kl_loss)\n",
    "            pval.value += 1\n",
    "            pval.description = \"Testing %s/%s\" % (pval.value, pval.value)\n",
    "        val_loss = val_loss/pval.value\n",
    "        val_reconstruction_loss = val_reconstruction_loss/pval.value\n",
    "        val_kl_loss = val_kl_loss/pval.value\n",
    "        \n",
    "        loss_dict['val_kl'].append(val_kl_loss)\n",
    "        loss_dict['val_recon'].append(val_reconstruction_loss)\n",
    "\n",
    "        # plot\n",
    "        if epoch_counter%10 == 0:\n",
    "            Xplot = sess.run(vae.output_image,\n",
    "                    feed_dict={vae.x: Xtest[:10,:],\n",
    "                                vae.is_train: False})\n",
    "            fig = res_plot(np.concatenate((Xtest[:10,:], Xplot), axis=0), 2, 10)\n",
    "            plt.savefig(os.path.join(FLAG_save_dir, 'recons', '{}.png'.format(str(epoch_counter).zfill(3))), \n",
    "                        bbox_inches='tight')\n",
    "            plt.close(fig)\n",
    "\n",
    "            #### produce 32 random images\n",
    "            samples = sess.run(vae.random_sample_images, feed_dict={vae.random_sample: np.random.randn(32, vae.n_dim),\n",
    "                                                                   vae.is_train: False})\n",
    "            fig = res_plot(samples, 4, 8)\n",
    "            plt.savefig(os.path.join(FLAG_save_dir,'samples', '{}.png'.format(str(epoch_counter).zfill(3))), \n",
    "                        bbox_inches='tight')\n",
    "            plt.close(fig)\n",
    "\n",
    "        # shuffle Xtrain and Ytrain in the next epoch\n",
    "        idx = np.random.permutation(Xtrain.shape[0])\n",
    "        Xtrain= Xtrain[idx,:,:,:]\n",
    "\n",
    "        # epoch end\n",
    "        epoch_counter += 1\n",
    "\n",
    "        ptrain.value = 0\n",
    "        pval.value = 0\n",
    "        bar_train.finish()\n",
    "        bar_val.finish()\n",
    "\n",
    "        print(\"Epoch %s (%s), %s sec >> train loss: %.4f, train recon loss: %.4f, train kl loss: %.4f, val loss: %.4f, val recon loss: %.4f, val kl loss: %.4f\" % (epoch_counter, patience_counter, round(time.time()-stime,2), train_loss, train_reconstruction_loss, train_kl_loss, val_loss, val_reconstruction_loss ,val_kl_loss))\n",
    "        if epoch_counter%4==0:\n",
    "            saver.save(sess,os.path.join(FLAG_save_dir,'model.ckpt'),global_step=epoch_counter)\n",
    "np.save(os.path.join(FLAG_save_dir, \"history_dict.npy\"), loss_dict)"
   ]
  },
  {
   "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.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}