DVAE_test.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import tensorflow as tf\n",
    "import matplotlib.pyplot as plt\n",
    "import os\n",
    "import pickle\n",
    "import math\n",
    "\n",
    "trunc_normal= tf.truncated_normal_initializer(stddev=0.02)\n",
    "normal = tf.random_normal_initializer(stddev=0.02)\n",
    "glorot_normal = tf.glorot_normal_initializer()\n",
    "\n",
    "from NN_architectures import *"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "# some constants\n",
    "\n",
    "LEARNING_RATE = 0.001\n",
    "BETA1 = 0.5\n",
    "BATCH_SIZE = 128\n",
    "EPOCHS = 10\n",
    "SAVE_SAMPLE_PERIOD = 428\n",
    "\n",
    "#task='TRAIN'\n",
    "task='TEST'\n",
    "PATH='DVAE_test'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "global e_sizes, d_sizes\n",
    "e_sizes = {'dense_layers':[(1000, False, 1, tf.nn.relu, glorot_normal),\n",
    "                     (750, True, 1, tf.nn.relu, glorot_normal),\n",
    "                     (250, False, 1, tf.nn.relu, glorot_normal)],\n",
    "           'last_layer_weight_init':glorot_normal,\n",
    "           'z':100,\n",
    "}\n",
    "d_sizes = { 'dense_layers':[(250, False, 1, tf.nn.relu, glorot_normal),\n",
    "                     (750, True, 1, tf.nn.relu, glorot_normal),\n",
    "                     (1000, False, 1, tf.nn.relu, glorot_normal),],\n",
    "           'last_layer_weight_init':glorot_normal,\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def mnist():\n",
    "\n",
    "    from tensorflow.examples.tutorials.mnist import input_data\n",
    "    mnist = input_data.read_data_sets(\"MNIST_data/\", one_hot=False)\n",
    "    \n",
    "    X_train = mnist.train.images\n",
    "    X_train = (X_train>0.5).astype(np.float32)\n",
    "    \n",
    "    X_test = mnist.test.images\n",
    "    X_test = (X_test>0.5).astype(np.float32)\n",
    "\n",
    "    dim = X_train.shape[1]\n",
    "    \n",
    "    tf.reset_default_graph()\n",
    "    vae = DVAE(dim, e_sizes, d_sizes,\n",
    "              lr=LEARNING_RATE, beta1=BETA1,\n",
    "              batch_size=BATCH_SIZE, epochs=EPOCHS,\n",
    "              save_sample= SAVE_SAMPLE_PERIOD, path=PATH, img_height=28, img_width=28 )\n",
    "    \n",
    "    vars_to_train= tf.trainable_variables()\n",
    "    \n",
    "    if task == 'TRAIN':\n",
    "        init_op = tf.global_variables_initializer()\n",
    "        \n",
    "    if task == 'TEST':\n",
    "        vars_all = tf.global_variables()\n",
    "        vars_to_init = list(set(vars_all)-set(vars_to_train))\n",
    "        init_op = tf.variables_initializer(vars_to_init)\n",
    "    \n",
    "    # Add ops to save and restore all the variables.\n",
    "    saver = tf.train.Saver()\n",
    "    \n",
    "    with tf.Session() as sess:\n",
    "        \n",
    "        sess.run(init_op)\n",
    "\n",
    "        if task=='TRAIN':\n",
    "            print('\\n Training...')\n",
    "            \n",
    "            if os.path.exists(PATH+'/'+PATH+'.ckpt.index'):\n",
    "                saver.restore(sess,PATH+'/'+PATH+'.ckpt')\n",
    "                print('Model restored.')\n",
    "            \n",
    "            vae.set_session(sess)\n",
    "            vae.fit(X_train)\n",
    "            \n",
    "            save_path = saver.save(sess, PATH+'/'+PATH+'.ckpt')\n",
    "            print(\"Model saved in path: %s\" % save_path)\n",
    "        \n",
    "        if task=='TEST':\n",
    "            print('\\n Evaluate model on test set...')\n",
    "            saver.restore(sess,PATH+'/'+PATH+'.ckpt')\n",
    "            print('Model restored.')\n",
    "            \n",
    "            vae.set_session(sess)\n",
    "            \n",
    "        done = False\n",
    "        while not done:\n",
    "            \n",
    "            i = np.random.choice(len(X_test))\n",
    "            x = X_test[i]\n",
    "            im = vae.posterior_predictive_sample([x])\n",
    "            #print('True energy deposit: '+str(np.sum(x))+'\\n'+\n",
    "            #      'HCAL reconstructed energy deposit: '+str(np.sum(reco))+'\\n'+\n",
    "            #      'NN Simulated energy deposit: '+str(np.sum(im)))\n",
    "            \n",
    "            plt.subplot(1,2,1)\n",
    "            \n",
    "            plt.imshow(x.reshape(\n",
    "                                28,\n",
    "                                28),\n",
    "                                cmap='gray'\n",
    "                      )\n",
    "            plt.title('Original')\n",
    "            plt.axis('off')\n",
    "            \n",
    "            plt.subplot(1,2,2)\n",
    "            plt.imshow(im.reshape(28,28), cmap='gray')\n",
    "            plt.title('Sampled')\n",
    "            plt.axis('off')\n",
    "            \n",
    "            fig = plt.gcf()\n",
    "            fig.set_size_inches(20, 10)\n",
    "            plt.savefig(PATH+'/sample_%d.png' % i,dpi=100)\n",
    "            plt.show()\n",
    "            \n",
    "\n",
    "            ans = input(\"Generate another?\")\n",
    "            if ans and ans[0] in ('n' or 'N'):\n",
    "                done = True\n",
    "            \n",
    "        done = False\n",
    "        i=0\n",
    "        while not done:\n",
    "            \n",
    "            im = vae.prior_predictive_sample()\n",
    "            #print('True energy deposit: '+str(np.sum(x))+'\\n'+\n",
    "            #      'HCAL reconstructed energy deposit: '+str(np.sum(reco))+'\\n'+\n",
    "            #      'NN Simulated energy deposit: '+str(np.sum(im)))\n",
    "            \n",
    "            \n",
    "\n",
    "            plt.imshow(im.reshape(28,28), cmap='gray')\n",
    "            plt.title('Sampled')\n",
    "            \n",
    "            fig = plt.gcf()\n",
    "            fig.set_size_inches(5, 5)\n",
    "            plt.savefig(PATH+'/sampled_from_dist_%d.png' % i,dpi=100)\n",
    "            plt.show()\n",
    "            i+=1\n",
    "            \n",
    "            ans = input(\"Generate another?\")\n",
    "            if ans and ans[0] in ('n' or 'N'):\n",
    "                done = True\n",
    "            "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Extracting MNIST_data/train-images-idx3-ubyte.gz\n",
      "Extracting MNIST_data/train-labels-idx1-ubyte.gz\n",
      "Extracting MNIST_data/t10k-images-idx3-ubyte.gz\n",
      "Extracting MNIST_data/t10k-labels-idx1-ubyte.gz\n",
      "\n",
      " Evaluate model on test set...\n",
      "INFO:tensorflow:Restoring parameters from DVAE_test/DVAE_test.ckpt\n",
      "Model restored.\n"
     ]
    },
    {
     "data": {
      "image/png": "iVBORw0KGgoAAAANSUhEUgAABIEAAAIqCAYAAACzPVrWAAAABHNCSVQICAgIfAhkiAAAAAlwSFlzAAALEgAACxIB0t1+/AAAADl0RVh0U29mdHdhcmUAbWF0cGxvdGxpYiB2ZXJzaW9uIDIuMS4xLCBodHRwOi8vbWF0cGxvdGxpYi5vcmcvAOZPmwAAGx1JREFUeJzt3X+s3Xd93/H327kOcRyy1DQLFEiLBl5DpoEtdSsR0jrV29pKUaStTEsJFKRp0iZ10kSRFgKlZlOzjWhCYiygMI2uJGRq2kVtl46yTQmDUaRViJppZQTFLAtJnB9zHJwfXN/72R/n0jrBSZz043uu7+vxkCI59x6/zufavsn3Pu/3XvcYowAAAADY3nYs+wAAAAAAnHkiEAAAAEAAEQgAAAAggAgEAAAAEEAEAgAAAAggAgEAAAAEEIGAl6y739fdn5z92NPYGt39+hlbAACJuvuXu/vTm/1zgeUSgYA/1t3v6u5D3f1Edz/Q3Td290XP9fgxxq+MMf7u6Wy/mMcCAGxn3f3W7v7v3f1Ydz/a3V/s7h9b9rmA7U8EAqqqqrvfU1X/vKreW1V/pqp+vKp+uKo+193nnuLxK5t7QgCAs193X1hVv1NVH62qPVX16qo6WFVPL/NcQAYRCPjexcjBqvqFMcZ/GmOsjjEOV9XfrkUIumbjtt/buvvT3X2sqt717FuBu/ud3f2t7n6kuz/Q3Ye7+8DG6/74sd39Ixtf0vXz3f1/uvvh7r7upJ2/1N1f6u6j3X1/d/+rU4UoAICz0N6qqjHGZ8YYa2OMJ8cYvzfG+MPu/nPd/V83rqUe7u6bT74re+Pa6r3d/Yfdfby7/013X9Ldv9vdj3f3f+7uH9h47Peut/5ed39745rqPc91qO7+8Y27k45291e7+ydOet3ruvuujef4XFX94Jn75QHOJBEIqKq6oqrOq6rfPPmFY4zvVNXvVtVf23jRVVV1W1VdVFU3n/zY7n5jVf3rqnp7Vb2qFncTvfoFnvetVfXnq+onq+qXuvuyjZevVdU/qsUFxls2Xv8PXsLbBQCw1fzvqlrr7l/t7p/+XrTZ0FV1fVX9UFVdVlWvrapfftbP/1u1uDbbW1VX1uJa7X21uG7aUVX/8FmP/6tV9Yaq+utV9Y+/9wm6k3X3q6vqP1bVP63F3Um/WFW/0d0Xbzzklqr6g43n+CdV9fMv+q0GtgQRCKha/A/94THGiVO87v76k8/2fGmMcfsYY32M8eSzHvezVfXbY4wvjDG+W1W/VFXjBZ734MZnv75aVV+tqjdVVY0x/mCM8ftjjBMbdyR9oqr+ykt70wAAto4xxrFafCJsVNVNVfVQd/9Wd18yxrh7jPG5McbTY4yHqupf1vdfA310jPHgGOO+qvpvVfXlMcZXxhhPV9V/qKp9z3r8wTHG8THGoar6t1V19SmOdU1V3THGuGPjOu9zVfU/qupnuvvSqvqxqvrAxrk+X1W/PeUXA9h0IhBQVfVwVf3gc3yfn1dtvL6q6t7n2fihk18/xniiqh55ged94KQfP1FVF1RVdffe7v6djW9OfayqfqXcdgwAbBNjjP81xnjXGOM1VfUXanEd9ZHu/rPdfWt337dxDfTp+v5roAdP+vGTp/j3C571+JOv37618VzP9sNV9baNLwU72t1HaxGqXrXx+P83xjj+rB3gLCQCAVVVX6rFNyP8mye/sLt3V9VPV9V/2XjR893Zc39Vveakn7urql7xEs9zY1X9UVW9YYxxYS1uce6XuAUAsGWNMf6oqj5Vixh0fS2ut/7ixjXQNfWnvwZ67Uk/vrSqvn2Kx9xbVb82xrjopH92jzH+WS2u8X5g47rw5B3gLCQCATXGeKwW3xj6o939U929s7t/pKp+var+b1X92mnM3FZVV3b3FRvfxPlgvfSLlpdX1bGq+k53/2hV/f2XuAMAsKV0949293u6+zUb//7aWnyJ1u/X4hroO1V1dOP79Lx3wlN+oLvP7+7Lq+rdVfXvT/GYT9fiOu5vdPc53X1ed/9Ed79mjPGtWnxp2MHuPre731qL70UEnIVEIKCqqsYY/6IWd9zcUIsA8+VafFboJze+xvyFfv7/rKpfqKpba/EZo8er6ki9tL/u9Ber6uc2Nm6qU1+sAACcjR6vqr9cVV/u7uO1iD9fq6r31OKTaPur6rFafKPm33yukRfhrqq6uxZ3dt8wxvi9Zz9gjHFvLf4CkPdV1UO1uAZ8b/3Jx4s/t3HmR6vqg1X17yacC1iCHuOFvm8rwIvX3RdU1dFafEnXPcs+DwBAko27uu+pqp3P8Zd/AIHcCQRM091XbtxuvLsWdxQdqqrDyz0VAAAAVSIQMNdVtfhmg9+uqjdU1d8ZbjcEAADYEnw5GAAAAEAAdwIBAAAABBCBAAAAAAKsbOaTdbevPQOAbW6M0cs+A8/kGgwAtr/TuQZzJxAAAABAABEIAAAAIIAIBAAAABBABAIAAAAIIAIBAAAABBCBAAAAAAKIQAAAAAABRCAAAACAACIQAAAAQAARCAAAACCACAQAAAAQQAQCAAAACCACAQAAAAQQgQAAAAACiEAAAAAAAUQgAAAAgAAiEAAAAEAAEQgAAAAggAgEAAAAEEAEAgAAAAggAgEAAAAEEIEAAAAAAohAAAAAAAFEIAAAAIAAIhAAAABAABEIAAAAIIAIBAAAABBABAIAAAAIIAIBAAAABBCBAAAAAAKIQAAAAAABRCAAAACAACIQAAAAQAARCAAAACDAyrIPAMD2NMZY9hGeobuXfQQAYAvZsWPOPRHnnXfelJ1Z106rq6tTdk6cODFlh63FnUAAAAAAAUQgAAAAgAAiEAAAAEAAEQgAAAAggAgEAAAAEEAEAgAAAAggAgEAAAAEEIEAAAAAAohAAAAAAAFEIAAAAIAAIhAAAABAABEIAAAAIIAIBAAAABBABAIAAAAIIAIBAAAABBCBAAAAAAKIQAAAAAABVpZ9AADmGGMs+whnRHcv+wgAAM9pbW1tys6sa7kTJ05M2WF7cicQAAAAQAARCAAAACCACAQAAAAQQAQCAAAACCACAQAAAAQQgQAAAAACiEAAAAAAAUQgAAAAgAAiEAAAAEAAEQgAAAAggAgEAAAAEEAEAgAAAAggAgEAAAAEEIEAAAAAAohAAAAAAAFEIAAAAIAAIhAAAABAABEIAAAAIMDKsg8AkG6MsewjnBHdvewjAADbzM6dO6dtXX755VN2HnnkkSk7R44cmbIDz8edQAAAAAABRCAAAACAACIQAAAAQAARCAAAACCACAQAAAAQQAQCAAAACCACAQAAAAQQgQAAAAACiEAAAAAAAUQgAAAAgAAiEAAAAEAAEQgAAAAggAgEAAAAEEAEAgAAAAggAgEAAAAEEIEAAAAAAohAAAAAAAFWln0AgLPRGGPZR/g+3b3sIwAAnNLLXvayKTtXXHHFlJ2qqv3790/Zuf3226fsnDhxYsoOPB93AgEAAAAEEIEAAAAAAohAAAAAAAFEIAAAAIAAIhAAAABAABEIAAAAIIAIBAAAABBABAIAAAAIIAIBAAAABBCBAAAAAAKIQAAAAAABRCAAAACAACIQAAAAQAARCAAAACCACAQAAAAQQAQCAAAACCACAQAAAARYWfYBADbTGGPZR/g+3b3sIwAAnNLOnTun7Fx99dVTdvbv3z9lp6rqK1/5ypSde++9d8rO2tralB14Pu4EAgAAAAggAgEAAAAEEIEAAAAAAohAAAAAAAFEIAAAAIAAIhAAAABAABEIAAAAIIAIBAAAABBABAIAAAAIIAIBAAAABBCBAAAAAAKIQAAAAAABRCAAAACAACIQAAAAQAARCAAAACCACAQAAAAQQAQCAAAACNBjjM17su7NezJgW9nM/1adju5e9hFgyxpjeAfZYlyDQZ5Z1ypXXnnllJ0bbrhhys4999wzZaeq6pprrpmy8/DDD0/Z2WrXu5x9TucazJ1AAAAAAAFEIAAAAIAAIhAAAABAABEIAAAAIIAIBAAAABBABAIAAAAIIAIBAAAABBCBAAAAAAKIQAAAAAABRCAAAACAACIQAAAAQAARCAAAACCACAQAAAAQQAQCAAAACCACAQAAAAQQgQAAAAACiEAAAAAAAVaWfQBgextjLPsIAABxrrnmmik7H/vYx6bsPProo1N2rrvuuik7VfPO5HqXs4k7gQAAAAACiEAAAAAAAUQgAAAAgAAiEAAAAEAAEQgAAAAggAgEAAAAEEAEAgAAAAggAgEAAAAEEIEAAAAAAohAAAAAAAFEIAAAAIAAIhAAAABAABEIAAAAIIAIBAAAABBABAIAAAAIIAIBAAAABBCBAAAAAAKsLPsAwNY0xlj2Ec6I7l72EQAAntOb3vSmKTs33XTTlJ319fUpOx/+8Ien7Bw6dGjKTlXV2tratC04W7gTCAAAACCACAQAAAAQQAQCAAAACCACAQAAAAQQgQAAAAACiEAAAAAAAUQgAAAAgAAiEAAAAEAAEQgAAAAggAgEAAAAEEAEAgAAAAggAgEAAAAEEIEAAAAAAohAAAAAAAFEIAAAAIAAIhAAAABAABEIAAAAIECPMTbvybo378kg1Ga+T2+m7l72EYDTNMbwDrvFuAaDM2/37t1Tdr75zW9O2bnkkkum7Lz73e+esvOZz3xmys7q6uqUnaqq9fX1aVuwFZzONZg7gQAAAAACiEAAAAAAAUQgAAAAgAAiEAAAAEAAEQgAAAAggAgEAAAAEEAEAgAAAAggAgEAAAAEEIEAAAAAAohAAAAAAAFEIAAAAIAAIhAAAABAABEIAAAAIIAIBAAAABBABAIAAAAIIAIBAAAABBCBAAAAAAKIQAAAAAABVpZ9AGB76+5lHwEA4Dnt3Llzys6dd945Zefiiy+esvPJT35yys7NN988ZWd1dXXKzkyzrlNn7ayvr0/ZgefjTiAAAACAACIQAAAAQAARCAAAACCACAQAAAAQQAQCAAAACCACAQAAAAQQgQAAAAACiEAAAAAAAUQgAAAAgAAiEAAAAEAAEQgAAAAggAgEAAAAEEAEAgAAAAggAgEAAAAEEIEAAAAAAohAAAAAAAFEIAAAAIAAK8s+ALAwxlj2EZ6hu5d9BACAU5p5nXL99ddP2dm/f/+UnS984QtTdq699topO2tra1N2duyYc//BRRddNGWnquriiy+esnPs2LEpOw899NCUnVm/Z1vt4xPmcCcQAAAAQAARCAAAACCACAQAAAAQQAQCAAAACCACAQAAAAQQgQAAAAACiEAAAAAAAUQgAAAAgAAiEAAAAEAAEQgAAAAggAgEAAAAEEAEAgAAAAggAgEAAAAEEIEAAAAAAohAAAAAAAFEIAAAAIAAIhAAAABAgB5jbN6TdW/ek8Em2Mz3n9PV3cs+whmxFX+tt6vt+meIzTPG8Idoi3ENxlYx6/8xl1122ZSdqqrbbrttys6xY8em7Fx77bVTdg4dOjRlZ8+ePVN2rrrqqik7Bw4cmLJTVXXuuedO2fnsZz87Zefmm2+esvPAAw9M2VldXZ2yw+Y5nWswdwIBAAAABBCBAAAAAAKIQAAAAAABRCAAAACAACIQAAAAQAARCAAAACCACAQAAAAQQAQCAAAACCACAQAAAAQQgQAAAAACiEAAAAAAAUQgAAAAgAAiEAAAAEAAEQgAAAAggAgEAAAAEEAEAgAAAAggAgEAAAAEWFn2AYCtaYyx7CMAANvMnj17pux88IMfnLJTVfX5z39+ys7HP/7xKTtHjx6dsvOOd7xjys7b3/72KTsPPvjglJ3zzz9/yk5V1crKnA+Hn3zyySk7r3jFK6bsHDlyZMoO25M7gQAAAAACiEAAAAAAAUQgAAAAgAAiEAAAAEAAEQgAAAAggAgEAAAAEEAEAgAAAAggAgEAAAAEEIEAAAAAAohAAAAAAAFEIAAAAIAAIhAAAABAABEIAAAAIIAIBAAAABBABAIAAAAIIAIBAAAABBCBAAAAAAKsLPsAsAxjjGUf4YzZrm9bdy/7CGfMrN+z7fxrBMBy7dy5c8rOddddN2Vn165dU3aqqj70oQ9N2XnwwQen7Ozbt2/KzlVXXTVlZ3V1dcrOLbfcMmVnfX19yk5V1YUXXjhl5/7775+yc999903ZWVtbm7LD9uROIAAAAIAAIhAAAABAABEIAAAAIIAIBAAAABBABAIAAAAIIAIBAAAABBCBAAAAAAKIQAAAAAABRCAAAACAACIQAAAAQAARCAAAACCACAQAAAAQQAQCAAAACCACAQAAAAQQgQAAAAACiEAAAAAAAUQgAAAAgAAryz4AsDV197KPsKWNMZZ9BADYNHv37p2ys2/fvik7t99++5Sdqqr77rtvys6ePXum7Bw4cGDKzr333jtl58Ybb5yy841vfGPKztra2pSdqqpLL710ys4FF1wwZefxxx+fsjPz14jtx51AAAAAAAFEIAAAAIAAIhAAAABAABEIAAAAIIAIBAAAABBABAIAAAAIIAIBAAAABBCBAAAAAAKIQAAAAAABRCAAAACAACIQAAAAQAARCAAAACCACAQAAAAQQAQCAAAACCACAQAAAAQQgQAAAAACiEAAAAAAAVaWfQB4McYYyz7Cltfdyz7ClrYV/wz5PQPgTNm1a9eUnbe97W1TdlZW5nz4ceedd07Zqao6ceLElJ3Xv/71U3bOOeecKTu33HLLlJ3Dhw9P2bnkkkum7Kyurk7Zqap685vfPGXnjjvumLLz3e9+d8rOVrzeZetwJxAAAABAABEIAAAAIIAIBAAAABBABAIAAAAIIAIBAAAABBCBAAAAAAKIQAAAAAABRCAAAACAACIQAAAAQAARCAAAACCACAQAAAAQQAQCAAAACCACAQAAAAQQgQAAAAACiEAAAAAAAUQgAAAAgAAiEAAAAEAAEQgAAAAgwMqyDwAsdPeyj3BGjDGWfYQzYrv+fgGwvezcuXPKzvHjx6fsHD58eMrOY489NmWnqmrHjjmfF9+9e/eUna997WtTdp544okpO+985zun7Lz85S+fsnP++edP2amq+shHPjJl55FHHpmys76+PmUHno87gQAAAAACiEAAAAAAAUQgAAAAgAAiEAAAAEAAEQgAAAAggAgEAAAAEEAEAgAAAAggAgEAAAAEEIEAAAAAAohAAAAAAAFEIAAAAIAAIhAAAABAABEIAAAAIIAIBAAAABBABAIAAAAIIAIBAAAABBCBAAAAAAKsLPsAwMIYY9lHiNDdyz4CAGya1dXVKTtHjhyZsrN79+4pOwcPHpyyU1V16623Ttl55StfOWXnjW9845Sdq6++esrO3r17p+zcfffdU3be//73T9mpmvfnen19fcoObAZ3AgEAAAAEEIEAAAAAAohAAAAAAAFEIAAAAIAAIhAAAABAABEIAAAAIIAIBAAAABBABAIAAAAIIAIBAAAABBCBAAAAAAKIQAAAAAABRCAAAACAACIQAAAAQAARCAAAACCACAQAAAAQQAQCAAAACCACAQAAAARYWfYByDDGWPYRWJLuXvYRACDW6urqlJ0vfvGLU3be8pa3TNk5cODAlJ2qqssvv3zKzl133TVlZ9euXVN21tfXp+zceuutU3Y+8YlPTNl54IEHpuxUVT399NNTdnysw9nEnUAAAAAAAUQgAAAAgAAiEAAAAEAAEQgAAAAggAgEAAAAEEAEAgAAAAggAgEAAAAEEIEAAAAAAohAAAAAAAFEIAAAAIAAIhAAAABAABEIAAAAIIAIBAAAABBABAIAAAAIIAIBAAAABBCBAAAAAAKIQAAAAAABeoyxeU/WvXlPxpYy689Zd0/ZmWXm+89We9sAXqoxhv+gbTGuwXLt2DHnc767du2asrNv374pO6973eum7FRVHThwYMrOpz71qSk7X//616fsHD9+fMrOU089NWVndXV1ys7M6+/N/FgYNsPpXIO5EwgAAAAggAgEAAAAEEAEAgAAAAggAgEAAAAEEIEAAAAAAohAAAAAAAFEIAAAAIAAIhAAAABAABEIAAAAIIAIBAAAABBABAIAAAAIIAIBAAAABBCBAAAAAAKIQAAAAAABRCAAAACAACIQAAAAQAARCAAAACBAjzE278m6N+/JAIClGGP0ss/AM7kG40/rnHPOmbJz7rnnTtmZdZ6ZnnrqqSk7a2trU3Y28+M8YGs4nWswdwIBAAAABBCBAAAAAAKIQAAAAAABRCAAAACAACIQAAAAQAARCAAAACCACAQAAAAQQAQCAAAACCACAQAAAAQQgQAAAAACiEAAAAAAAUQgAAAAgAAiEAAAAEAAEQgAAAAggAgEAAAAEEAEAgAAAAggAgEAAAAE6DHG5j1Z9+Y9GQCwFGOMXvYZeCbXYACw/Z3ONZg7gQAAAAACiEAAAAAAAUQgAAAAgAAiEAAAAEAAEQgAAAAggAgEAAAAEEAEAgAAAAggAgEAAAAEEIEAAAAAAohAAAAAAAFEIAAAAIAAIhAAAABAABEIAAAAIIAIBAAAABBABAIAAAAIIAIBAAAABBCBAAAAAAKIQAAAAAABRCAAAACAACIQAAAAQAARCAAAACCACAQAAAAQQAQCAAAACCACAQAAAAQQgQAAAAACiEAAAAAAAUQgAAAAgAAiEAAAAEAAEQgAAAAggAgEAAAAEEAEAgAAAAggAgEAAAAEEIEAAAAAAohAAAAAAAFEIAAAAIAAIhAAAABAABEIAAAAIIAIBAAAABBABAIAAAAIIAIBAAAABBCBAAAAAAKIQAAAAAABRCAAAACAACIQAAAAQAARCAAAACCACAQAAAAQQAQCAAAACCACAQAAAAQQgQAAAAACiEAAAAAAAUQgAAAAgAAiEAAAAEAAEQgAAAAggAgEAAAAEEAEAgAAAAggAgEAAAAEEIEAAAAAAohAAAAAAAFEIAAAAIAAIhAAAABAABEIAAAAIIAIBAAAABBABAIAAAAIIAIBAAAABBCBAAAAAAKIQAAAAAABRCAAAACAACIQAAAAQAARCAAAACCACAQAAAAQQAQCAAAACCACAQAAAAQQgQAAAAACiEAAAAAAAUQgAAAAgAAiEAAAAEAAEQgAAAAggAgEAAAAEEAEAgAAAAggAgEAAAAEEIEAAAAAAohAAAAAAAFEIAAAAIAAIhAAAABAABEIAAAAIIAIBAAAABBABAIAAAAIIAIBAAAABBCBAAAAAAKIQAAAAAABeoyx7DMAAAAAcIa5EwgAAAAggAgEAAAAEEAEAgAAAAggAgEAAAAEEIEAAAAAAohAAAAAAAFEIAAAAIAAIhAAAABAABEIAAAAIIAIBAAAABBABAIAAAAIIAIBAAAABBCBAAAAAAKIQAAAAAABRCAAAACAACIQAAAAQAARCAAAACCACAQAAAAQQAQCAAAACCACAQAAAAQQgQAAAAACiEAAAAAAAUQgAAAAgAD/HzAaneSNgJ6VAAAAAElFTkSuQmCC\n",
      "text/plain": [
       "<matplotlib.figure.Figure at 0x7ffaaf220c18>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "if __name__=='__main__':\n",
    "\n",
    "    if task == 'TRAIN':\n",
    "        if not os.path.exists(PATH):\n",
    "            os.mkdir(PATH)\n",
    "    \n",
    "        elif os.path.exists(PATH):\n",
    "            if os.path.exists(PATH+'/checkpoint'):\n",
    "                ans = input('A previous checkpoint already exists, choose the action to perform \\n \\n 1) Overwrite the current model saved at '+PATH+'/checkpoint \\n 2) Start training a new model \\n 3) Restore and continue training the previous model \\n ')\n",
    "                \n",
    "                if ans == '1':\n",
    "                    print('Overwriting existing model in '+PATH)\n",
    "                    for file in os.listdir(PATH):\n",
    "                        file_path = os.path.join(PATH, file)\n",
    "                        try:\n",
    "                            if os.path.isfile(file_path):\n",
    "                                os.unlink(file_path)\n",
    "                            #elif os.path.isdir(file_path): shutil.rmtree(file_path)\n",
    "                        except Exception as e:\n",
    "                            print(e)\n",
    "                            \n",
    "                elif ans == '2':\n",
    "                    PATH = input('Specify the name of the model, a new directory will be created.\\n')\n",
    "                    os.mkdir(PATH)    \n",
    "        \n",
    "        mnist()\n",
    "   \n",
    "    elif task == 'TEST': \n",
    "        if not os.path.exists(PATH+'/checkpoint'):\n",
    "            print('No checkpoint to test')\n",
    "        else:\n",
    "            mnist()"
   ]
  },
  {
   "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.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}