cb_vae_cifar.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import tensorflow as tf\n",
    "import os\n",
    "import struct\n",
    "from sklearn.neighbors import KNeighborsClassifier\n",
    "from six.moves import cPickle\n",
    "from utils import *\n",
    "import matplotlib.pyplot as plt\n",
    "%matplotlib inline"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "d = 50 #latent dimension\n",
    "lr = 0.001 # learning rate\n",
    "beta_reg = 1.0 # for beta-VAE, 1.0 corresponds to usual VAE\n",
    "gamma = 0.0 # warping parameter\n",
    "N = 50000 # number of train samples in MNIST\n",
    "N_test = 10000 # number of test samples in MNIST\n",
    "batch_size = 50\n",
    "max_epochs = 1 # number of epochs used for training"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# VAE IMPLEMENTATION\n",
    "tf.reset_default_graph()\n",
    "\n",
    "eps = tf.placeholder(tf.float32, shape=(None, d))  # variables to reparametrize to sample from approx posterior\n",
    "z_gen = tf.placeholder(tf.float32, shape=(None, d))  # used only for generative samples (samples from prior)\n",
    "X = tf.placeholder(tf.float32, shape=(None, 3072))\n",
    "labels = tf.placeholder(tf.float32, shape=(None, 10))\n",
    "# the following two placeholders should be obtained by running the graph and generating samples from the model\n",
    "# they are used to compute the inception score\n",
    "class_distr_gen = tf.placeholder(tf.float32, shape=(10))  # this is intended to be the generative label distribution\n",
    "X_gen = tf.placeholder(tf.float32, shape=(None, 3072))  # these are generated samples\n",
    "keep_prob = tf.placeholder(tf.float32, name='keep_prob')\n",
    "\n",
    "# the following lines are only used to have tensorflow versions of the mean and inverse mean functions\n",
    "input_vals = tf.placeholder(tf.float32, shape=(None, 3072))\n",
    "input_means = tf.clip_by_value(cont_bern_mean(input_vals), 1e-4, 1-1e-4)\n",
    "input_lams = tf.clip_by_value(cont_bern_lam(input_vals), 1e-4, 1-1e-4)\n",
    "\n",
    "with tf.variable_scope('ELBO', reuse=tf.AUTO_REUSE):\n",
    "    mu_Z, sigma_Z = encoder_cifar(X, d)\n",
    "Z = mu_Z + sigma_Z * eps  # [batch_size, d]\n",
    "\n",
    "with tf.variable_scope('ELBO', reuse=tf.AUTO_REUSE):\n",
    "    lam = decoder_cifar_cb(Z)\n",
    "    lam = tf.clip_by_value(lam, 1e-4, 1 - 1e-4)\n",
    "    # sample from the model\n",
    "    lam_gen = decoder_cifar_cb(z_gen)\n",
    "    lam_gen = tf.clip_by_value(lam_gen, 1e-4, 1 - 1e-4)\n",
    "\n",
    "log_norm_const = cont_bern_log_norm(lam)\n",
    "log_p_all = tf.reduce_sum(X * tf.log(lam) + (1 - X) * tf.log(1 - lam) + log_norm_const, 1)\n",
    "log_p = tf.reduce_mean(log_p_all)\n",
    "\n",
    "# computing an IW estimate of the log likelihood requires having k epsilon samples per data point in the batch,\n",
    "# so that eps would have to be shaped [None, d, k], which would complicate the rest of the graph.\n",
    "# since trainin with IWAE is not required here, only the log importance weights are computed in the graph and the\n",
    "# log likelihood estimate is computed outside the graph by calling it several times for the same batch but with\n",
    "# different random epsilons.\n",
    "log_iw = log_p_all + tf.reduce_sum(-0.5 * tf.square(Z), 1)\n",
    "log_iw = log_iw + tf.reduce_sum(tf.log(1e-8 + sigma_Z) + tf.square(mu_Z - Z) / (2.0 * tf.square(sigma_Z)), 1)\n",
    "\n",
    "KL = 0.5 * tf.reduce_sum(tf.square(mu_Z) + tf.square(sigma_Z) - tf.log(1e-8 + tf.square(sigma_Z)) - 1.0, 1)\n",
    "KL = tf.reduce_mean(KL)\n",
    "\n",
    "ELBO = log_p - beta_reg * KL\n",
    "cost = - ELBO\n",
    "\n",
    "log_p_all_cheat = tf.reduce_sum(X * tf.log(lam) + (1 - X) * tf.log(1 - lam), 1)\n",
    "log_p_cheat = tf.reduce_mean(log_p_all_cheat)\n",
    "ELBO_cheat = log_p_cheat - beta_reg * KL\n",
    "cost_cheat = - ELBO_cheat\n",
    "log_iw_cheat = log_iw - log_p_all + log_p_all_cheat\n",
    "\n",
    "mean_from_lam = cont_bern_lam(lam)\n",
    "mean_from_lam = tf.clip_by_value(mean_from_lam, 1e-4, 1 - 1e-4)\n",
    "\n",
    "log_norm_const_from_mean = cont_bern_log_norm(mean_from_lam)\n",
    "\n",
    "log_p_all_from_mean = tf.reduce_sum(X * tf.log(mean_from_lam) + (1 - X) * tf.log(1 - mean_from_lam) +\n",
    "                                    log_norm_const_from_mean, 1)\n",
    "log_p_from_mean = tf.reduce_mean(log_p_all_from_mean)\n",
    "ELBO_from_mean = log_p_from_mean - beta_reg *  KL\n",
    "log_iw_from_mean = log_iw - log_p_all + log_p_all_from_mean\n",
    "\n",
    "log_p_all_from_mean_cheat = tf.reduce_sum(X * tf.log(mean_from_lam) + (1 - X) * tf.log(1 - mean_from_lam), 1)\n",
    "log_p_from_mean_cheat = tf.reduce_mean(log_p_all_from_mean_cheat)\n",
    "ELBO_from_mean_cheat = log_p_from_mean - KL\n",
    "log_iw_from_mean_cheat = log_iw - log_p_all + log_p_all_from_mean_cheat\n",
    "\n",
    "with tf.variable_scope('classifier', reuse=tf.AUTO_REUSE):\n",
    "    class_logits = classifier_cifar(X, keep_prob)\n",
    "    class_logits_gen = classifier_cifar(X_gen, keep_prob)\n",
    "acc = 1.0 - tf.reduce_mean(tf.abs(tf.sign(tf.cast(tf.argmax(class_logits, 1) - tf.argmax(labels, 1), tf.float32))))\n",
    "class_loss = tf.reduce_mean(tf.nn.softmax_cross_entropy_with_logits_v2(logits=class_logits, labels=labels))\n",
    "class_probs_gen = tf.clip_by_value(tf.nn.softmax(class_logits_gen), 1e-4, 1 - 1e-4)\n",
    "is_kl = tf.reduce_sum(class_probs_gen * (tf.log(class_probs_gen) - tf.log(class_distr_gen)), axis=1)\n",
    "log_is = tf.reduce_mean(is_kl)\n",
    "\n",
    "optim = tf.train.AdamOptimizer(learning_rate=lr)\n",
    "all_params_ELBO = tf.trainable_variables(scope='ELBO')\n",
    "grads_and_vars = optim.compute_gradients(cost, all_params_ELBO)\n",
    "clipped_grads_and_vars = [(tf.clip_by_value(grad, -1.0, 1.0), var) for grad, var in grads_and_vars]\n",
    "optimizer = optim.apply_gradients(clipped_grads_and_vars)\n",
    "\n",
    "optim_cheat = tf.train.AdamOptimizer(learning_rate=lr)\n",
    "grads_and_vars_cheat = optim_cheat.compute_gradients(cost_cheat, all_params_ELBO)\n",
    "clipped_grads_and_vars_cheat = [(tf.clip_by_value(grad, -1.0, 1.0), var) for grad, var in grads_and_vars_cheat]\n",
    "optimizer_cheat = optim_cheat.apply_gradients(clipped_grads_and_vars_cheat)\n",
    "\n",
    "optim_class = tf.train.AdamOptimizer(learning_rate=lr)\n",
    "all_params_class = tf.trainable_variables(scope='classifier')\n",
    "grads_and_vars_class = optim.compute_gradients(class_loss, all_params_class)\n",
    "clipped_grads_and_vars_class = [(tf.clip_by_value(grad, -1.0, 1.0), var) for grad, var in grads_and_vars_class]\n",
    "optimizer_class = optim_class.apply_gradients(clipped_grads_and_vars_class)\n",
    "\n",
    "init_op = tf.global_variables_initializer()\n",
    "sess = tf.Session()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "# HELPER FUNCTIONS TO EVALUATE METRICS ONCE THE MODEL IS TRAINED\n",
    "\n",
    "def epoch_metrics_cb(N, d, batch_size, X_eval, k=100):\n",
    "    # like epoch_metrics, but also computes the ELBO and log lik estimates after applying the continuous Bernoulli\n",
    "    # inverse mean function to the output of the decoder\n",
    "    assert np.shape(X_eval)[0] % batch_size == 0\n",
    "    ELBO_val = 0.0\n",
    "    ELBO_cheat_val = 0.0\n",
    "    IW_ELBO = 0.0\n",
    "    IW_cheat = 0.0\n",
    "    ELBO_fm_val = 0.0\n",
    "    ELBO_fm_cheat_val = 0.0\n",
    "    IW_fm_ELBO = 0.0\n",
    "    IW_fm_cheat = 0.0\n",
    "    local_perms = PermManager(np.shape(X_eval)[0], batch_size)\n",
    "    while local_perms.epoch < 1:\n",
    "        eps_batch = np.random.normal(size=[batch_size, d])\n",
    "        batch = local_perms.get_indices()\n",
    "        X_batch = np.reshape(X_eval[batch, :], [batch_size, -1])\n",
    "        ELBO_val_batch, ELBO_cheat_val_batch, ELBO_fm_val_batch, ELBO_fm_cheat_val_batch = sess.run([ELBO, ELBO_cheat,\n",
    "                                                                                                     ELBO_from_mean,\n",
    "                                                                                                     ELBO_from_mean_cheat],\n",
    "                                                                                                    {eps: eps_batch,\n",
    "                                                                                                     X: X_batch})\n",
    "        ELBO_val += ELBO_val_batch\n",
    "        ELBO_cheat_val += ELBO_cheat_val_batch\n",
    "        ELBO_fm_val += ELBO_fm_val_batch\n",
    "        ELBO_fm_cheat_val += ELBO_fm_cheat_val_batch\n",
    "        increase_IW_ELBO = []\n",
    "        increase_IW_cheat = []\n",
    "        increase_IW_fm_ELBO = []\n",
    "        increase_IW_fm_cheat = []\n",
    "        for i in xrange(k):\n",
    "            eps_batch = np.random.normal(size=[batch_size, d])\n",
    "            inc_IW_ELBO, inc_IW_cheat, inc_IW_fm_ELBO, inc_IW_fm_cheat = sess.run([log_iw, log_iw_cheat,\n",
    "                                                                                    log_iw_from_mean,\n",
    "                                                                                    log_iw_from_mean_cheat],\n",
    "                                                                                   {eps: eps_batch, X: X_batch})\n",
    "            increase_IW_ELBO.append(inc_IW_ELBO)\n",
    "            increase_IW_cheat.append(inc_IW_cheat)\n",
    "            increase_IW_fm_ELBO.append(inc_IW_fm_ELBO)\n",
    "            increase_IW_fm_cheat.append(inc_IW_fm_cheat)\n",
    "        increase_IW_ELBO = np.array(increase_IW_ELBO)\n",
    "        a_ELBO = np.max(increase_IW_ELBO, axis=0)\n",
    "        IW_ELBO += np.mean(a_ELBO + np.log(np.sum(np.exp(increase_IW_ELBO - a_ELBO), 0)) - np.log(k))\n",
    "        increase_IW_cheat = np.array(increase_IW_cheat)\n",
    "        a_cheat = np.max(increase_IW_cheat, axis=0)\n",
    "        IW_cheat += np.mean(a_cheat + np.log(np.sum(np.exp(increase_IW_cheat - a_cheat), 0)) - np.log(k))\n",
    "        increase_IW_fm_ELBO = np.array(increase_IW_fm_ELBO)\n",
    "        a_fm_ELBO = np.max(increase_IW_fm_ELBO, axis=0)\n",
    "        IW_fm_ELBO += np.mean(a_fm_ELBO + np.log(np.sum(np.exp(increase_IW_fm_ELBO - a_fm_ELBO), 0)) - np.log(k))\n",
    "        increase_IW_fm_cheat = np.array(increase_IW_fm_cheat)\n",
    "        a_fm_cheat = np.max(increase_IW_fm_cheat, axis=0)\n",
    "        IW_fm_cheat += np.mean(a_fm_cheat + np.log(np.sum(np.exp(increase_IW_fm_cheat - a_fm_cheat), 0)) - np.log(k))\n",
    "    ELBO_val = ELBO_val * batch_size / N\n",
    "    ELBO_cheat_val = ELBO_cheat_val * batch_size / N\n",
    "    IW_ELBO = IW_ELBO * batch_size / N\n",
    "    IW_cheat = IW_cheat * batch_size / N\n",
    "    ELBO_fm_val = ELBO_fm_val * batch_size / N\n",
    "    ELBO_fm_cheat_val = ELBO_fm_cheat_val * batch_size / N\n",
    "    IW_fm_ELBO = IW_fm_ELBO * batch_size / N\n",
    "    IW_fm_cheat = IW_fm_cheat * batch_size / N\n",
    "    del local_perms\n",
    "    return ELBO_val, ELBO_cheat_val, IW_ELBO, IW_cheat, ELBO_fm_val, ELBO_fm_cheat_val, IW_fm_ELBO, IW_fm_cheat\n",
    "\n",
    "\n",
    "def compute_IS(batch_size, samples):\n",
    "    # computes the inception score using samples\n",
    "    N_samples = np.shape(samples)[0]\n",
    "    class_distr_val = np.zeros(10)\n",
    "    local_perms = PermManager(N_samples, batch_size)\n",
    "    while local_perms.epoch < 1:\n",
    "        batch = local_perms.get_indices()\n",
    "        samples_batch = np.reshape(samples[batch], [batch_size, -1])\n",
    "        class_distr_val += np.sum(sess.run(class_probs_gen, {X_gen: samples_batch, keep_prob: 1.0}), axis=0)\n",
    "    class_distr_val = class_distr_val / N_samples\n",
    "    log_is_val = 0.0\n",
    "    local_perms = PermManager(N_samples, batch_size)\n",
    "    while local_perms.epoch < 1:\n",
    "        batch = local_perms.get_indices()\n",
    "        samples_batch = np.reshape(samples[batch], [batch_size, -1])\n",
    "        log_is_val += sess.run(log_is, {X_gen: samples_batch, class_distr_gen: class_distr_val, keep_prob:1.0})\n",
    "    log_is_val = log_is_val * batch_size / N\n",
    "    is_val = np.exp(log_is_val)\n",
    "    return is_val\n",
    "\n",
    "\n",
    "def mean_from_lam2(lam_vals):\n",
    "    # tensorflow version of mean function\n",
    "    return sess.run(input_means, {input_vals: lam_vals})\n",
    "\n",
    "\n",
    "def lam_from_means2(batch_size, mean_vals):\n",
    "    # tensorflow version of inverse mean function.\n",
    "    # batches over samples\n",
    "    N_loc = np.shape(mean_vals)[0]\n",
    "    assert N_loc % batch_size == 0\n",
    "    lams = []\n",
    "    local_perms = PermManager(N_loc, batch_size, perm=np.arange(N_loc))\n",
    "    while local_perms.epoch < 1:\n",
    "        batch = local_perms.get_indices()\n",
    "        lams.append(sess.run(input_lams, {input_vals: mean_vals[batch]}))\n",
    "    return np.concatenate(lams)\n",
    "\n",
    "\n",
    "def sample_from_model_cb(batch_size, N):\n",
    "    # gives back N samples (i.e. outputs of decoder) from the model\n",
    "    assert N % batch_size == 0\n",
    "    K = N / batch_size\n",
    "    samples = []\n",
    "    for i in xrange(K):\n",
    "        samples.append(sess.run(lam_gen, {z_gen: np.random.normal(size=[batch_size, d])}))\n",
    "    return np.concatenate(samples)\n",
    "\n",
    "\n",
    "def k_nn_acc(k, batch_size, X_eval, digits_eval, X_train, digits_train):\n",
    "    # computes the knn metric from the model's latent variables\n",
    "    assert np.shape(X_eval)[0] % batch_size == 0\n",
    "    assert np.shape(X_eval)[0] == np.shape(digits_eval)[0]\n",
    "    assert np.shape(X_train)[0] % batch_size == 0\n",
    "    assert np.shape(X_train)[0] == np.shape(digits_train)[0]\n",
    "    local_perms = PermManager(np.shape(X_train)[0], batch_size, perm=np.arange(np.shape(X_train)[0]))\n",
    "    mu_vals_train = []\n",
    "    while local_perms.epoch < 1:\n",
    "        batch = local_perms.get_indices()\n",
    "        X_batch = np.reshape(X_train[batch, :], [batch_size, -1])\n",
    "        mu_vals_train.append(sess.run(mu_Z, {X: X_batch}))\n",
    "    mu_vals_train = np.concatenate(mu_vals_train)\n",
    "    classifier = KNeighborsClassifier(n_neighbors=k)\n",
    "    classifier.fit(mu_vals_train, digits_train)\n",
    "    local_perms = PermManager(np.shape(X_eval)[0], batch_size, perm=np.arange(np.shape(X_eval)[0]))\n",
    "    mu_vals_eval = []\n",
    "    while local_perms.epoch < 1:\n",
    "        batch = local_perms.get_indices()\n",
    "        X_batch = np.reshape(X_eval[batch, :], [batch_size, -1])\n",
    "        mu_vals_eval.append(sess.run(mu_Z, {X: X_batch}))\n",
    "    mu_vals_eval = np.concatenate(mu_vals_eval)\n",
    "    digit_pred = classifier.predict(mu_vals_eval)\n",
    "    del local_perms\n",
    "    return float(np.sum(digit_pred == digits_eval)) / np.shape(X_eval)[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# LOAD AND WARP CIFAR-10\n",
    "batch_1 = unpickle('cifar-10-batches-py/data_batch_1')\n",
    "batch_2 = unpickle('cifar-10-batches-py/data_batch_2')\n",
    "batch_3 = unpickle('cifar-10-batches-py/data_batch_3')\n",
    "batch_4 = unpickle('cifar-10-batches-py/data_batch_4')\n",
    "batch_5 = unpickle('cifar-10-batches-py/data_batch_5')\n",
    "batch_test = unpickle('cifar-10-batches-py/test_batch')\n",
    "train_img = np.concatenate( [batch_1['data'], batch_2['data'], batch_3['data'], batch_4['data'], batch_5['data']])\n",
    "train_img = np.array(np.transpose(np.reshape(train_img, [-1, 3, 32, 32]), [0, 2, 3, 1]), dtype='float32')\n",
    "train_img = (train_img + np.array(np.random.random((50000, 32, 32, 3)), dtype='float32')) / 256.0\n",
    "train_img = warp(train_img, gamma)\n",
    "test_img = np.array(np.transpose(np.reshape(batch_test['data'], [-1, 3, 32, 32]), [0, 2, 3, 1]), dtype='float32')\n",
    "test_img = (test_img + np.array(np.random.random((10000, 32, 32, 3)), dtype='float32')) / 256.0\n",
    "test_img = warp(test_img, gamma)\n",
    "train_digits = np.concatenate([batch_1['labels'], batch_2['labels'], batch_3['labels'], batch_4['labels'],\n",
    "                               batch_5['labels']])\n",
    "train_lbl = np.array(make_one_hot(train_digits, 10), dtype='float32')\n",
    "test_digits = np.array(batch_test['labels'])\n",
    "test_lbl = np.array(make_one_hot(test_digits, 10), dtype='float32')\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# TRAIN CLASSIFIER (USED TO COMPUTE INCEPTION SCORES)\n",
    "\n",
    "sess.run(init_op)\n",
    "\n",
    "perms = PermManager(N, batch_size)\n",
    "while True:\n",
    "    start_epoch = perms.epoch\n",
    "    batch = perms.get_indices()\n",
    "    X_batch = np.reshape(train_img[batch, :], [batch_size, -1])\n",
    "    labels_batch = train_lbl[batch]\n",
    "    _, c = sess.run([optimizer_class, class_loss], {X: X_batch, labels: labels_batch, keep_prob: 0.5})\n",
    "    if perms.epoch >= max_epochs:\n",
    "        break"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# TRAIN CONTINUOUS BERNOULLI VAE\n",
    "\n",
    "# sess.run(init_op)  # uncomment to restart variables\n",
    "\n",
    "perms = PermManager(N, batch_size)\n",
    "while True:\n",
    "    start_epoch = perms.epoch\n",
    "    eps_batch = np.random.normal(size=[batch_size, d])\n",
    "    batch = perms.get_indices()\n",
    "    X_batch = np.reshape(train_img[batch, :], [batch_size, -1])\n",
    "    _, c = sess.run([optimizer, cost], {eps: eps_batch, X: X_batch})\n",
    "    if perms.epoch >= max_epochs:\n",
    "        break"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# epoch metrics computes ELBOs and log likelihoods, both including and ignoring normalizing constants, as well as\n",
    "# transforming the decoder output with the inverse mean function\n",
    "print epoch_metrics_cb(N_test, d, batch_size, test_img)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# compute the knn metric to measure usefulness of latents\n",
    "print k_nn_acc(15, batch_size, test_img, test_digits, train_img, train_digits)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# SHOW MODEL SAMPLES\n",
    "\n",
    "ind = 0\n",
    "samples = sample_from_model_cb(100, 100)\n",
    "\n",
    "plt.figure(figsize=(4, 4))\n",
    "plt.imshow(np.reshape(samples[ind], [32, 32, 3]), norm=None, vmin=0.0, vmax=1.0, cmap='gray')\n",
    "plt.axis('off')\n",
    "plt.tight_layout()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# COMPUTE INCEPTION SCORES\n",
    "\n",
    "samples = sample_from_model_cb(batch_size, N)\n",
    "# uncomment to select which IS to compute\n",
    "# new_samples = np.reshape(train_img, [N, 784]) # data\n",
    "# new_samples = sample_cont_bern(samples) # CB from decoder output\n",
    "# new_samples = sample_cont_bern(mean_from_lam2(samples)) # CB from mu(decoder output)\n",
    "new_samples = samples # decoder output\n",
    "# new_samples = mean_from_lam2(samples) # mu(decoder output)\n",
    "# new_samples = sample_bern(samples) # B from decoder output\n",
    "# new_samples = sample_bern(mean_from_lam2(samples)) # B from mu(decoder output)\n",
    "# new_samples = sample_cont_bern(lam_from_means2(batch_size, samples)) # CB from mu^{-1}(decoder output)\n",
    "# new_samples = lam_from_means2(batch_size, samples) # mu^{-1}(decoder output)\n",
    "# new_samples = sample_bern(lam_from_means2(batch_size, samples)) # B from mu^{-1}(decoder output)\n",
    "\n",
    "print compute_IS(batch_size, new_samples)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# TRAIN BERNOULLI VAE\n",
    "\n",
    "sess.run(init_op)  # uncomment to restart variables, should be done if CB VAE was previously trained\n",
    "\n",
    "perms = PermManager(N, batch_size)\n",
    "while True:\n",
    "    start_epoch = perms.epoch\n",
    "    eps_batch = np.random.normal(size=[batch_size, d])\n",
    "    batch = perms.get_indices()\n",
    "    X_batch = np.reshape(train_img[batch, :], [batch_size, -1])\n",
    "    _, c = sess.run([optimizer_cheat, cost_cheat], {eps: eps_batch, X: X_batch})\n",
    "    if perms.epoch >= max_epochs:\n",
    "        break"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# SAME AS SOME CELLS ABOVE, MEANT TO BE RUN AFTER HAVING TRAINED THE BERNOULLI VAE FOR COMPARISON\n",
    "print epoch_metrics_cb(N_test, d, batch_size, test_img)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print k_nn_acc(15, batch_size, test_img, test_digits, train_img, train_digits)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ind = 0\n",
    "samples = sample_from_model_cb(100, 100)\n",
    "\n",
    "plt.figure(figsize=(4, 4))\n",
    "plt.imshow(np.reshape(samples[ind], [32, 32, 3]), norm=None, vmin=0.0, vmax=1.0, cmap='gray')\n",
    "plt.axis('off')\n",
    "plt.tight_layout()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# NOTE THAT CLASSIFIER SHOULD BE RETRAINED BEFORE RUNNING THIS IF THE BERNOULLI VAE WAS THE LAST TRAINED ONE\n",
    "\n",
    "samples = sample_from_model_cb(batch_size, N)\n",
    "# uncomment to select which IS to compute\n",
    "# new_samples = np.reshape(train_img, [N, 784]) # data\n",
    "# new_samples = sample_cont_bern(samples) # CB from decoder output\n",
    "# new_samples = sample_cont_bern(mean_from_lam2(samples)) # CB from mu(decoder output)\n",
    "new_samples = samples # decoder output\n",
    "# new_samples = mean_from_lam2(samples) # mu(decoder output)\n",
    "# new_samples = sample_bern(samples) # B from decoder output\n",
    "# new_samples = sample_bern(mean_from_lam2(samples)) # B from mu(decoder output)\n",
    "# new_samples = sample_cont_bern(lam_from_means2(batch_size, samples)) # CB from mu^{-1}(decoder output)\n",
    "# new_samples = lam_from_means2(batch_size, samples) # mu^{-1}(decoder output)\n",
    "# new_samples = sample_bern(lam_from_means2(batch_size, samples)) # B from mu^{-1}(decoder output)\n",
    "\n",
    "print compute_IS(batch_size, new_samples)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 2",
   "language": "python",
   "name": "python2"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 2
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython2",
   "version": "2.7.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}