{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Variational AutoEncoder (VAE) with CelebA\n",
    "=========================================\n",
    "---\n",
    "Formation Introduction au Deep Learning  (FIDLE) - S. Arias, E. Maldonado, JL. Parouty - CNRS/SARI/DEVLOG - 2020  \n",
    "\n",
    "## Episode 1 - Train a model\n",
    "\n",
    " - Defining a VAE model\n",
    " - Build the model\n",
    " - Train it\n",
    " - Follow the learning process with Tensorboard\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Step 1 - Setup environment\n",
    "### 1.1 - Python stuff"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import tensorflow as tf\n",
    "import numpy as np\n",
    "import os,sys\n",
    "from importlib import reload\n",
    "\n",
    "import modules.vae\n",
    "import modules.data_generator\n",
    "\n",
    "reload(modules.data_generator)\n",
    "reload(modules.vae)\n",
    "\n",
    "from modules.vae  import VariationalAutoencoder\n",
    "from modules.data_generator import DataGenerator\n",
    "\n",
    "sys.path.append('..')\n",
    "import fidle.pwk as ooo\n",
    "reload(ooo)\n",
    "\n",
    "ooo.init()\n",
    "\n",
    "VariationalAutoencoder.about()\n",
    "DataGenerator.about()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 1.2 - The good place"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "place, dataset_dir = ooo.good_place( { 'GRICAD' : f'{os.getenv(\"SCRATCH_DIR\",\"\")}/PROJECTS/pr-fidle/datasets/celeba',\n",
    "                                       'IDRIS'  : f'{os.getenv(\"WORK\",\"\")}/datasets/celeba'    } )\n",
    "\n",
    "# ---- train/test datasets\n",
    "\n",
    "train_dir    = f'{dataset_dir}/clusters-M.train'\n",
    "test_dir     = f'{dataset_dir}/clusters-M.test'"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Step 2 - DataGenerator and validation data\n",
    "Ok, everything's perfect, now let's instantiate our generator for the entire dataset."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "data_gen = DataGenerator(train_dir, 32, k_size=1)\n",
    "x_test   = np.load(f'{test_dir}/images-000.npy')\n",
    "\n",
    "print(f'Data generator : {len(data_gen)} batchs of {data_gen.batch_size} images, or {data_gen.dataset_size} images')\n",
    "print(f'x_test : {len(x_test)} images')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Step 3 - Get VAE model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "tag = f'CelebA.052-M.{os.getenv(\"SLURM_JOB_ID\",\"unknown\")}'\n",
    "\n",
    "input_shape = (192, 160, 3)\n",
    "z_dim       = 200\n",
    "verbose     = 0\n",
    "\n",
    "encoder= [ {'type':'Conv2D',          'filters':32, 'kernel_size':(3,3), 'strides':2, 'padding':'same', 'activation':'relu'},\n",
    "           {'type':'Dropout',         'rate':0.25},\n",
    "           {'type':'Conv2D',          'filters':64, 'kernel_size':(3,3), 'strides':2, 'padding':'same', 'activation':'relu'},\n",
    "           {'type':'Dropout',         'rate':0.25},\n",
    "           {'type':'Conv2D',          'filters':64, 'kernel_size':(3,3), 'strides':2, 'padding':'same', 'activation':'relu'},\n",
    "           {'type':'Dropout',         'rate':0.25},\n",
    "           {'type':'Conv2D',          'filters':64, 'kernel_size':(3,3), 'strides':2, 'padding':'same', 'activation':'relu'},\n",
    "           {'type':'Dropout',         'rate':0.25},\n",
    "         ]\n",
    "\n",
    "decoder= [ {'type':'Conv2DTranspose', 'filters':64, 'kernel_size':(3,3), 'strides':2, 'padding':'same', 'activation':'relu'},\n",
    "           {'type':'Dropout',         'rate':0.25},\n",
    "           {'type':'Conv2DTranspose', 'filters':64, 'kernel_size':(3,3), 'strides':2, 'padding':'same', 'activation':'relu'},\n",
    "           {'type':'Dropout',         'rate':0.25},\n",
    "           {'type':'Conv2DTranspose', 'filters':32, 'kernel_size':(3,3), 'strides':2, 'padding':'same', 'activation':'relu'},\n",
    "           {'type':'Dropout',         'rate':0.25},\n",
    "           {'type':'Conv2DTranspose', 'filters':3,  'kernel_size':(3,3), 'strides':2, 'padding':'same', 'activation':'sigmoid'}\n",
    "         ]\n",
    "\n",
    "vae = modules.vae.VariationalAutoencoder(input_shape    = input_shape, \n",
    "                                         encoder_layers = encoder, \n",
    "                                         decoder_layers = decoder,\n",
    "                                         z_dim          = z_dim, \n",
    "                                         verbose        = verbose,\n",
    "                                         run_tag        = tag)\n",
    "vae.save(model=None)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Step 4 - Compile it"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "optimizer     = tf.keras.optimizers.Adam(1e-4)\n",
    "r_loss_factor = 10000\n",
    "\n",
    "vae.compile(optimizer, r_loss_factor)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Step 5 - Train\n",
    "For 10 epochs, adam optimizer :  \n",
    "- Run time at IDRIS : 1299.77 sec. - 0:21:39\n",
    "- Run time at GRICAD : 2092.77 sec. - 0:34:52"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "epochs            = 20\n",
    "initial_epoch     = 0"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "vae.train(data_generator    = data_gen,\n",
    "          x_test            = x_test,\n",
    "          epochs            = epochs,\n",
    "          initial_epoch     = initial_epoch\n",
    "         )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "----\n",
    "That's all folks !"
   ]
  }
 ],
 "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.7.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}