VAE/01-VAE-with-MNIST.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<img width=\"800px\" src=\"../fidle/img/00-Fidle-header-01.svg\"></img>\n",
    "\n",
    "# <!-- TITLE --> [VAE1] - First VAE, using functional API (MNIST dataset)\n",
    "<!-- DESC --> Construction and training of a VAE, using functional APPI, with a latent space of small dimension.\n",
    "\n",
    "<!-- AUTHOR : Jean-Luc Parouty (CNRS/SIMaP) -->\n",
    "\n",
    "## Objectives :\n",
    " - Understanding and implementing a **variational autoencoder** neurals network (VAE)\n",
    " - Understanding **Keras functional API**\n",
    "\n",
    "The calculation needs being important, it is preferable to use a very simple dataset such as MNIST to start with.  \n",
    "...MNIST with a small scale if you haven't a GPU ;-)\n",
    "\n",
    "## What we're going to do :\n",
    "\n",
    " - Defining a VAE model\n",
    " - Build the model\n",
    " - Train it\n",
    " - Have a look on the train process\n",
    "\n",
    "## Acknowledgements :\n",
    "Thanks to **François Chollet** who is at the base of this example (and the creator of Keras !!).  \n",
    "See : https://keras.io/examples/generative/vae\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Step 1 - Init python stuff"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import tensorflow as tf\n",
    "from tensorflow import keras\n",
    "from tensorflow.keras import layers\n",
    "from tensorflow.keras.callbacks import TensorBoard\n",
    "\n",
    "from modules.layers    import SamplingLayer, VariationalLossLayer\n",
    "from modules.callbacks import ImagesCallback, BestModelCallback\n",
    "from modules.datagen   import MNIST\n",
    "\n",
    "import sys\n",
    "sys.path.append('..')\n",
    "import fidle.pwk as pwk\n",
    "\n",
    "run_dir = './run/VAE1.001'\n",
    "datasets_dir = pwk.init('VAE1', run_dir)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Step 2 - Parameters\n",
    "`scale` : With scale=1, we need 1'30s on a GPU V100 ...and >20' on a CPU !\\\n",
    "`latent_dim` : 2 dimensions is small, but usefull to draw !\\\n",
    "`fit_verbosity` : verbosity during training : 0 = silent, 1 = progress bar, 2 = one line per epoch\n",
    "\n",
    "\n",
    "`loss_weights` : Our **loss function** is the weighted sum of two loss:\n",
    " - `r_loss` which measures the loss during reconstruction.  \n",
    " - `kl_loss` which measures the dispersion.  \n",
    "\n",
    "The weights are defined by: `loss_weights=[k1,k2]` where : `total_loss = k1*r_loss + k2*kl_loss`  \n",
    "In practice, a value of \\[1,.001\\] gives good results here.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "latent_dim    = 2\n",
    "loss_weights  = [1,.001]\n",
    "\n",
    "scale         = 1\n",
    "seed          = 123\n",
    "\n",
    "batch_size    = 64\n",
    "epochs        = 10\n",
    "fit_verbosity = 1"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Override parameters (batch mode) - Just forget this cell"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pwk.override('latent_dim', 'loss_weights', 'scale', 'seed', 'batch_size', 'epochs', 'fit_verbosity')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Step 3 - Prepare data\n",
    "`MNIST.get_data()` return : `x_train,y_train, x_test,y_test`,  \\\n",
    "but we only need x_train for our training."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "x_data, y_data, _,_ = MNIST.get_data(seed=seed, scale=scale, train_prop=1 )\n",
    "\n",
    "pwk.plot_images(x_data[:20], None, indices='all', columns=10, x_size=1,y_size=1,y_padding=0, save_as='01-original')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Step 4 - Build model\n",
    "In this example, we will use the **functional API.**  \n",
    "For this, we will use two custom layers :\n",
    " - `SamplingLayer`, which generates a vector z from the parameters z_mean and z_log_var - See : [SamplingLayer.py](./modules/layers/SamplingLayer.py)\n",
    " - `VariationalLossLayer`, which allows us to calculate the loss function, loss - See : [VariationalLossLayer.py](./modules/layers/VariationalLossLayer.py)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Encoder"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "inputs    = keras.Input(shape=(28, 28, 1))\n",
    "x         = layers.Conv2D(32, 3, strides=1, padding=\"same\", activation=\"relu\")(inputs)\n",
    "x         = layers.Conv2D(64, 3, strides=2, padding=\"same\", activation=\"relu\")(x)\n",
    "x         = layers.Conv2D(64, 3, strides=2, padding=\"same\", activation=\"relu\")(x)\n",
    "x         = layers.Conv2D(64, 3, strides=1, padding=\"same\", activation=\"relu\")(x)\n",
    "x         = layers.Flatten()(x)\n",
    "x         = layers.Dense(16, activation=\"relu\")(x)\n",
    "\n",
    "z_mean    = layers.Dense(latent_dim, name=\"z_mean\")(x)\n",
    "z_log_var = layers.Dense(latent_dim, name=\"z_log_var\")(x)\n",
    "z         = SamplingLayer()([z_mean, z_log_var])\n",
    "\n",
    "encoder = keras.Model(inputs, [z_mean, z_log_var, z], name=\"encoder\")\n",
    "# encoder.summary()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Decoder"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "inputs  = keras.Input(shape=(latent_dim,))\n",
    "x       = layers.Dense(7 * 7 * 64, activation=\"relu\")(inputs)\n",
    "x       = layers.Reshape((7, 7, 64))(x)\n",
    "x       = layers.Conv2DTranspose(64, 3, strides=1, padding=\"same\", activation=\"relu\")(x)\n",
    "x       = layers.Conv2DTranspose(64, 3, strides=2, padding=\"same\", activation=\"relu\")(x)\n",
    "x       = layers.Conv2DTranspose(32, 3, strides=2, padding=\"same\", activation=\"relu\")(x)\n",
    "outputs = layers.Conv2DTranspose(1,  3, padding=\"same\", activation=\"sigmoid\")(x)\n",
    "\n",
    "decoder = keras.Model(inputs, outputs, name=\"decoder\")\n",
    "\n",
    "# decoder.summary()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### VAE\n",
    "\n",
    "We will calculate the loss with a specific layer: `VariationalLossLayer` - See : [VariationalLossLayer.py](./modules/layers/VariationalLossLayer.py)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "inputs = keras.Input(shape=(28, 28, 1))\n",
    "\n",
    "z_mean, z_log_var, z = encoder(inputs)\n",
    "outputs              = decoder(z)\n",
    "\n",
    "outputs = VariationalLossLayer(loss_weights=loss_weights)([inputs, z_mean, z_log_var, outputs])\n",
    "\n",
    "vae=keras.Model(inputs,outputs)\n",
    "\n",
    "vae.compile(optimizer='adam', loss=None)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Step 5 - Train\n",
    "### 5.1 - Using two nice custom callbacks :-)\n",
    "Two custom callbacks are used:\n",
    " - `ImagesCallback` : qui va sauvegarder des images durant l'apprentissage - See [ImagesCallback.py](./modules/callbacks/ImagesCallback.py)\n",
    " - `BestModelCallback` : qui sauvegardera le meilleur model - See [BestModelCallback.py](./modules/callbacks/BestModelCallback.py)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "callback_images      = ImagesCallback(x=x_data, z_dim=latent_dim, nb_images=5, from_z=True, from_random=True, run_dir=run_dir)\n",
    "callback_bestmodel   = BestModelCallback( run_dir + '/models/best_model.h5' )\n",
    "callback_tensorboard = TensorBoard(log_dir=run_dir + '/logs', histogram_freq=1)\n",
    "\n",
    "callbacks_list = [callback_images, callback_bestmodel]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 5.2 - Let's train !\n",
    "With `scale=1`, need 1'15 on a GPU (V100 at IDRIS) ...or 20' on a CPU  "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pwk.chrono_start()\n",
    "\n",
    "history = vae.fit(x_data, epochs=epochs, batch_size=batch_size, callbacks=callbacks_list, verbose=fit_verbosity)\n",
    "\n",
    "pwk.chrono_show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Step 6 - Training review\n",
    "### 6.1 - History"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pwk.plot_history(history,  plot={\"Loss\":['loss','r_loss', 'kl_loss']}, save_as='history')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 6.2 - Reconstruction during training\n",
    "At the end of each epoch, our callback saved some reconstructed images.  \n",
    "Where :  \n",
    "Original image -> encoder -> z -> decoder -> Reconstructed image"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "images_z, images_r = callback_images.get_images( range(0,epochs,2) )\n",
    "\n",
    "pwk.subtitle('Original images :')\n",
    "pwk.plot_images(x_data[:5], None, indices='all', columns=5, x_size=2,y_size=2, save_as=None)\n",
    "\n",
    "pwk.subtitle('Encoded/decoded images')\n",
    "pwk.plot_images(images_z, None, indices='all', columns=5, x_size=2,y_size=2, save_as='02-reconstruct')\n",
    "\n",
    "pwk.subtitle('Original images :')\n",
    "pwk.plot_images(x_data[:5], None, indices='all', columns=5, x_size=2,y_size=2, save_as=None)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 6.3 - Generation (latent -> decoder)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pwk.subtitle('Generated images from latent space')\n",
    "pwk.plot_images(images_r, None, indices='all', columns=5, x_size=2,y_size=2, save_as='03-generated')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Annexe - Model Save and reload "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "filename = run_dir+'/models/my_model.h5'\n",
    "\n",
    "# ---- Save our model :\n",
    "\n",
    "vae.save(filename)\n",
    "\n",
    "# ---- Reload it\n",
    "\n",
    "vae = keras.models.load_model(filename, custom_objects={'SamplingLayer': SamplingLayer, 'VariationalLossLayer':VariationalLossLayer})\n",
    "\n",
    "# ---- Retrieve a layer\n",
    "\n",
    "decoder = vae.get_layer('decoder')\n",
    "\n",
    "img = decoder( np.array([[-1,.1]]))\n",
    "pwk.plot_images(np.array(img), x_size=2,y_size=2, save_as='04-example')\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pwk.end()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "---\n",
    "<img width=\"80px\" src=\"../fidle/img/00-Fidle-logo-01.svg\"></img>"
   ]
  }
 ],
 "metadata": {
  "interpreter": {
   "hash": "8e38643e33497db9a306e3f311fa98cb1e65371278ca73ee4ea0c76aa5a4f387"
  },
  "kernelspec": {
   "display_name": "Python 3.9.7 64-bit ('fidle-cpu': conda)",
   "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.9.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}