mnist_vae.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "Ndo4ERqnwQOU"
   },
   "source": [
    "# (Convolutional) Autoencoder - Toy MNIST Example\n",
    "\n",
    "##### Code closely from [Tensorflow Tutorial](https://www.tensorflow.org/tutorials/generative/cvae)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "YfIk2es3hJEd",
    "tags": []
   },
   "outputs": [],
   "source": [
    "from IPython import display\n",
    "\n",
    "import glob\n",
    "import imageio\n",
    "import PIL\n",
    "import time\n",
    "from typing import List, Tuple, Dict, Optional\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "import tensorflow as tf\n",
    "import tensorflow_probability as tfp\n",
    "tfk = tf.keras\n",
    "tfkl = tf.keras.layers"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "iYn4MdZnKCey"
   },
   "source": [
    "## Load the MNIST dataset\n",
    "Each MNIST image is originally a vector of 784 integers, each of which is between 0-255 and represents the intensity of a pixel. Model each pixel with a Bernoulli distribution in our model, and statically binarize the dataset."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "a4fYMGxGhrna",
    "tags": []
   },
   "outputs": [],
   "source": [
    "(train_images, train_labels), (test_images, test_labels) = tfk.datasets.mnist.load_data()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "NFC2ghIdiZYE",
    "tags": []
   },
   "outputs": [],
   "source": [
    "def preprocess_images(images):\n",
    "    images = images.reshape((images.shape[0], 28, 28, 1)) / 255.\n",
    "    return np.where(images > .5, 1.0, 0.0).astype('float32')\n",
    "\n",
    "train_images = preprocess_images(train_images)\n",
    "test_images = preprocess_images(test_images)\n",
    "\n",
    "train_dataset = tf.data.Dataset.from_tensor_slices(train_images).shuffle(len(train_images))\n",
    "test_dataset = tf.data.Dataset.from_tensor_slices(test_images).shuffle(len(test_images))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Define Model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "VGLbvBEmjK0a",
    "tags": []
   },
   "outputs": [],
   "source": [
    "class CVAE(tfk.Model):\n",
    "    \"\"\"Convolutional variational autoencoder.\"\"\"\n",
    "\n",
    "    def __init__(self, latent_dim: int, filters: List[int] = [32, 64], input_shape: Tuple[int, int, int] = (28, 28, 1)) -> None:\n",
    "        super().__init__()\n",
    "        self.latent_dim = latent_dim\n",
    "        enc_convs = [tfkl.Conv2D(filters=f, kernel_size=3, strides=2, activation='relu') for f in filters]\n",
    "        dec_convs = [tfkl.Conv2DTranspose(filters=f, kernel_size=3, strides=2, activation='relu', padding='same') for f in filters[::-1]]\n",
    "        self.encoder = tf.keras.Sequential(\n",
    "            [\n",
    "                tfkl.InputLayer(input_shape=input_shape),\n",
    "                *enc_convs,\n",
    "                tfkl.Flatten(),\n",
    "                tfkl.Dense(latent_dim + latent_dim),\n",
    "            ]\n",
    "        )\n",
    "\n",
    "        self.decoder = tf.keras.Sequential(\n",
    "            [\n",
    "                tfkl.InputLayer(input_shape=(latent_dim,)),\n",
    "                tfkl.Dense(units=7*7*32, activation=tf.nn.relu),\n",
    "                tfkl.Reshape(target_shape=(7, 7, 32)),\n",
    "                *dec_convs,\n",
    "                tfkl.Conv2DTranspose(\n",
    "                    filters=1, kernel_size=3, strides=1, padding='same'),\n",
    "            ]\n",
    "        )\n",
    "\n",
    "    # tf.function adds the function to the computational graph for compilation\n",
    "    @tf.function\n",
    "    def sample(self, eps: Optional[tf.Tensor] = None) -> tf.Tensor:\n",
    "        if eps is None:\n",
    "            eps = tf.random.normal(shape=(100, self.latent_dim))\n",
    "        return self.decode(eps, apply_sigmoid=True)\n",
    "\n",
    "    def encode(self, x: tf.Tensor) -> Tuple[tf.Tensor, ...]:\n",
    "        \"\"\"Encoder of VAE\"\"\"\n",
    "        mean, logvar = tf.split(self.encoder(x), num_or_size_splits=2, axis=1)\n",
    "        eps = tf.random.normal(shape=mean.shape)\n",
    "        z = eps * tf.exp(logvar * .5) + mean\n",
    "        return z, mean, logvar\n",
    "\n",
    "    def decode(self, z: tf.Tensor, apply_sigmoid: bool = False) -> tf.Tensor:\n",
    "        \"\"\"Decoder of VAE\"\"\"\n",
    "        logits = self.decoder(z)\n",
    "        if apply_sigmoid:\n",
    "            probs = tf.sigmoid(logits)\n",
    "            return probs\n",
    "        return logits"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Helper Functions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "iWCn_PVdEJZ7",
    "tags": []
   },
   "outputs": [],
   "source": [
    "def log_normal_pdf(sample: tf.Tensor, mean: tf.Tensor, logvar: tf.Tensor, raxis: int = 1) -> tf.Tensor:\n",
    "    log2pi = tf.math.log(2. * np.pi)\n",
    "    return tf.reduce_sum(\n",
    "        -.5 * ((sample - mean) ** 2. * tf.exp(-logvar) + logvar + log2pi),\n",
    "        axis=raxis)\n",
    "\n",
    "# Analytic KL for Multivariate Gaussian\n",
    "def kl_divergence(mean: tf.Tensor, sd: tf.Tensor) -> tf.Tensor:\n",
    "    return - 0.5 * tf.reduce_sum(\n",
    "            1 + tf.math.log(tf.math.square(sd)) - tf.math.square(mean) - tf.math.square(sd),\n",
    "            axis=1)\n",
    "\n",
    "def compute_loss(model: tfk.Model, \n",
    "                 x: tf.Tensor, \n",
    "                 loss: str = \"cross_entropy\", \n",
    "                 analytic_kl: bool = False, \n",
    "                 beta: float = 1.) -> Tuple[tf.Tensor, ...]:\n",
    "\n",
    "    # forward pass\n",
    "    z, mean, logvar = model.encode(x)\n",
    "    x_logit = model.decode(z)\n",
    "    \n",
    "    # compute neg log likelihood\n",
    "    if loss == \"cross_entropy\":\n",
    "        cross_ent = tf.nn.sigmoid_cross_entropy_with_logits(logits=x_logit, labels=x)\n",
    "        logpx_z = -tf.reduce_sum(cross_ent, axis=[1, 2, 3])\n",
    "    elif loss == \"mse\":\n",
    "        logpx_z = tf.reduce_mean((x - x_logit)**2 / np.prod(x.shape), axis=[1, 2, 3])\n",
    "        \n",
    "    # compute kl divergence\n",
    "    if analytic_kl:\n",
    "        sd = tf.exp(logvar * .5)\n",
    "        kl_div = kl_divergence(mean, sd)\n",
    "    else:\n",
    "        logpz = log_normal_pdf(z, 0., 0.)\n",
    "        logqz_x = log_normal_pdf(z, mean, logvar)\n",
    "        kl_div = logpz - logqz_x\n",
    "    kl_div *= beta\n",
    "    \n",
    "    return -tf.reduce_mean(logpx_z + kl_div), logpx_z, kl_div\n",
    "\n",
    "@tf.function\n",
    "def train_step(model: tfk.Model, \n",
    "               x: tf.Tensor, \n",
    "               optimizer: tfk.optimizers.Optimizer, \n",
    "               loss: str = \"cross_entropy\",\n",
    "               analytic_kl: bool = False,\n",
    "               beta: float = 1.) -> Tuple[tf.Tensor, ...]:\n",
    "    \"\"\"Executes one training step and returns the loss.\n",
    "\n",
    "    This function computes the loss and gradients, and uses the latter to\n",
    "    update the model's parameters.\n",
    "    \"\"\"\n",
    "    with tf.GradientTape() as tape:\n",
    "        elbo, nll, kl_div = compute_loss(model, x)\n",
    "    gradients = tape.gradient(elbo, model.trainable_variables)\n",
    "    optimizer.apply_gradients(zip(gradients, model.trainable_variables))\n",
    "    return elbo, nll, kl_div"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Hyperparameter"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "EPOCHS = 10\n",
    "BATCH_SIZE = 32\n",
    "LR = 1e-4\n",
    "LATENT_DIM = 2\n",
    "LOSS = \"mse\" # \"cross_entropy\"\n",
    "ANALYTIC_KL = True\n",
    "BETA = 1.\n",
    "FILTERS = [32, 64]\n",
    "\n",
    "train_dataset = train_dataset.batch(BATCH_SIZE)\n",
    "test_dataset = test_dataset.batch(BATCH_SIZE)\n",
    "num_examples_to_generate = 16"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model = CVAE(LATENT_DIM, filters=FILTERS)\n",
    "optimizer = tfk.optimizers.Adam(LR)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "RmdVsmvhPxyy",
    "tags": []
   },
   "outputs": [],
   "source": [
    "def generate_images(model: tfk.Model, test_sample: tf.Tensor) -> None:\n",
    "    z, mean, logvar = model.encode(test_sample)\n",
    "    predictions = model.sample(z)\n",
    "    fig = plt.figure(figsize=(4, 4))\n",
    "\n",
    "    for i in range(predictions.shape[0]):\n",
    "        plt.subplot(4, 4, i + 1)\n",
    "        plt.imshow(predictions[i, :, :, 0], cmap='gray')\n",
    "        plt.axis('off')\n",
    "    \n",
    "    plt.tight_layout()\n",
    "    plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "swCyrbqQQ-Ri",
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Pick a sample of the test set for generating output images\n",
    "assert BATCH_SIZE >= num_examples_to_generate\n",
    "for test_batch in test_dataset.take(1):\n",
    "    test_sample = test_batch[0:num_examples_to_generate, :, :, :]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "2M7LmLtGEMQJ",
    "tags": []
   },
   "outputs": [],
   "source": [
    "generate_images(model, test_sample)\n",
    "\n",
    "for epoch in range(1, EPOCHS + 1):\n",
    "    start_time = time.time()\n",
    "    for train_x in train_dataset:\n",
    "        elbo, nll, kl_div = train_step(model, train_x, optimizer, loss=LOSS, analytic_kl=ANALYTIC_KL, beta=BETA)\n",
    "    end_time = time.time()\n",
    "\n",
    "    for test_x in test_dataset:\n",
    "        elbo, nll, kl_div = compute_loss(model, test_x, loss=LOSS, analytic_kl=ANALYTIC_KL, beta=BETA)\n",
    "\n",
    "    display.clear_output(wait=False)\n",
    "    print(f'Epoch: {epoch}, Test set ELBO: {elbo.numpy():.2f}, time elapse for current epoch: {end_time - start_time:.2f}')\n",
    "    generate_images(model, test_sample)\n",
    "    \n",
    "    # model.save_weights('./checkpoints/mnist_model')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Plot Latent Space\n",
    "\n",
    "If `LATENT_DIM > 2` we must use a dimensionality reduction method for which we chose [t-SNE](https://towardsdatascience.com/an-introduction-to-t-sne-with-python-example-5a3a293108d1) to plot in 2D. \n",
    "\n",
    "__Question__: Why use t-SNE and not PCA?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model.load_weights('./checkpoints/mnist_model')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from sklearn.manifold import TSNE\n",
    "\n",
    "def plot_latent_space(model: tfk.Model, test_images: np.ndarray, test_labels: np.ndarray) -> None:\n",
    "    zs = np.array([model.encode(test_x[None,...,None])[0].numpy()[0].tolist() for test_x in test_images])\n",
    "    if zs.shape[-1] > 2: # check if latent_dim > 2 -> dimensionality reduction with tsne for plotting purposes\n",
    "        tsne = TSNE()\n",
    "        zs = tsne.fit_transform(zs)\n",
    "    zs = {n: zs[test_labels == n] for n in range(10)}\n",
    "    for n, zs_n in zs.items():\n",
    "        plt.scatter(zs_n[:,0], zs_n[:,1], label=n)\n",
    "    plt.legend()\n",
    "    plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plot_latent_space(model, test_images, test_labels)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "PeunRU6TSumT"
   },
   "source": [
    "#### Underneath only works if `LATENT_DIM = 2`\n",
    "\n",
    "### Display a 2D manifold of digits from the latent space\n",
    "\n",
    "Running the code below will show a continuous distribution of the different digit classes, with each digit morphing into another across the 2D latent space. Use [TensorFlow Probability](https://www.tensorflow.org/probability) to generate a standard normal distribution for the latent space."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "cellView": "code",
    "id": "mNcaaYPBS3mj"
   },
   "outputs": [],
   "source": [
    "def plot_latent_images(model: tfk.Model, n: int, digit_size: int = 28) -> None:\n",
    "    \"\"\"Plots n x n digit images decoded from the latent space.\"\"\"\n",
    "\n",
    "    norm = tfp.distributions.Normal(0, 1)\n",
    "    grid_x = norm.quantile(np.linspace(0.05, 0.95, n))\n",
    "    grid_y = norm.quantile(np.linspace(0.05, 0.95, n))\n",
    "    image_width = digit_size*n\n",
    "    image_height = image_width\n",
    "    image = np.zeros((image_height, image_width))\n",
    "\n",
    "    for i, yi in enumerate(grid_x):\n",
    "        for j, xi in enumerate(grid_y):\n",
    "            z = np.array([[xi, yi]])\n",
    "            x_decoded = model.sample(z)\n",
    "            digit = tf.reshape(x_decoded[0], (digit_size, digit_size))\n",
    "            image[i * digit_size: (i + 1) * digit_size,\n",
    "                j * digit_size: (j + 1) * digit_size] = digit.numpy()\n",
    "\n",
    "    plt.figure(figsize=(10, 10))\n",
    "    plt.imshow(image, cmap='Greys_r')\n",
    "    plt.axis('Off')\n",
    "    plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "F-ZG69QCZnGY",
    "tags": []
   },
   "outputs": [],
   "source": [
    "plot_latent_images(model, 20)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Choose Your Own Latent Space"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "z = np.array([[-2, -2]])\n",
    "pred = model.decode(z, apply_sigmoid=True)\n",
    "plt.imshow(pred[0,...,0].numpy(), cmap=\"gray\")\n",
    "plt.axis('off')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "colab": {
   "collapsed_sections": [],
   "name": "cvae.ipynb",
   "toc_visible": true
  },
  "kernelspec": {
   "display_name": "vae-tutorial",
   "language": "python",
   "name": "vae-tutorial"
  },
  "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.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}