assignment04-large.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!nvidia-smi  # check GPU usage -- can ignore this"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# ignore this cell -- stuff for our server\n",
    "import os\n",
    "os.environ[\"CUDA_VISIBLE_DEVICES\"] = \"0\"\n",
    "os.environ['HTTP_PROXY']='http://proxy:3128/'\n",
    "os.environ['HTTPS_PROXY']='http://proxy:3128/'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# imports\n",
    "import tensorflow as tf\n",
    "import tensorflow_probability as tfp\n",
    "import numpy as np\n",
    "from matplotlib import pyplot as plt\n",
    "from matplotlib import image as mpimage\n",
    "\n",
    "from data.utils import parse_image_example\n",
    "from modeling.layers import ConvNormAct, ResidualBlock\n",
    "\n",
    "tfkl = tf.keras.layers"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "batch_size = 128\n",
    "train_data = tf.data.TFRecordDataset(\"data/flickr_64_train.TFR\").shuffle(60000).map(parse_image_example).batch(batch_size)\n",
    "test_data = tf.data.TFRecordDataset(\"data/flickr_64_test.TFR\").map(parse_image_example).batch(batch_size)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "test_images = np.concatenate([batch for batch in iter(test_data)], axis=0)\n",
    "\n",
    "plt.figure(figsize=(15,15))\n",
    "for ind, img in enumerate(test_images[:64]):\n",
    "    plt.subplot(8, 8, ind+1)\n",
    "    plt.imshow(img)\n",
    "    plt.axis(\"off\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "test_images.max()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# a few options for loss functions.\n",
    "# note: all of these losses return a batch_size vector, i.e. no averaging over the batch dimension at this point.\n",
    "# also note that we generally SUM over dimensions of the image!\n",
    "\n",
    "def squared_loss(y_true, y_pred):\n",
    "    batch_shape = tf.shape(y_true)[0]\n",
    "    y_true = tf.reshape(y_true, [batch_shape, -1])\n",
    "    y_pred = tf.reshape(y_pred, [batch_shape, -1])\n",
    "    \n",
    "    return tf.reduce_sum((y_true - y_pred)**2, axis=-1)\n",
    "\n",
    "\n",
    "def logloss(y_true, y_pred, epsilon=0.):\n",
    "    # this is what we get if assuming a gaussian likelihood and choosing optimal sigma PER IMAGE\n",
    "    batch_shape = tf.shape(y_true)[0]\n",
    "    y_true = tf.reshape(y_true, [batch_shape, -1])\n",
    "    y_pred = tf.reshape(y_pred, [batch_shape, -1])\n",
    "    \n",
    "    return tf.cast(tf.shape(y_true)[1], tf.float32) * tf.math.log(tf.norm(y_true - y_pred, axis=-1) + epsilon)\n",
    "\n",
    "\n",
    "def logloss2(y_true, y_pred, epsilon=0.):\n",
    "    # this is what we get if assuming a gaussian likelihood and choosing optimal sigma PER PIXEL.\n",
    "    # very unstable >:(\n",
    "    batch_shape = tf.shape(y_true)[0]\n",
    "    y_true = tf.reshape(y_true, [batch_shape, -1])\n",
    "    y_pred = tf.reshape(y_pred, [batch_shape, -1])\n",
    "    \n",
    "    return tf.reduce_sum(tf.math.log(tf.abs(y_true - y_pred) + epsilon), axis=-1)\n",
    "\n",
    "\n",
    "def bernoulli_loss(y_true, y_pred):\n",
    "    batch_shape = tf.shape(y_true)[0]\n",
    "    y_true = tf.reshape(y_true, [batch_shape, -1])\n",
    "    y_pred = tf.reshape(y_pred, [batch_shape, -1])\n",
    "    \n",
    "    xent = tf.nn.sigmoid_cross_entropy_with_logits(labels=y_true, logits=y_pred)\n",
    "    \n",
    "    return tf.reduce_sum(xent, axis=-1)\n",
    "\n",
    "\n",
    "def continuous_bernoulli_loss(y_true, y_pred):\n",
    "    # this is the loss for the continuous bernoulli distribution\n",
    "    # it's really just binary cross-entropy plus one more term corresponding to the normalization constant\n",
    "    batch_shape = tf.shape(y_true)[0]\n",
    "    y_true = tf.reshape(y_true, [batch_shape, -1])\n",
    "    y_pred = tf.reshape(y_pred, [batch_shape, -1])\n",
    "    \n",
    "    base = tf.nn.sigmoid_cross_entropy_with_logits(labels=y_true, logits=y_pred)\n",
    "    log_normalizer = continuous_bernoulli_log_normalizer(tf.clip_by_value(tf.nn.sigmoid(y_pred), 1e-4, 1-1e-4))\n",
    "    \n",
    "    return tf.reduce_sum(base - log_normalizer, axis=-1)\n",
    "\n",
    "\n",
    "def continuous_bernoulli_log_normalizer(lam, l_lim=0.49, u_lim=0.51):\n",
    "    # taken from https://github.com/cunningham-lab/cb_and_cc\n",
    "    cut_lam = tf.where(tf.logical_or(tf.less(lam, l_lim), tf.greater(lam, u_lim)), lam, l_lim * tf.ones_like(lam))\n",
    "    log_norm = tf.math.log(tf.abs(2.0 * tf.math.atanh(1 - 2.0 * cut_lam))) - tf.math.log(tf.abs(1 - 2.0 * cut_lam))\n",
    "    taylor = tf.math.log(2.0) + 4.0 / 3.0 * tf.pow(lam - 0.5, 2) + 104.0 / 45.0 * tf.pow(lam - 0.5, 4)\n",
    "    return tf.where(tf.logical_or(tf.less(lam, l_lim), tf.greater(lam, u_lim)), log_norm, taylor)\n",
    "\n",
    "\n",
    "def continuous_bernoulli_expected_value(lam, l_lim=0.49, u_lim=0.51):\n",
    "    # if using continuous bernoulli, the expected value is a bit more complicated\n",
    "    cut_lam = tf.where(tf.logical_or(tf.less(lam, l_lim), tf.greater(lam, u_lim)), lam, l_lim * tf.ones_like(lam))\n",
    "    expected = cut_lam / (2*cut_lam - 1) + 1 / (2*tf.math.atanh(1 - 2*cut_lam))\n",
    "    return tf.where(tf.logical_or(tf.less(lam, l_lim), tf.greater(lam, u_lim)), expected, 0.5*tf.ones_like(expected))\n",
    "\n",
    "\n",
    "def kl_loss_function(means, log_variances):\n",
    "    batch_shape = tf.shape(means)[0]\n",
    "    means = tf.reshape(means, [batch_shape, -1])\n",
    "    log_variances = tf.reshape(log_variances, [batch_shape, -1])\n",
    "    \n",
    "    return 0.5 * tf.reduce_sum(means**2 - 1 + tf.exp(log_variances) - log_variances, axis=-1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class VAE(tf.keras.Model):\n",
    "    def __init__(self, inputs, encoder, decoder, reconstruction_loss_fn, beta=1., **kwargs):\n",
    "        super().__init__(inputs, decoder(self.sample_codes(*encoder(inputs))), **kwargs)\n",
    "        self.encoder = encoder\n",
    "        self.decoder = decoder\n",
    "        self.reconstruction_loss_fn = reconstruction_loss_fn\n",
    "        \n",
    "        self.loss_tracker = tf.keras.metrics.Mean(\"loss\")\n",
    "        self.kl_tracker = tf.keras.metrics.Mean(\"kld\")\n",
    "        self.recon_tracker = tf.keras.metrics.Mean(\"recon_loss\")\n",
    "        \n",
    "        # to get \"standard\" VAE, use beta=1\n",
    "        self.beta = beta\n",
    "        \n",
    "    def train_step(self, data):\n",
    "        with tf.GradientTape() as tape:\n",
    "            means, log_variances = encoder(data, training=True)\n",
    "            sampled_codes = self.sample_codes(means, log_variances)\n",
    "            reconstructions = decoder(sampled_codes, training=True)\n",
    "            \n",
    "            recon_loss = self.reconstruction_loss_fn(data, reconstructions)\n",
    "            \n",
    "            kl_loss = kl_loss_function(means, log_variances)\n",
    "            \n",
    "            total_loss = tf.reduce_mean(recon_loss + self.beta*kl_loss)\n",
    "            \n",
    "        variables = self.trainable_variables\n",
    "        gradients = tape.gradient(total_loss, variables)\n",
    "        self.optimizer.apply_gradients(zip(gradients, variables))\n",
    "        \n",
    "        self.loss_tracker.update_state(total_loss)\n",
    "        self.kl_tracker.update_state(kl_loss)\n",
    "        self.recon_tracker.update_state(recon_loss)\n",
    "        \n",
    "        return {\"loss\": self.loss_tracker.result(),\n",
    "                \"kld\": self.kl_tracker.result(),\n",
    "                \"recon_loss\": self.recon_tracker.result()}\n",
    "    \n",
    "    def test_step(self, data):\n",
    "        means, log_variances = encoder(data, training=False)\n",
    "        sampled_codes = self.sample_codes(means, log_variances)\n",
    "        reconstructions = decoder(sampled_codes, training=False)\n",
    "\n",
    "        recon_loss = self.reconstruction_loss_fn(data, reconstructions)\n",
    "\n",
    "        kl_loss = kl_loss_function(means, log_variances)\n",
    "\n",
    "        total_loss = tf.reduce_mean(recon_loss + self.beta*kl_loss)\n",
    "        \n",
    "        self.loss_tracker.update_state(total_loss)\n",
    "        self.kl_tracker.update_state(kl_loss)\n",
    "        self.recon_tracker.update_state(recon_loss)\n",
    "        \n",
    "        return {\"loss\": self.loss_tracker.result(),\n",
    "                \"kld\": self.kl_tracker.result(),\n",
    "                \"recon_loss\": self.recon_tracker.result()}\n",
    "    \n",
    "    @property\n",
    "    def metrics(self):\n",
    "        return [self.loss_tracker, self.kl_tracker, self.recon_tracker]\n",
    "    \n",
    "    \n",
    "    def sample_codes(self, means, log_variances):\n",
    "        stddevs = tf.exp(0.5*log_variances)\n",
    "        return stddevs * tf.random.normal(tf.shape(means)) + means"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def residual_stack(inputs, filters, strides, blocks_per_level, mode, name):\n",
    "    outputs = inputs\n",
    "    for level_ind, (level_filters, level_stride) in enumerate(zip(filters, strides)):\n",
    "        for block_ind in range(blocks_per_level):\n",
    "            outputs = ResidualBlock(level_filters,\n",
    "                                    mode, \n",
    "                                    strides=level_stride if block_ind == (blocks_per_level - 1) else 1,\n",
    "                                    name=\"_\".join([name, str(level_ind+1), str(block_ind+1)]))(outputs)\n",
    "        \n",
    "    return outputs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# build the model\n",
    "likelihood = \"continuous_bernoulli\"\n",
    "\n",
    "loss_dictionary = {\"bernoulli\": bernoulli_loss, \n",
    "                   \"continuous_bernoulli\": continuous_bernoulli_loss, \n",
    "                   \"gaussian_fixed_sigma\": squared_loss,\n",
    "                   \"gaussian_image_sigma\": logloss, \n",
    "                   \"gaussian_pixel_sigma\": logloss2}\n",
    "\n",
    "\n",
    "if likelihood not in loss_dictionary.keys():\n",
    "    raise ValueError(\"Invalid likelihood!\")\n",
    "\n",
    "tf.keras.backend.clear_session()\n",
    "\n",
    "#normalization = lambda **kwargs: tfkl.GroupNormalization(groups=32, **kwargs)\n",
    "normalization = tfkl.BatchNormalization\n",
    "\n",
    "blocks_per_level = 4\n",
    "filters = [64, 128, 256, 512, 512]\n",
    "strides = [2, 2, 2, 2, 1]\n",
    "\n",
    "encoder_input = tf.keras.Input((64, 64, 3))\n",
    "encoder_output = residual_stack(encoder_input, filters, strides, blocks_per_level, \"conv\", \"encoder\")\n",
    "encoder_final = tfkl.Flatten()(encoder_output)\n",
    "encoder_final = tfkl.Dense(2*512)(encoder_final)\n",
    "\n",
    "means, log_variances = tf.split(encoder_final, 2, axis=-1)\n",
    "encoder = tf.keras.Model(encoder_input, [means, log_variances], name=\"encoder\")\n",
    "code_shape = encoder.output_shape[0][1:]\n",
    "\n",
    "\n",
    "decoder_input = tf.keras.Input(code_shape)\n",
    "\n",
    "decoder_front = tfkl.Dense(4*4*512)(decoder_input)\n",
    "decoder_front = tfkl.Reshape((4, 4, 512))(decoder_front)\n",
    "\n",
    "decoder_output = residual_stack(decoder_front, reversed(filters), strides, blocks_per_level, \"transpose\", \"decoder\")\n",
    "decoder_final = tfkl.Conv2D(3, 1)(decoder_output)\n",
    "\n",
    "decoder = tf.keras.Model(decoder_input, decoder_final, name=\"decoder\")\n",
    "\n",
    "# beta is set here! IMPORTANT PARAMETER!\n",
    "model = VAE(encoder_input, encoder, decoder, loss_dictionary[likelihood], beta=4., name=\"vae\")\n",
    "model.summary(expand_nested=True)\n",
    "\n",
    "# this initializes the final layer weights to 0.\n",
    "# seems to help with exploding loss in the beginning.\n",
    "# in some sense this is the \"perfect\" initialization for the KL divergence,\n",
    "# as all means will be 0 and all log variances as well(-> variances will be 1)\n",
    "model.layers[1].layers[-2].weights[0].assign(tf.zeros_like(model.layers[1].layers[-2].weights[0]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "n_steps = 100000\n",
    "n_data = 60000\n",
    "n_epochs = n_steps // (n_data // batch_size)\n",
    "\n",
    "learning_rate = 0.0002\n",
    "#decay_fn = tf.keras.optimizers.schedules.CosineDecay(learning_rate, n_steps)\n",
    "optimizer = tf.optimizers.Adam(learning_rate)\n",
    "\n",
    "model.compile(optimizer=optimizer)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "def map_for_likelihood(inputs, likelihood):\n",
    "    # depending on what likelihood we choose for our data, we may need to model outputs to something else.\n",
    "    # - for bernoulli likelihood, technically we don't have to do anything.\n",
    "    #   however, we usually don't want a sigmoid output in our model when we use cross-entropy loss.\n",
    "    #   so we do the sigmoid here.\n",
    "    # - for continuous bernoulli, we need to map the sigmoid result further.\n",
    "    # - for gaussian, we don't need to do anything, but since the output is unconstrained,\n",
    "    #   it makes sense to clip to [0, 1]\n",
    "    if \"bernoulli\" in likelihood:\n",
    "        outputs = tf.nn.sigmoid(inputs)\n",
    "        if likelihood == \"continuous_bernoulli\":\n",
    "            outputs = continuous_bernoulli_expected_value(\n",
    "                tf.clip_by_value(outputs, 1e-4, 1-1e-4))\n",
    "    else:  # gaussian\n",
    "        outputs = np.clip(inputs, 0, 1)\n",
    "    return outputs\n",
    "\n",
    "\n",
    "class ImageGenCallback(tf.keras.callbacks.Callback):\n",
    "    def __init__(self, frequency, **kwargs):\n",
    "        super().__init__(**kwargs)\n",
    "        self.frequency = frequency\n",
    "        \n",
    "    def on_epoch_end(self, epoch, logs=None):\n",
    "        if not epoch % self.frequency:\n",
    "            noise = tf.random.normal((64,) + code_shape)\n",
    "            generated_batch = map_for_likelihood(self.model.decoder(noise), likelihood)\n",
    "        \n",
    "            plt.figure(figsize=(15,15))\n",
    "            for ind, image in enumerate(generated_batch):\n",
    "                plt.subplot(8, 8, ind+1)\n",
    "                plt.imshow(image)\n",
    "                plt.axis(\"off\")\n",
    "            plt.suptitle(\"Random generations\")\n",
    "            plt.show()\n",
    "\n",
    "\n",
    "class ReconstructionCallback(tf.keras.callbacks.Callback):\n",
    "    def __init__(self, frequency, **kwargs):\n",
    "        super().__init__(**kwargs)\n",
    "        self.frequency = frequency\n",
    "    \n",
    "    def on_epoch_end(self, epoch, logs=None):\n",
    "        if not epoch % self.frequency:\n",
    "            cropped_test = test_images[:32]\n",
    "            generated_batch = map_for_likelihood(self.model(cropped_test), likelihood)\n",
    "        \n",
    "            plt.figure(figsize=(15,15))\n",
    "            for ind, (original, reconstruction) in enumerate(zip(cropped_test, generated_batch)):\n",
    "                comparison = np.concatenate((original, reconstruction), axis=1)\n",
    "                plt.subplot(8, 4, ind+1)\n",
    "                plt.imshow(comparison)\n",
    "                plt.axis(\"off\")\n",
    "            plt.suptitle(\"Test set reconstructions\")\n",
    "            plt.show()\n",
    "\n",
    "do_train = True\n",
    "\n",
    "# In general I'm not super happy with the training :(\n",
    "# early stopping tends to kick in rather... early, so not many steps are done.\n",
    "# maybe 40 epochs or so, depending on beta.\n",
    "# I think the architecture is still not otpimal.\n",
    "# however, tuning it takes quite long.\n",
    "\n",
    "\n",
    "if do_train:\n",
    "    lr_schedule = tf.keras.callbacks.ReduceLROnPlateau(factor=0.2, patience=5, verbose=1)\n",
    "    early_stop = tf.keras.callbacks.EarlyStopping(patience=20, restore_best_weights=True, verbose=1)\n",
    "    reconstruct = ReconstructionCallback(10)\n",
    "    image_gen_callback = ImageGenCallback(10)\n",
    "\n",
    "    history = model.fit(train_data, epochs=n_epochs, validation_data=test_data,\n",
    "                        callbacks=[lr_schedule, early_stop, reconstruct, image_gen_callback])\n",
    "    model.save_weights(\"weights/weights_assignment04_large.hdf5\")\n",
    "\n",
    "    # note: the if using the continuous bernoulli loss, the loss will likely be < 0.\n",
    "    # this might seem a bit weird, but is actually not an issue.\n",
    "    # recall that this loss is the negative log likelihood.\n",
    "    # - if the NLL is negative, that means the log likelihood is positive\n",
    "    # - if the log likelihood is > 0, that implies that the likelihood is > 1\n",
    "    # - a p > 1 is nothing unusual for continuous distributions, where we are using *density* functions\n",
    "else:\n",
    "    model.load_weights(\"weights/weights_assignment04_large.hdf5\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model.evaluate(test_data)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# compare some inputs and reconstructions\n",
    "reconstructions = map_for_likelihood(model.predict(test_data), likelihood)\n",
    "\n",
    "plt.figure(figsize=(15, 15))\n",
    "for ind, (original, reconstruction) in enumerate(zip(test_images[:32], reconstructions[:32])):\n",
    "    plt.subplot(8, 4, ind+1)\n",
    "    concat = np.concatenate((original, reconstruction), axis=1)\n",
    "    plt.imshow(concat, vmin=0, vmax=1)\n",
    "    plt.axis(\"off\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "random_codes = tf.random.normal((64,) + code_shape)\n",
    "generated = map_for_likelihood(decoder(random_codes), likelihood)\n",
    "\n",
    "plt.figure(figsize=(15, 15))\n",
    "for ind, image in enumerate(generated):\n",
    "    plt.subplot(8, 8, ind+1)\n",
    "    plt.imshow(image, vmin=0, vmax=1)\n",
    "    plt.axis(\"off\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# can check if these are close to 0 (i.e. did the KL divergence do its job)\n",
    "all_means, all_log_variances = encoder.predict(test_data)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.hist(all_means.reshape(-1), bins=100)\n",
    "plt.title(\"Distribution of latent means in test set\")\n",
    "plt.show()\n",
    "\n",
    "plt.hist(all_log_variances.reshape(-1), bins=100)\n",
    "plt.title(\"Distribution of latent log variances in test set\")\n",
    "plt.show()\n",
    "\n",
    "# should be close to 1\n",
    "plt.hist(np.exp(all_log_variances.reshape(-1)), bins=100)\n",
    "plt.title(\"Distribution of latent variances in test set\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "randind1 = np.random.randint(len(test_images))\n",
    "randind2 = np.random.randint(len(test_images))\n",
    "\n",
    "# annoying indexing since we only want 1 code but need a batch axis\n",
    "code1_mean, code1_logvar = encoder(test_images[randind1][None])\n",
    "code1 = model.sample_codes(code1_mean, code1_logvar)\n",
    "\n",
    "code2_mean, code2_logvar = encoder(test_images[randind2][None])\n",
    "code2 = model.sample_codes(code2_mean, code2_logvar)\n",
    "\n",
    "n_interpolation = 64\n",
    "interpolation_coeffs = np.linspace(0, 1, n_interpolation)\n",
    "\n",
    "\n",
    "interpolation = \"slerp\"  # or \"linear\"\n",
    "# due to \"reasons\", spherical interpolation may be more appropriate\n",
    "\n",
    "# note, the 0 indexing in here is just to get rid of the batch axis\n",
    "interpolated_codes = []\n",
    "for coefficient in interpolation_coeffs:\n",
    "    if interpolation == \"linear\":\n",
    "        interpolated_codes.append(coefficient*code2[0] + (1-coefficient)*code1[0])\n",
    "    elif interpolation == \"slerp\":\n",
    "        angle = np.arccos(np.dot(code1[0], code2[0]) / (np.linalg.norm(code1)*np.linalg.norm(code2)))\n",
    "        interpolated_codes.append(np.sin((1-coefficient)*angle)/np.sin(angle) * code1[0] \n",
    "                                  + np.sin(coefficient*angle)/np.sin(angle) * code2[0])\n",
    "    else:\n",
    "        raise ValueError(\"invalid interpolation\")\n",
    "interpolated_codes = np.array(interpolated_codes)\n",
    "\n",
    "\n",
    "# see the images we are dealing with here\n",
    "plt.imshow(test_images[randind1])\n",
    "plt.axis(\"off\")\n",
    "plt.show()\n",
    "\n",
    "plt.imshow(test_images[randind2])\n",
    "plt.axis(\"off\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# show: first the original, to the right of it the reconstruction from the mean,\n",
    "# then 14 random code reconstructions\n",
    "# I only do it for the 1st image here.\n",
    "repeated_image = tf.tile(test_images[randind1:randind1+1], [14, 1, 1, 1])\n",
    "\n",
    "# I ran out of names so it's a\n",
    "a = map_for_likelihood(model(repeated_image), likelihood)\n",
    "recon_from_mean = map_for_likelihood(decoder(code1_mean), likelihood)\n",
    "a = np.concatenate([test_images[randind1][None], recon_from_mean, a])\n",
    "\n",
    "plt.figure(figsize=(8, 8))\n",
    "for ind, image in enumerate(a):\n",
    "    plt.subplot(4, 4, ind+1)\n",
    "    plt.imshow(image, vmin=0, vmax=1)\n",
    "    plt.axis(\"off\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# actually look at the interpolations\n",
    "# NOTE that the 2D grid layout implies there's something 2D happening here.\n",
    "# it's not. it's just saves space to display the images this way.\n",
    "# it's a 1d interpolation from top left to bottom right.\n",
    "# it would be clearer to plot it all in one row or column.\n",
    "interpolated_images = map_for_likelihood(decoder(interpolated_codes), likelihood)\n",
    "plt.figure(figsize=(15, 15))\n",
    "for ind in range(64):\n",
    "    plt.subplot(8, 8, ind+1)\n",
    "    plt.imshow(interpolated_images[ind])\n",
    "    plt.axis(\"off\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "# here's a fun thing we can do to check the usage of the latent space.\n",
    "# we iterate over the latent dimensions. \n",
    "# each time, we take the original code for an image and fix all dimensions except the current one.\n",
    "# the current dimension is moved in some range (here, -2 to 2), sampling multiple points along the way.\n",
    "# each resulting code is decoded and displayed.\n",
    "# I compute the difference in pixel space between successive images.\n",
    "# for some dimensions, there is next to no difference.\n",
    "# I conclude that these dimensions are effectively unused.\n",
    "# for dimensions where the average change is larger than some threshold, I plot the \"walk\" over that dimension.\n",
    "\n",
    "# an issue here is that I'm only doing this for a single image. ideally, we would sample images randomly, or\n",
    "# do the entire process for many images.\n",
    "\n",
    "diffs = []\n",
    "for latent_dim in range(code_shape[-1]):\n",
    "\n",
    "    code_repeated = np.tile(code1, [n_interpolation, 1])\n",
    "    value_range = np.linspace(-2, 2, n_interpolation)\n",
    "\n",
    "    for dim in range(len(code_repeated)):\n",
    "        code_repeated[dim, latent_dim] = value_range[dim]\n",
    "\n",
    "    interpolated_images = map_for_likelihood(decoder(code_repeated), likelihood)\n",
    "\n",
    "    avg_abs_diff = np.abs(interpolated_images[0] - interpolated_images[-1]).sum() / np.prod(interpolated_images[0].shape)\n",
    "    diffs.append(avg_abs_diff)\n",
    "    # this value is somewhat arbitrary!\n",
    "    if avg_abs_diff < 0.05:\n",
    "        print(\"Dim {} not used\".format(latent_dim))\n",
    "    else:\n",
    "        print(avg_abs_diff)\n",
    "        plt.figure(figsize=(15, 15))\n",
    "        for ind in range(64):\n",
    "            plt.subplot(8, 8, ind+1)\n",
    "            plt.imshow(interpolated_images[ind])\n",
    "            plt.axis(\"off\")\n",
    "        plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# here I sort the differences computed above.\n",
    "# we can see that there is a sharp dropoff in difference at some point.\n",
    "# this may be seen as a kind of cut-off point which allows us to (very approximately)\n",
    "# read off the number of dimensions actually used in the latent space.\n",
    "# this will vary heavily with beta (higher beta -> fewer dimensions used).\n",
    "plt.loglog(sorted(diffs, reverse=True))\n",
    "plt.xlabel(\"Dimension index\")\n",
    "plt.ylabel(\"Average difference\")\n",
    "plt.title(\"Average differences between successive images in latent space walk per dimension.\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.8.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}