vae-cifar.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/homes/em626/sbrml/VAEs/vae-venv-cpu1/lib/python3.5/site-packages/tensorflow/python/framework/dtypes.py:516: FutureWarning: Passing (type, 1) or '1type' as a synonym of type is deprecated; in a future version of numpy, it will be understood as (type, (1,)) / '(1,)type'.\n",
      "  _np_qint8 = np.dtype([(\"qint8\", np.int8, 1)])\n",
      "/homes/em626/sbrml/VAEs/vae-venv-cpu1/lib/python3.5/site-packages/tensorflow/python/framework/dtypes.py:517: FutureWarning: Passing (type, 1) or '1type' as a synonym of type is deprecated; in a future version of numpy, it will be understood as (type, (1,)) / '(1,)type'.\n",
      "  _np_quint8 = np.dtype([(\"quint8\", np.uint8, 1)])\n",
      "/homes/em626/sbrml/VAEs/vae-venv-cpu1/lib/python3.5/site-packages/tensorflow/python/framework/dtypes.py:518: FutureWarning: Passing (type, 1) or '1type' as a synonym of type is deprecated; in a future version of numpy, it will be understood as (type, (1,)) / '(1,)type'.\n",
      "  _np_qint16 = np.dtype([(\"qint16\", np.int16, 1)])\n",
      "/homes/em626/sbrml/VAEs/vae-venv-cpu1/lib/python3.5/site-packages/tensorflow/python/framework/dtypes.py:519: FutureWarning: Passing (type, 1) or '1type' as a synonym of type is deprecated; in a future version of numpy, it will be understood as (type, (1,)) / '(1,)type'.\n",
      "  _np_quint16 = np.dtype([(\"quint16\", np.uint16, 1)])\n",
      "/homes/em626/sbrml/VAEs/vae-venv-cpu1/lib/python3.5/site-packages/tensorflow/python/framework/dtypes.py:520: FutureWarning: Passing (type, 1) or '1type' as a synonym of type is deprecated; in a future version of numpy, it will be understood as (type, (1,)) / '(1,)type'.\n",
      "  _np_qint32 = np.dtype([(\"qint32\", np.int32, 1)])\n",
      "/homes/em626/sbrml/VAEs/vae-venv-cpu1/lib/python3.5/site-packages/tensorflow/python/framework/dtypes.py:525: FutureWarning: Passing (type, 1) or '1type' as a synonym of type is deprecated; in a future version of numpy, it will be understood as (type, (1,)) / '(1,)type'.\n",
      "  np_resource = np.dtype([(\"resource\", np.ubyte, 1)])\n",
      "/homes/em626/sbrml/VAEs/vae-venv-cpu1/lib/python3.5/site-packages/tensorboard/compat/tensorflow_stub/dtypes.py:541: FutureWarning: Passing (type, 1) or '1type' as a synonym of type is deprecated; in a future version of numpy, it will be understood as (type, (1,)) / '(1,)type'.\n",
      "  _np_qint8 = np.dtype([(\"qint8\", np.int8, 1)])\n",
      "/homes/em626/sbrml/VAEs/vae-venv-cpu1/lib/python3.5/site-packages/tensorboard/compat/tensorflow_stub/dtypes.py:542: FutureWarning: Passing (type, 1) or '1type' as a synonym of type is deprecated; in a future version of numpy, it will be understood as (type, (1,)) / '(1,)type'.\n",
      "  _np_quint8 = np.dtype([(\"quint8\", np.uint8, 1)])\n",
      "/homes/em626/sbrml/VAEs/vae-venv-cpu1/lib/python3.5/site-packages/tensorboard/compat/tensorflow_stub/dtypes.py:543: FutureWarning: Passing (type, 1) or '1type' as a synonym of type is deprecated; in a future version of numpy, it will be understood as (type, (1,)) / '(1,)type'.\n",
      "  _np_qint16 = np.dtype([(\"qint16\", np.int16, 1)])\n",
      "/homes/em626/sbrml/VAEs/vae-venv-cpu1/lib/python3.5/site-packages/tensorboard/compat/tensorflow_stub/dtypes.py:544: FutureWarning: Passing (type, 1) or '1type' as a synonym of type is deprecated; in a future version of numpy, it will be understood as (type, (1,)) / '(1,)type'.\n",
      "  _np_quint16 = np.dtype([(\"quint16\", np.uint16, 1)])\n",
      "/homes/em626/sbrml/VAEs/vae-venv-cpu1/lib/python3.5/site-packages/tensorboard/compat/tensorflow_stub/dtypes.py:545: FutureWarning: Passing (type, 1) or '1type' as a synonym of type is deprecated; in a future version of numpy, it will be understood as (type, (1,)) / '(1,)type'.\n",
      "  _np_qint32 = np.dtype([(\"qint32\", np.int32, 1)])\n",
      "/homes/em626/sbrml/VAEs/vae-venv-cpu1/lib/python3.5/site-packages/tensorboard/compat/tensorflow_stub/dtypes.py:550: FutureWarning: Passing (type, 1) or '1type' as a synonym of type is deprecated; in a future version of numpy, it will be understood as (type, (1,)) / '(1,)type'.\n",
      "  np_resource = np.dtype([(\"resource\", np.ubyte, 1)])\n"
     ]
    }
   ],
   "source": [
    "import tensorflow as tf\n",
    "import tensorflow_datasets as tfds\n",
    "import tensorflow_probability as tfp\n",
    "\n",
    "tfk = tf.keras\n",
    "tfl = tfk.layers\n",
    "tfd = tfp.distributions\n",
    "tfs = tf.summary\n",
    "\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "from tqdm import tqdm_notebook as tqdm"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "batch_size = 64\n",
    "buffer_size = 500\n",
    "\n",
    "def normalise_image(img):\n",
    "    return tf.cast(img, tf.float32) / 255.\n",
    "\n",
    "train_dataset = tfds.load(name=\"cifar10\", split=tfds.Split.ALL)\n",
    "train_dataset = train_dataset.shuffle(buffer_size=buffer_size).repeat()\n",
    "train_dataset = train_dataset.map(lambda x: x['image'])\n",
    "train_dataset = train_dataset.map(normalise_image)\n",
    "train_dataset = train_dataset.batch(batch_size)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Encoder(tfl.Layer):\n",
    "    \n",
    "    def __init__(self, latent_size, name='encoder', **kwargs):\n",
    "        \n",
    "        self.latent_size = latent_size\n",
    "        \n",
    "        super(Encoder, self).__init__(name=name, **kwargs)\n",
    "        \n",
    "    \n",
    "    def build(self, input_shape):\n",
    "        \n",
    "        self.layers = [\n",
    "            \n",
    "            tfl.Conv2D(filters=64,\n",
    "                       kernel_size=(5, 5),\n",
    "                       padding='same'),\n",
    "            \n",
    "            tf.nn.relu,\n",
    "            \n",
    "            tfl.Conv2D(filters=128,\n",
    "                       kernel_size=(5, 5),\n",
    "                       strides=(2, 2),\n",
    "                       padding='same'),\n",
    "            \n",
    "            tf.nn.relu,\n",
    "            \n",
    "            tfl.Conv2D(filters=256,\n",
    "                       kernel_size=(5, 5),\n",
    "                       strides=(2, 2),\n",
    "                       padding='same'),\n",
    "            \n",
    "            tf.nn.relu,\n",
    "            \n",
    "            tfl.Conv2D(filters=512,\n",
    "                       kernel_size=(5, 5),\n",
    "                       strides=(2, 2),\n",
    "                       padding='same'),\n",
    "            \n",
    "            tf.nn.relu,\n",
    "            \n",
    "            tfl.Reshape((4 * 4 * 512,)),\n",
    "            \n",
    "            tfl.Dense(512),\n",
    "            \n",
    "            tf.nn.relu\n",
    "        ]\n",
    "        \n",
    "        self.loc_head = tfl.Dense(self.latent_size)\n",
    "        \n",
    "        self.log_scale_head = tfl.Dense(self.latent_size)\n",
    "    \n",
    "    \n",
    "    def call(self, tensor):\n",
    "        \n",
    "        for layer in self.layers:\n",
    "            tensor = layer(tensor)\n",
    "            \n",
    "        loc = self.loc_head(tensor)\n",
    "        scale = tf.nn.softplus(self.log_scale_head(tensor))\n",
    "    \n",
    "        self.posterior = tfd.Normal(loc=loc, scale=scale)\n",
    "        \n",
    "        return self.posterior.sample()\n",
    "    \n",
    "    \n",
    "class Decoder(tfl.Layer):\n",
    "    \n",
    "    def __init__(self, name='decoder', **kwargs):\n",
    "        \n",
    "        super(Decoder, self).__init__(name=name, **kwargs)\n",
    "        \n",
    "    \n",
    "    def build(self, input_shape):\n",
    "        \n",
    "        self.layers = [\n",
    "            \n",
    "            tfl.Dense(512),\n",
    "            \n",
    "            tf.nn.relu,\n",
    "            \n",
    "            tfl.Dense(4 * 4 * 512),\n",
    "            \n",
    "            tf.nn.relu,\n",
    "            \n",
    "            tfl.Reshape((4, 4, 512)),\n",
    "            \n",
    "            tfl.Conv2DTranspose(filters=256,\n",
    "                                kernel_size=(5, 5),\n",
    "                                strides=(2, 2),\n",
    "                                padding='same'),\n",
    "            \n",
    "            tf.nn.relu,\n",
    "            \n",
    "            tfl.Conv2DTranspose(filters=128,\n",
    "                                kernel_size=(5, 5),\n",
    "                                strides=(2, 2),\n",
    "                                padding='same'),\n",
    "            \n",
    "            tf.nn.relu,\n",
    "            \n",
    "            tfl.Conv2DTranspose(filters=64,\n",
    "                                kernel_size=(5, 5),\n",
    "                                strides=(2, 2),\n",
    "                                padding='same'),\n",
    "            \n",
    "            tf.nn.relu,\n",
    "            \n",
    "            tfl.Conv2DTranspose(filters=3,\n",
    "                                kernel_size=(5, 5),\n",
    "                                padding='same'),\n",
    "            \n",
    "            tf.nn.sigmoid\n",
    "            \n",
    "        ]\n",
    "    \n",
    "\n",
    "    def call(self, tensor):\n",
    "        \n",
    "        for layer in self.layers:\n",
    "            tensor = layer(tensor)\n",
    "        \n",
    "        return tensor\n",
    "    \n",
    "    \n",
    "class VAE(tfk.Model):\n",
    "    \n",
    "    def __init__(self, latent_size, name='vae', **kwargs):\n",
    "        \n",
    "        self.latent_size = latent_size\n",
    "        \n",
    "        self.log_noise = tf.Variable(0.0)\n",
    "        \n",
    "        super(VAE, self).__init__(name=name, **kwargs)\n",
    "        \n",
    "        \n",
    "    @property\n",
    "    def log_prob(self):\n",
    "        \n",
    "        return tf.reduce_mean(self._log_prob)\n",
    "    \n",
    "    @property\n",
    "    def kl_divergence(self):\n",
    "        \n",
    "        kl_each_latent = tfd.kl_divergence(self.posterior, self.prior)\n",
    "        \n",
    "        kl_each_example = tf.reduce_sum(kl_each_latent, axis=-1)\n",
    "        \n",
    "        return tf.reduce_mean(kl_each_example)\n",
    "    \n",
    "    @property\n",
    "    def posterior(self):\n",
    "        \n",
    "        return self.encoder.posterior\n",
    "    \n",
    "    \n",
    "    def build(self, input_shape):\n",
    "        \n",
    "        self.encoder = Encoder(self.latent_size)\n",
    "        self.decoder = Decoder()\n",
    "\n",
    "        self.prior = tfd.Normal(loc=tf.zeros(self.latent_size),\n",
    "                                scale=tf.ones(self.latent_size))\n",
    "    \n",
    "    \n",
    "    def call(self, tensor):\n",
    "        \n",
    "        latents = self.encoder(tensor)\n",
    "        \n",
    "        loc = self.decoder(latents)\n",
    "        \n",
    "        scale = tf.exp(self.log_noise)\n",
    "        \n",
    "        self.likelihood = tfd.Normal(loc=loc, scale=scale)\n",
    "        self._log_prob = tf.reduce_sum(self.likelihood.log_prob(tensor), axis=(1, 2, 3))\n",
    "        \n",
    "        return loc"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "808a2900318c47ffa1bd013e8e070321",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(IntProgress(value=0, max=1000000), HTML(value='')))"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "train_steps = int(1e6)\n",
    "beta = 1e0\n",
    "learn_rate = 1e-3\n",
    "log_freq = 10\n",
    "\n",
    "vae = VAE(64)\n",
    "\n",
    "optimizer = tfk.optimizers.Adam(learn_rate)\n",
    "\n",
    "train_summary_writer = tfs.create_file_writer('summaries/train/cifar10')\n",
    "\n",
    "with train_summary_writer.as_default():\n",
    "\n",
    "    for batch in tqdm(train_dataset.take(train_steps), total=train_steps):\n",
    "        \n",
    "        with tf.GradientTape() as tape:\n",
    "            \n",
    "            reconstructions = vae(batch)\n",
    "            \n",
    "            likelihood = vae.log_prob\n",
    "            \n",
    "            kl_divergence = tf.reduce_sum(vae.kl_divergence)\n",
    "\n",
    "            neg_elbo = - likelihood + beta * tf.reduce_sum(vae.kl_divergence)\n",
    "\n",
    "        gradients = tape.gradient(neg_elbo, vae.trainable_variables)\n",
    "\n",
    "        optimizer.apply_gradients(zip(gradients, vae.trainable_variables))\n",
    "\n",
    "        if tf.equal(optimizer.iterations % log_freq, 0):\n",
    "\n",
    "            tfs.scalar('ELBO', - neg_elbo, step=optimizer.iterations)\n",
    "            tfs.scalar('Likelihood', likelihood, step=optimizer.iterations)\n",
    "            tfs.scalar('KL-divergence', kl_divergence, step=optimizer.iterations)\n",
    "            tfs.scalar('log-noise', vae.log_noise, step=optimizer.iterations)\n",
    "            tfs.image('Original', batch, step=optimizer.iterations)\n",
    "            tfs.image('Reconstruction', reconstructions, step=optimizer.iterations)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "vae-venv-cpu1",
   "language": "python",
   "name": "vae-venv-cpu1"
  },
  "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.5.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}