Keras_MMD_Variational_Autoencoder.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "Keras-MMD-Variational-Autoencoder.ipynb",
      "version": "0.3.2",
      "provenance": [],
      "collapsed_sections": []
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "accelerator": "GPU"
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "IjrC9Djc0cmB",
        "colab_type": "text"
      },
      "source": [
        "A keras verison of the MMD-Variational-Autoencoder is implemented here."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "kkDnfW850c-X",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "import numpy as np\n",
        "from matplotlib import pyplot as plt\n",
        "import math, os\n",
        "from keras.layers.advanced_activations import LeakyReLU\n",
        "from keras.layers import Conv2D, Conv2DTranspose, Dense, Input, Reshape, Flatten\n",
        "from keras.models import Model\n",
        "from keras.utils import plot_model\n",
        "import pdb\n",
        "from keras import backend as K\n",
        "from keras.datasets import mnist"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "kSIuBibJ5GJK",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "'Set parameters'\n",
        "z_dim = 2\n",
        "epoch = 100\n",
        "batch_size = 200"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "8FtUIupU04Vg",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "'Build the encoder'\n",
        "encoder_input = Input(shape=(28, 28, 1))\n",
        "en_conv1 = Conv2D(filters = 64, kernel_size=4, strides=2, padding='same')(encoder_input)\n",
        "en_conv1 = LeakyReLU(0.1)(en_conv1)\n",
        "\n",
        "en_conv2 = Conv2D(filters = 128, kernel_size=4, strides=2, padding='same')(en_conv1)\n",
        "en_conv2 = LeakyReLU(0.1)(en_conv2)\n",
        "\n",
        "en_conv2 = Flatten()(en_conv2)\n",
        "en_fc1 = Dense(1024)(en_conv2)\n",
        "en_fc1 = LeakyReLU(0.1)(en_fc1)\n",
        "\n",
        "encoder_output = Dense(z_dim)(en_fc1)\n",
        "\n",
        "encoder = Model(encoder_input, encoder_output, name='encoder')\n",
        "# plot_model(encoder, show_shapes=True)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "7S68R3EH5Rla",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "'Build the decoder'\n",
        "decoder_input = Input(shape=(z_dim,))\n",
        "de_fc1 = Dense(1024, activation='relu')(decoder_input)\n",
        "de_fc2 = Dense(7*7*128, activation='relu')(de_fc1)\n",
        "de_fc2 = Reshape((7, 7, 128))(de_fc2)\n",
        "de_conv1 = Conv2DTranspose(filters=64, kernel_size=4, strides=2, activation='relu', padding='same')(de_fc2)\n",
        "decoder_output = Conv2DTranspose(filters=1, kernel_size=4, strides=2, activation='sigmoid', padding='same')(de_conv1)\n",
        "\n",
        "decoder = Model(decoder_input, decoder_output, name='decoder')\n",
        "# plot_model(decoder, show_shapes=True)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "0LWHNzFW7Itw",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "'Build the computation graph for training'\n",
        "train_z = encoder(encoder_input)\n",
        "train_xr = decoder(train_z)\n",
        "autoencoder = Model(encoder_input, train_xr)\n",
        "\n",
        "# plot_model(autoencoder, show_shapes=True)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "eN0u4ZS9QD2x",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "# Build the computation graph for generating samples\n",
        "'our decoder can do this job, no more action needed'"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "ru7saOsZRA1y",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "'MMD functions'\n",
        "def compute_kernel(x, y):\n",
        "    x_size = K.shape(x)[0]\n",
        "    y_size = K.shape(y)[0]\n",
        "    dim = K.shape(x)[1]\n",
        "    tiled_x = K.tile(K.reshape(x, [x_size, 1, dim]), [1, y_size, 1])\n",
        "    tiled_y = K.tile(K.reshape(y, [1, y_size, dim]), [x_size, 1, 1])\n",
        "    return K.exp(-K.mean(K.square(tiled_x - tiled_y), axis=2) / K.cast(dim, 'float32'))\n",
        "\n",
        "def compute_mmd(x, y):\n",
        "    x_kernel = compute_kernel(x, x)\n",
        "    y_kernel = compute_kernel(y, y)\n",
        "    xy_kernel = compute_kernel(x, y)\n",
        "    return K.mean(x_kernel) + K.mean(y_kernel) - 2 * K.mean(xy_kernel)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "_-1dBS2ORR-h",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "def custom_loss(train_z, train_xr, train_x):\n",
        "    \"\"\"\n",
        "    Critical loss builder\n",
        "\t\t:param train_z: latent code\n",
        "\t\t:param train_xr: reconstructed data\n",
        "\t\t:param train_x: training data, the input data\n",
        "\t\t:return: new loss\n",
        "\t\t\"\"\"\n",
        "    'So, we first get the mmd loss'\n",
        "    'First, sample from random noise'\n",
        "    batch_size = K.shape(train_z)[0]\n",
        "    latent_dim = K.int_shape(train_z)[1]\n",
        "    true_samples = K.random_normal(shape=(batch_size, latent_dim), mean=0., stddev=1.)\n",
        "    'calculate mmd loss'\n",
        "    loss_mmd = compute_mmd(true_samples, train_z)\n",
        "\n",
        "    'Then, also get the reconstructed loss'\n",
        "    loss_nll = K.mean(K.square(train_xr - train_x))\n",
        "\n",
        "    'Add them together, then you can get the final loss'\n",
        "    loss = loss_nll + loss_mmd\n",
        "    return loss"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "-znLI5QeUWuL",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "'Time to prepare the data'\n",
        "# Loads the training and test data sets (ignoring class labels)\n",
        "(x_train, _), (x_test, y_test) = mnist.load_data()\n",
        "\n",
        "# Scales the training and test data to range between 0 and 1.\n",
        "max_value = float(x_train.max())\n",
        "x_train = x_train.astype('float32') / max_value\n",
        "x_test = x_test.astype('float32') / max_value\n",
        "x_train = x_train.reshape((len(x_train), 28, 28, 1))\n",
        "x_test = x_test.reshape((len(x_test), 28, 28, 1))"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "3ysa_TrrUlBl",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "'compile the model, prepare for training'\n",
        "loss = custom_loss(train_z, train_xr, encoder_input)\n",
        "autoencoder.add_loss(loss)\n",
        "autoencoder.compile(optimizer='adam')"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "t1eWneXGZ713",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "'Train the model'\n",
        "autoencoder.fit(x_train,\n",
        "\t\t\t\tepochs=epoch,\n",
        "\t\t\t\tbatch_size=batch_size,\n",
        "\t\t\t\tshuffle=True,\n",
        "\t\t\t\tvalidation_data=(x_test, None)\n",
        "\t\t\t\t)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "ttOUsiTYcneY",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "# Convert a numpy array of shape [batch_size, height, width, 1] into a displayable array \n",
        "# of shape [height*sqrt(batch_size, width*sqrt(batch_size))] by tiling the images\n",
        "def convert_to_display(samples):\n",
        "    cnt, height, width = int(math.floor(math.sqrt(samples.shape[0]))), samples.shape[1], samples.shape[2]\n",
        "    samples = np.transpose(samples, axes=[1, 0, 2, 3])\n",
        "    samples = np.reshape(samples, [height, cnt, cnt, width])\n",
        "    samples = np.transpose(samples, axes=[1, 0, 2, 3])\n",
        "    samples = np.reshape(samples, [height*cnt, width*cnt])\n",
        "    return samples"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "3uzscFTqdihh",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "'Time to evaluate'\n",
        "random_noise = np.random.normal(size=(100, z_dim))\n",
        "samples = decoder.predict(random_noise)\n",
        "plt.imshow(convert_to_display(samples), cmap='Greys_r')\n",
        "plt.show()"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "5DDNIFwJemNL",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "# If latent z is 2-dimensional we visualize it by plotting latent z of different digits in different colors\n",
        "if z_dim == 2:\n",
        "    z = encoder.predict(x_test)\n",
        "    label = y_test\n",
        "    plt.scatter(z[:, 0], z[:, 1], c=label)\n",
        "    plt.colorbar()\n",
        "    plt.show()"
      ],
      "execution_count": 0,
      "outputs": []
    }
  ]
}