notebooks/unsupervised/variational-autoencoder/vae-cifar10.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Variational Auto Encoder"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Dataset\n",
    "DATASET = 'cifar10'\n",
    "CHANNELS = 3\n",
    "INPUT_SHAPE = (32, 32, CHANNELS)\n",
    "VALID_SIZE = 0\n",
    "\n",
    "# Optimization\n",
    "LATENT_DIM = 2048\n",
    "REC_LOSS_W = 1\n",
    "KL_LOSS_W = 1\n",
    "\n",
    "# Training\n",
    "EPOCHS = 200\n",
    "BATCH_SIZE = 64\n",
    "LR = 0.001\n",
    "\n",
    "CVS = [32, 64, 128]\n",
    "CV_PARAMS = dict(strides=2, padding='same', activation='relu')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Setup"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import time\n",
    "from collections import namedtuple\n",
    "\n",
    "import numpy as np\n",
    "import tensorflow as tf\n",
    "\n",
    "import seaborn as sns\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "from tensorflow.keras import Model\n",
    "from tensorflow.keras.layers import (Layer, Conv2D, Conv2DTranspose, Dense,\n",
    "                                     Dropout, BatchNormalization,\n",
    "                                     Activation, GlobalAveragePooling2D,\n",
    "                                     Reshape, Flatten)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(\"Num GPUs Available: \", len(tf.config.experimental.list_physical_devices('GPU')))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sns.set_style(\"whitegrid\", {'axes.grid' : False})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def plot(y, titles=None, rows=1, i0=0):\n",
    "    for i, image in enumerate(y):\n",
    "        if image is None:\n",
    "            plt.subplot(rows, len(y), i0+i+1)\n",
    "            plt.axis('off')\n",
    "            continue\n",
    "\n",
    "        t = titles[i] if titles else None\n",
    "        plt.subplot(rows, len(y), i0+i+1, title=t)\n",
    "        plt.imshow(image)\n",
    "        plt.axis('off')\n",
    "\n",
    "def plot_to_image(figure):\n",
    "    import io\n",
    "    buf = io.BytesIO()\n",
    "    plt.savefig(buf, format='png')\n",
    "    plt.close(figure)\n",
    "    buf.seek(0)\n",
    "    image = tf.image.decode_png(buf.getvalue(), channels=4)\n",
    "    image = tf.expand_dims(image, 0)\n",
    "    return image\n",
    "\n",
    "class TensorBoardImage(tf.keras.callbacks.Callback):\n",
    "    def __init__(self, tag, logs_dir):\n",
    "        super().__init__() \n",
    "        self.tag = tag\n",
    "        self.logs_dir = logs_dir\n",
    "\n",
    "    def on_train_begin(self, logs=None):    \n",
    "        self.writer = tf.summary.create_file_writer(self.logs_dir)\n",
    "\n",
    "    def on_epoch_end(self, epoch, logs=None):\n",
    "        d_zu, d_zlv, d_z = encoder.predict(Data.x[:16])\n",
    "        rec = decoder.predict(d_z)\n",
    "        rec = np.expand_dims(np.hstack(rec), 0)\n",
    "        with self.writer.as_default():\n",
    "            tf.summary.image(self.tag, rec, step=epoch)\n",
    "\n",
    "    def on_train_end(self, logs=None):\n",
    "        self.writer.close()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Load Dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "CLASSES = np.asarray('airplane automobile bird cat deer dog frog horse ship truck'.split())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "Data = namedtuple('Data', 'x y xv yv xt yt')\n",
    "\n",
    "def load_data():\n",
    "    (x, y), (xt, yt) = tf.keras.datasets.cifar10.load_data()\n",
    "    x = (x.astype(\"float32\") / 255) # .mean(axis=-1, keepdims=True)\n",
    "    xt = (xt.astype(\"float32\") / 255) # .mean(axis=-1, keepdims=True)\n",
    "    y, yt = y.ravel(), yt.ravel()\n",
    "\n",
    "    if VALID_SIZE:\n",
    "        _valid_samples = int(VALID_SIZE * len(x))\n",
    "\n",
    "        x, xv = x[_valid_samples:], x[:_valid_samples]\n",
    "        y, yv = x[_valid_samples:], x[:_valid_samples]\n",
    "    else:\n",
    "        xv, yv = None, None\n",
    "        \n",
    "    print('Training')\n",
    "    print('  samples:', len(x))\n",
    "    print('  labels examples:', y[:10])\n",
    "    if VALID_SIZE:\n",
    "        print('Validating')\n",
    "        print('  samples:', len(xv))\n",
    "        print('  labels examples:', yv[:10])\n",
    "    print('Testing')\n",
    "    print('  samples:', len(xt))\n",
    "    print('  labels examples:', yt[:10])\n",
    "    \n",
    "    plot(x[:4])\n",
    "    \n",
    "    return Data(x=x, y=y, xv=xv, yv=yv, xt=xt, yt=yt)\n",
    "\n",
    "Data = load_data()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "x_ds = (tf.data\n",
    "        .Dataset.from_tensor_slices(Data.x)\n",
    "        .shuffle(len(Data.x))\n",
    "        .batch(BATCH_SIZE, drop_remainder=True))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Defining Model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "RUN_ID = int(time.time())\n",
    "\n",
    "LOGS = (f'/tf/logs/d:{DATASET} e:{EPOCHS} b:{BATCH_SIZE} lr:{LR} '\n",
    "        f'arch:({\",\".join(map(str, CVS))}) latent:{LATENT_DIM} '\n",
    "        f'rec_w:{round(REC_LOSS_W, 6)} kl_w:{round(KL_LOSS_W, 6)}'\n",
    "        f'/{RUN_ID}')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Sampling(Layer):\n",
    "    def call(self, inputs):\n",
    "        z_mean, z_log_var = inputs\n",
    "        batch = tf.shape(z_mean)[0]\n",
    "        dim = tf.shape(z_mean)[1]\n",
    "        epsilon = tf.keras.backend.random_normal(shape=(batch, dim))\n",
    "        return z_mean + tf.exp(0.5 * z_log_var) * epsilon\n",
    "\n",
    "def encode(x, latent_dim=LATENT_DIM):\n",
    "    y = x\n",
    "    for ix, c in enumerate(CVS):\n",
    "        y = Conv2D(c, 3, name=f'cv{ix}', **CV_PARAMS)(y)\n",
    "    y = Flatten(name='ft')(y)\n",
    "    y = Dense(latent_dim, activation=\"relu\", name='fc1')(y)\n",
    "    zu = Dense(latent_dim, name='zu')(y)\n",
    "    zlv = Dense(latent_dim, name='zlv')(y)\n",
    "\n",
    "    return zu, zlv"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "x = tf.keras.Input(shape=INPUT_SHAPE, name='images')\n",
    "zu, zlv = encode(x)\n",
    "z = Sampling(name='zs')([zu, zlv])\n",
    "\n",
    "encoder = Model(x, [zu, zlv, z], name='encoder')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "encoder.summary()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def decode(z, act='sigmoid'):\n",
    "    size = 4\n",
    "    filters = CVS[-1]\n",
    "\n",
    "    z = Dense(size*size*filters, activation=\"relu\", name='fc1')(z)\n",
    "    z = Reshape((size, size, filters), name='rs')(z)\n",
    "    for ix, c in enumerate(reversed(CVS)):\n",
    "        z = Conv2DTranspose(c, 3, name=f'cvt{ix}', **CV_PARAMS)(z)\n",
    "    z = Conv2DTranspose(CHANNELS, 3, activation=act, padding='same', name='decoded')(z)\n",
    "\n",
    "    return z\n",
    "\n",
    "lvs = tf.keras.Input(shape=(LATENT_DIM,), name='latent_vars')\n",
    "ty = decode(lvs)\n",
    "\n",
    "decoder = Model(lvs, ty, name='decoder')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "decoder.summary()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from tensorflow.keras import metrics, losses\n",
    "\n",
    "class VAE(Model):\n",
    "    def __init__(self, encoder, decoder, **kwargs):\n",
    "        super(VAE, self).__init__(**kwargs)\n",
    "        self.encoder = encoder\n",
    "        self.decoder = decoder\n",
    "        self.total_loss_tracker = metrics.Mean(name=\"total_loss\")\n",
    "        self.reconstruction_loss_tracker = metrics.Mean(name=\"reconstruction_loss\")\n",
    "        self.kl_loss_tracker = metrics.Mean(name=\"kl_loss\")\n",
    "\n",
    "    @property\n",
    "    def metrics(self):\n",
    "        return [\n",
    "            self.total_loss_tracker,\n",
    "            self.reconstruction_loss_tracker,\n",
    "            self.kl_loss_tracker,\n",
    "        ]\n",
    "\n",
    "    def train_step(self, data):\n",
    "        with tf.GradientTape() as tape:\n",
    "            z_mean, z_log_var, z = self.encoder(data)\n",
    "            reconstruction = self.decoder(z)\n",
    "            reconstruction_loss = REC_LOSS_W * tf.reduce_mean(\n",
    "                tf.reduce_sum(losses.binary_crossentropy(data, reconstruction),\n",
    "                              axis=(1, 2))\n",
    "            )\n",
    "            kl_loss = -0.5 * (1 + z_log_var - tf.square(z_mean) - tf.exp(z_log_var))\n",
    "            kl_loss = KL_LOSS_W * tf.reduce_mean(tf.reduce_sum(kl_loss, axis=1))\n",
    "            total_loss = reconstruction_loss + kl_loss\n",
    "\n",
    "        grads = tape.gradient(total_loss, self.trainable_weights)\n",
    "        self.optimizer.apply_gradients(zip(grads, self.trainable_weights))\n",
    "\n",
    "        self.total_loss_tracker.update_state(total_loss)\n",
    "        self.reconstruction_loss_tracker.update_state(reconstruction_loss)\n",
    "        self.kl_loss_tracker.update_state(kl_loss)\n",
    "        \n",
    "        return {\n",
    "            \"loss\": self.total_loss_tracker.result(),\n",
    "            \"reconstruction_loss\": self.reconstruction_loss_tracker.result(),\n",
    "            \"kl_loss\": self.kl_loss_tracker.result(),\n",
    "        }"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model = VAE(encoder, decoder)\n",
    "model.compile(optimizer=tf.keras.optimizers.Adam(learning_rate=LR))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Summary"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "tf.keras.utils.plot_model(encoder, show_shapes=True, rankdir=\"LR\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "tf.keras.utils.plot_model(decoder, show_shapes=True, rankdir=\"LR\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Training"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(\"Num GPUs Available: \", len(tf.config.experimental.list_physical_devices('GPU')))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if os.path.exists(LOGS):\n",
    "    raise ValueError(f'Conflicting logs {LOGS}. Change or delete the target folder.')\n",
    "\n",
    "model.fit(\n",
    "    x_ds,\n",
    "    epochs=EPOCHS,\n",
    "    verbose=2,\n",
    "    callbacks=[\n",
    "        tf.keras.callbacks.EarlyStopping(\n",
    "            monitor='loss',\n",
    "            patience=max(1, EPOCHS // 10),\n",
    "            verbose=1),\n",
    "        tf.keras.callbacks.TerminateOnNaN(),\n",
    "        tf.keras.callbacks.TensorBoard(\n",
    "            LOGS,\n",
    "            histogram_freq=1,\n",
    "            embeddings_freq=3),\n",
    "        TensorBoardImage('reconstruction', LOGS + '/rec')\n",
    "    ]);"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Generating"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "file_writer = tf.summary.create_file_writer(LOGS + '/train')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "\n",
    "SAMPLES = 10\n",
    "\n",
    "def plot_latent_space(vae, n=SAMPLES, figsize=15):\n",
    "    # display a n*n 2D manifold of digits\n",
    "    preds = []\n",
    "    dh, dw = INPUT_SHAPE[:2]\n",
    "    scale = 1.0\n",
    "    figure = np.zeros((INPUT_SHAPE[0] * n, INPUT_SHAPE[1] * n, CHANNELS))\n",
    "    # linearly spaced coordinates corresponding to the 2D plot\n",
    "    # of digit classes in the latent space\n",
    "    grid_x = np.linspace(-scale, scale, n)\n",
    "    grid_y = np.linspace(-scale, scale, n)[::-1]\n",
    "\n",
    "    for i, yi in enumerate(grid_y):\n",
    "        for j, xi in enumerate(grid_x):\n",
    "            z_sample = np.array([[xi, yi] + [0] * (LATENT_DIM - 2)])\n",
    "            x_decoded = vae.decoder.predict(z_sample)\n",
    "            digit = x_decoded[0].reshape(INPUT_SHAPE)\n",
    "            figure[\n",
    "                i * dh : (i + 1) * dw,\n",
    "                j * dh : (j + 1) * dh,\n",
    "            ] = digit\n",
    "            preds.append(x_decoded)\n",
    "\n",
    "    plt_fig = plt.figure(figsize=(figsize, figsize))\n",
    "    plt.imshow(figure, cmap=\"Greys_r\")\n",
    "    plt.show()\n",
    "\n",
    "    return figure.reshape(1, *figure.shape)\n",
    "\n",
    "predictions = plot_latent_space(model)\n",
    "\n",
    "with file_writer.as_default():\n",
    "    tf.summary.image(f'{SAMPLES**2} samples generated', predictions, step=0);"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from sklearn.decomposition import PCA\n",
    "\n",
    "SAMPLES_PLOTTED = 10000\n",
    "\n",
    "def plot_label_clusters(data, labels):\n",
    "    zu, _, _ = encoder.predict(data[:SAMPLES_PLOTTED])\n",
    "    zu = PCA(n_components=2).fit_transform(zu)\n",
    "\n",
    "    fig = plt.figure(figsize=(12, 10))\n",
    "    sns.scatterplot(x=zu[:, 0], y=zu[:, 1], hue=labels[:SAMPLES_PLOTTED])\n",
    "    plt.xlabel(\"Pz[0]\")\n",
    "    plt.ylabel(\"Pz[1]\")\n",
    "    plt.show()\n",
    "\n",
    "    return fig\n",
    "\n",
    "xp, yp = ((Data.xv, Data.yv) if VALID_SIZE else (Data.x, Data.y))\n",
    "clu_fig = plot_label_clusters(xp, CLASSES[yp])\n",
    "plt.show()\n",
    "\n",
    "with file_writer.as_default():\n",
    "    tf.summary.image(\"PCA(z)\", plot_to_image(clu_fig), step=0);"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import warnings\n",
    "from math import ceil\n",
    "from itertools import combinations\n",
    "\n",
    "import pandas as pd\n",
    "\n",
    "SAMPLES = 10000\n",
    "DIMS = 8\n",
    "\n",
    "def plot_label_pairs_clusters(data, labels):\n",
    "    zu, _, _ = encoder.predict(data[:SAMPLES])\n",
    "    dim = zu.shape[1]\n",
    "\n",
    "    d = pd.DataFrame(zu[:, :DIMS])\n",
    "    d['labels'] = labels[:SAMPLES]\n",
    "    return sns.pairplot(d, hue='labels')\n",
    "\n",
    "xp, yp = ((Data.xv, Data.yv) if VALID_SIZE else (Data.x, Data.y))\n",
    "g = plot_label_pairs_clusters(xp, CLASSES[yp])\n",
    "plt.show()\n",
    "\n",
    "with file_writer.as_default():\n",
    "    tf.summary.image(\"latent_variables\", plot_to_image(g.fig), step=0);"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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.6.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}