DCVAE_test.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "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",
    "glorot_uniform = tf.glorot_uniform_initializer()\n",
    "\n",
    "from NN_architectures import *"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "# some constants\n",
    "\n",
    "LEARNING_RATE = 0.0005\n",
    "BETA1 = 0.5\n",
    "BATCH_SIZE = 128\n",
    "EPOCHS = 30\n",
    "SAVE_SAMPLE_PERIOD = 428\n",
    "\n",
    "task='TRAIN'\n",
    "#task='TEST'\n",
    "PATH='DCVAE_test'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "global e_sizes, d_sizes\n",
    "\n",
    "e_sizes = {\n",
    "        'conv_layers': [(2, 4, 2, False, 1, lrelu, glorot_normal),\n",
    "                        (32, 4, 2, True, 1, lrelu, glorot_normal),\n",
    "                         (64, 4, 2, False, 1, lrelu, glorot_normal)],\n",
    "        \n",
    "        'dense_layers': [(1024, False, 1, lrelu, glorot_normal)],\n",
    "        'last_layer_weight_init':glorot_normal,\n",
    "        'z': 200\n",
    "        \n",
    "}\n",
    "    \n",
    "d_sizes = {\n",
    "        'projection': 64,\n",
    "        'bn_after_project': False,\n",
    "        'dense_layers': [(1024, False, 1 , tf.nn.relu, glorot_normal)],\n",
    "        'conv_layers': [(64, 4, 2, False, 1,tf.nn.relu, glorot_normal),\n",
    "                        (32, 4, 2, True, 1, tf.nn.relu, glorot_normal),\n",
    "                        (1, 4, 2, False, 1, lambda x: x, glorot_normal)],\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "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",
    "    X_train = X_train.reshape(X_train.shape[0], 28, 28 ,1)\n",
    "    \n",
    "    X_test = mnist.test.images\n",
    "    X_test = (X_test>0.5).astype(np.float32)\n",
    "    X_test = X_test.reshape(X_test.shape[0], 28, 28 ,1)\n",
    "    \n",
    "    n_H = X_train.shape[1]\n",
    "    n_W = X_train.shape[2]\n",
    "    n_C = X_train.shape[-1]\n",
    "    \n",
    "    \n",
    "    tf.reset_default_graph()\n",
    "    vae = DCVAE(n_H, n_W, n_C, 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, )\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].reshape(n_H,n_W,n_C)\n",
    "            im = vae.posterior_predictive_sample([x]).reshape(n_H,n_W)\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",
    "                                n_W,\n",
    "                                n_H),\n",
    "                                cmap='gray'\n",
    "                      )\n",
    "            plt.title('Original')\n",
    "            plt.axis('off')\n",
    "            \n",
    "            plt.subplot(1,2,2)\n",
    "            plt.imshow(im, cmap='gray')\n",
    "            plt.title('Sampled')\n",
    "            plt.axis('off')\n",
    "            \n",
    "            fig = plt.gcf()\n",
    "            fig.set_size_inches(8, 4)\n",
    "            plt.savefig(PATH+'/sample_%d.png' % i,dpi=150)\n",
    "            plt.show()\n",
    "            \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": 5,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "ename": "KeyboardInterrupt",
     "evalue": "",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mKeyboardInterrupt\u001b[0m                         Traceback (most recent call last)",
      "\u001b[0;32m~/anaconda3/envs/deepnet/lib/python3.6/site-packages/ipykernel/kernelbase.py\u001b[0m in \u001b[0;36m_input_request\u001b[0;34m(self, prompt, ident, parent, password)\u001b[0m\n\u001b[1;32m    726\u001b[0m             \u001b[0;32mtry\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 727\u001b[0;31m                 \u001b[0mident\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mreply\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0msession\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mrecv\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mstdin_socket\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;36m0\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    728\u001b[0m             \u001b[0;32mexcept\u001b[0m \u001b[0mException\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/anaconda3/envs/deepnet/lib/python3.6/site-packages/jupyter_client/session.py\u001b[0m in \u001b[0;36mrecv\u001b[0;34m(self, socket, mode, content, copy)\u001b[0m\n\u001b[1;32m    795\u001b[0m         \u001b[0;32mtry\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 796\u001b[0;31m             \u001b[0mmsg_list\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0msocket\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mrecv_multipart\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mmode\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mcopy\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mcopy\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    797\u001b[0m         \u001b[0;32mexcept\u001b[0m \u001b[0mzmq\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mZMQError\u001b[0m \u001b[0;32mas\u001b[0m \u001b[0me\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/anaconda3/envs/deepnet/lib/python3.6/site-packages/zmq/sugar/socket.py\u001b[0m in \u001b[0;36mrecv_multipart\u001b[0;34m(self, flags, copy, track)\u001b[0m\n\u001b[1;32m    394\u001b[0m         \"\"\"\n\u001b[0;32m--> 395\u001b[0;31m         \u001b[0mparts\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;34m[\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mrecv\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mflags\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mcopy\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mcopy\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mtrack\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mtrack\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    396\u001b[0m         \u001b[0;31m# have first part already, only loop while more to receive\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32mzmq/backend/cython/socket.pyx\u001b[0m in \u001b[0;36mzmq.backend.cython.socket.Socket.recv\u001b[0;34m()\u001b[0m\n",
      "\u001b[0;32mzmq/backend/cython/socket.pyx\u001b[0m in \u001b[0;36mzmq.backend.cython.socket.Socket.recv\u001b[0;34m()\u001b[0m\n",
      "\u001b[0;32mzmq/backend/cython/socket.pyx\u001b[0m in \u001b[0;36mzmq.backend.cython.socket._recv_copy\u001b[0;34m()\u001b[0m\n",
      "\u001b[0;32m~/anaconda3/envs/deepnet/lib/python3.6/site-packages/zmq/backend/cython/checkrc.pxd\u001b[0m in \u001b[0;36mzmq.backend.cython.checkrc._check_rc\u001b[0;34m()\u001b[0m\n",
      "\u001b[0;31mKeyboardInterrupt\u001b[0m: ",
      "\nDuring handling of the above exception, another exception occurred:\n",
      "\u001b[0;31mKeyboardInterrupt\u001b[0m                         Traceback (most recent call last)",
      "\u001b[0;32m<ipython-input-5-3d0e91404e9b>\u001b[0m in \u001b[0;36m<module>\u001b[0;34m()\u001b[0m\n\u001b[1;32m      7\u001b[0m         \u001b[0;32melif\u001b[0m \u001b[0mos\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mpath\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mexists\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mPATH\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m      8\u001b[0m             \u001b[0;32mif\u001b[0m \u001b[0mos\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mpath\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mexists\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mPATH\u001b[0m\u001b[0;34m+\u001b[0m\u001b[0;34m'/checkpoint'\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m----> 9\u001b[0;31m                 \u001b[0mans\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0minput\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m'A previous checkpoint already exists, choose the action to perform \\n \\n 1) Overwrite the current model saved at '\u001b[0m\u001b[0;34m+\u001b[0m\u001b[0mPATH\u001b[0m\u001b[0;34m+\u001b[0m\u001b[0;34m'/checkpoint \\n 2) Start training a new model \\n 3) Restore and continue training the previous model \\n '\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     10\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     11\u001b[0m                 \u001b[0;32mif\u001b[0m \u001b[0mans\u001b[0m \u001b[0;34m==\u001b[0m \u001b[0;34m'1'\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/anaconda3/envs/deepnet/lib/python3.6/site-packages/ipykernel/kernelbase.py\u001b[0m in \u001b[0;36mraw_input\u001b[0;34m(self, prompt)\u001b[0m\n\u001b[1;32m    700\u001b[0m             \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_parent_ident\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    701\u001b[0m             \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_parent_header\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 702\u001b[0;31m             \u001b[0mpassword\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;32mFalse\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    703\u001b[0m         )\n\u001b[1;32m    704\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/anaconda3/envs/deepnet/lib/python3.6/site-packages/ipykernel/kernelbase.py\u001b[0m in \u001b[0;36m_input_request\u001b[0;34m(self, prompt, ident, parent, password)\u001b[0m\n\u001b[1;32m    730\u001b[0m             \u001b[0;32mexcept\u001b[0m \u001b[0mKeyboardInterrupt\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    731\u001b[0m                 \u001b[0;31m# re-raise KeyboardInterrupt, to truncate traceback\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 732\u001b[0;31m                 \u001b[0;32mraise\u001b[0m \u001b[0mKeyboardInterrupt\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    733\u001b[0m             \u001b[0;32melse\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    734\u001b[0m                 \u001b[0;32mbreak\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;31mKeyboardInterrupt\u001b[0m: "
     ]
    }
   ],
   "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": [
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}