{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "&emsp;&emsp;&emsp;&emsp;&emsp;\n",
    "&emsp;&emsp;&emsp;&emsp;&emsp;\n",
    "&emsp;&emsp;&emsp;&emsp;&emsp;\n",
    "&emsp;&emsp;&emsp;&emsp;&emsp;\n",
    "&emsp;&emsp;&emsp;&emsp;&emsp;\n",
    "&emsp;&emsp;&emsp;&ensp;\n",
    "[Home Page](../Start_Here.ipynb)\n",
    "\n",
    "\n",
    "[Previous Notebook](Part4.ipynb)\n",
    "&emsp;&emsp;&emsp;&emsp;&emsp;&emsp;\n",
    "&emsp;&emsp;&emsp;&emsp;&emsp;\n",
    "&emsp;&emsp;&emsp;&emsp;&emsp;\n",
    "&emsp;&emsp;&emsp;&emsp;&emsp;\n",
    "[1](Start_Here.ipynb)\n",
    "[2](Part2.ipynb)\n",
    "[3](Part3.ipynb)\n",
    "[4](Part4.ipynb)\n",
    "[5]\n",
    "&emsp;&emsp;&emsp;&emsp;&emsp;\n",
    "&emsp;&emsp;&emsp;&emsp;&emsp;\n",
    "&emsp;&emsp;&emsp;&emsp;&emsp;\n",
    "&emsp;&emsp;&emsp;&emsp;&emsp;\n",
    "&emsp;&emsp;&emsp;&emsp;&emsp;\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Exercise  :\n",
    "\n",
    "Now with that, you are introduced to the Problem Statement and understood the ways you need to solve it. Now is your time to Tweak, Tune and get your hands dirty with the model. \n",
    "\n",
    "Let us help you get started by pointing out some ways in which you can make the model more efficient. \n",
    "\n",
    "- Epochs: Yes, you can increase the number of epochs to train more on the dataset, but make sure you don't overfit it.\n",
    "- Learning Rate: This is a very critical thing to understand and tweak as lower values make learning slow, and higher values stop the model from learning.\n",
    "- Depth of the Model: We had a depth of 5 upsampling and 5 downsampling blocks, we can reduce and increase to learn complex functions.\n",
    "- Dropout Rate: This would regularize the model from overfitting and learn the features across all neurons properly.\n",
    "- Gatedness of the Model: You can set the parameter `gated` to `True` or `False` to see how it impacts the model.\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "Don't remember what an Epoch means? Worry not, head to our [Introduction to Deep Learning](../Start_Here/Intro_to_DL.ipynb) notebook to understand what an Epoch means and what the other terms mean.\n",
    "\n",
    "\n",
    "Note, before you start tweaking and training your model, it would be worthwhile to refer to these to see how they affect your model: \n",
    "\n",
    "[Epochs impact on overfitting](https://datascience.stackexchange.com/questions/27561/can-the-number-of-epochs-influence-overfitting ) \n",
    "\n",
    "[Depth of the Model](https://www.quora.com/Does-adding-more-layers-always-result-in-more-accuracy-in-convolutional-neural-networks)\n",
    "\n",
    "[Understand the Impact of Learning Rate on Neural Network Performance](https://machinelearningmastery.com/understand-the-dynamics-of-learning-rate-on-deep-learning-neural-networks/)\n",
    "\n",
    "[Introduction to Dropout for Regularizing Deep Neural Networks](https://machinelearningmastery.com/dropout-for-regularizing-deep-neural-networks/)\n",
    "\n",
    "[Introduction to Optimizers](https://algorithmia.com/blog/introduction-to-optimizers)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Import Necessary Libraries\n",
    "from __future__ import print_function\n",
    "\n",
    "import sys\n",
    "sys.path.append('/workspace/python/source_code')\n",
    "\n",
    "import numpy as np \n",
    "import time\n",
    "import importlib\n",
    "import os\n",
    "import math\n",
    "import matplotlib.pyplot as plt\n",
    "from matplotlib.pyplot import imshow\n",
    "\n",
    "# TensorFlow and tf.keras\n",
    "import tensorflow as tf\n",
    "from tensorflow.keras.models import *\n",
    "from tensorflow.keras.layers import *\n",
    "from tensorflow.keras.optimizers import *\n",
    "from tensorflow.keras.activations import *\n",
    "from tensorflow.keras.callbacks import ModelCheckpoint, LearningRateScheduler\n",
    "from tensorflow.keras import backend as keras\n",
    "from tensorflow.keras.preprocessing import image\n",
    "from tensorflow.keras.applications.imagenet_utils import preprocess_input\n",
    "from tensorflow.keras.initializers import glorot_uniform\n",
    "import tensorflow.keras.backend as K\n",
    "\n",
    "# Custom Utlities\n",
    "import model.flow_architecture as flow_architecture\n",
    "import utils.data_utils as data_utils\n",
    "\n",
    "import os\n",
    "os.environ[\"CUDA_VISIBLE_DEVICES\"]=\"0\"\n",
    "# reload(data_utils) # you need to execute this in case you modify the plotting scripts in data_utils"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "batch_size = 8\n",
    "dataset_size = 3000   # Number of elements in the train.tfrecords\n",
    "validation_size = 256 # Number of elements to use for validation\n",
    "\n",
    "# derive some quantities\n",
    "train_size = dataset_size - validation_size\n",
    "train_batches = int(train_size / batch_size)\n",
    "validation_batches= int(validation_size / batch_size)\n",
    "\n",
    "test_size = 28\n",
    "test_batches = int(test_size/batch_size)\n",
    "print('Number of batches in train/validation/test dataset:', train_batches, '/', validation_batches, '/', test_batches)\n",
    "\n",
    "def init_datasets():\n",
    "    dataset = tf.data.TFRecordDataset('data/train.tfrecords')\n",
    "    # Transform binary data into image arrays\n",
    "    dataset = dataset.map(data_utils.parse_flow_data)\n",
    "    \n",
    "    training_dataset = dataset.skip(validation_size).shuffle(buffer_size=512)\n",
    "    training_dataset = training_dataset.batch(batch_size, drop_remainder=True)\n",
    "#     training_dataset = training_dataset.repeat()\n",
    "\n",
    "    validation_dataset = dataset.take(validation_size).batch(batch_size, drop_remainder=True)\n",
    "#     validation_dataset = validation_dataset.repeat()\n",
    "\n",
    "    # Read test dataset\n",
    "    test_dataset = tf.data.TFRecordDataset('data/test.tfrecords')\n",
    "    test_dataset = test_dataset.map(data_utils.parse_flow_data) # Transform binary data into image arrays\n",
    "    test_dataset = test_dataset.batch(batch_size, drop_remainder = True)\n",
    " \n",
    "    return training_dataset, validation_dataset, test_dataset\n",
    "\n",
    "training_dataset, validation_dataset, test_dataset = init_datasets()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def loss_image(vflow_hat, vflow):\n",
    "    ''' Defines the loss for the predicted flow.\n",
    "    \n",
    "    Arguments:\n",
    "    vflow_hat -- predicted flow, shape (?, nh, nw, 2)\n",
    "    vflow   -- target flow from the simulation, shape (?, nh, nw, 2)\n",
    "    \n",
    "    Returns: the L2 loss\n",
    "    '''\n",
    "    ### define the square error loss (~ 1 line of code)\n",
    "    loss = tf.nn.l2_loss(vflow_hat - vflow)\n",
    "    ###\n",
    "                         \n",
    "    # Add a scalar to tensorboard\n",
    "    tf.summary.scalar('loss', loss)\n",
    "    \n",
    "    return loss"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def concat_elu(x):\n",
    "    \"\"\" like concatenated ReLU (http://arxiv.org/abs/1603.05201), but then with ELU \"\"\"\n",
    "    axis = len(x.get_shape())-1\n",
    "    return elu(concatenate([x, -x], axis=axis))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Helper Function to Set Non-linearity\n",
    "def set_nonlinearity(name):\n",
    "    if name == 'concat_elu':\n",
    "        return concat_elu\n",
    "    elif name == 'elu':\n",
    "        return tf.nn.elu\n",
    "    elif name == 'concat_relu':\n",
    "        return tf.nn.crelu\n",
    "    elif name == 'relu':\n",
    "        return tf.nn.relu\n",
    "    else:\n",
    "        raise('nonlinearity ' + name + ' is not supported')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def res_block(x, a=None, filter_size=16, nonlinearity=concat_elu, rate=0, stride=1, gated=False, name=\"resnet\"):\n",
    "    \"\"\" Residual block of 3x3 convolutions \"\"\"\n",
    "    # Copy of our Input\n",
    "    orig_x = x\n",
    "    # Get Shape of Input data\n",
    "    orig_x_int_shape = flow_architecture.int_shape(x)\n",
    "    \n",
    "    #### First Convolution Layer\n",
    "    # If Input has one channel Data ( i.e., Input is our Input Image )\n",
    "    if orig_x_int_shape[3] == 1:\n",
    "        x_1 = flow_architecture.conv_layer(x, 3, stride, filter_size, name + '_conv_1')\n",
    "    else:\n",
    "        x_1 = flow_architecture.conv_layer(nonlinearity(x), 3, stride, filter_size, name + '_conv_1')\n",
    "    \n",
    "    # a is fed during the Up-sampling part of the Network ( Refer Upsampling block )\n",
    "    if a is not None:\n",
    "        shape_a = flow_architecture.int_shape(a)\n",
    "        shape_x_1 = flow_architecture.int_shape(x_1)\n",
    "        paddings = [[0,0],[0, shape_x_1[1]-shape_a[1]], [0, shape_x_1[2]-shape_a[2]], [0, 0]]\n",
    "        a = tf.pad(a, paddings)\n",
    "        x_1 = x_1 + flow_architecture.nin(nonlinearity(a), filter_size, name + '_nin')\n",
    "    # Add Activation Function and Dropout\n",
    "    x_1 = nonlinearity(x_1)\n",
    "    x_1 = Dropout(rate=rate)(x_1)\n",
    "    \n",
    "    #### Second Convolution Layer \n",
    "    # Implemented Gated Residual Blocks \n",
    "    if not gated:\n",
    "        x_2 = flow_architecture.conv_layer(x_1, 3, 1, filter_size, name + '_conv_2')\n",
    "    else:\n",
    "        x_2 = flow_architecture.conv_layer(x_1, 3, 1, filter_size*2, name + '_conv_2')\n",
    "        x_2_1, x_2_2 = tf.split(axis=3,num_or_size_splits=2,value=x_2)\n",
    "        x_2 = x_2_1 * tf.nn.sigmoid(x_2_2)\n",
    "    \n",
    "    # During Down-sampling Apply Pooling layer for the Input to Match the Outout\n",
    "    if int(orig_x.get_shape()[2]) > int(x_2.get_shape()[2]):\n",
    "        orig_x = tf.nn.avg_pool(orig_x, [1,2,2,1], [1,2,2,1], padding='SAME')\n",
    "\n",
    "    # Pad Input Data\n",
    "    out_filter = filter_size\n",
    "    in_filter = int(orig_x.get_shape()[3])\n",
    "    if out_filter != in_filter:\n",
    "        orig_x = tf.pad( orig_x, [[0, 0], [0, 0], [0, 0], [(out_filter-in_filter), 0]])\n",
    "    # Output Input Data + Output of Convolution Layer ( Why ? Because this is a Residual Block )\n",
    "    return orig_x + x_2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def downsampling_res_blocks(x, nr_res_blocks, filter_size, nonlinearity, rate, gated,name_prefix, downsample):\n",
    "    ''' An optional downsampling step followed by one (or more) residual blocks '''\n",
    "    \n",
    "    # Set Parameters and Call our Residual Block Function \n",
    "    if downsample:\n",
    "        x = res_block(x, filter_size=filter_size, nonlinearity=nonlinearity, rate=rate, stride=2,gated=gated, name=name_prefix + \"downsample\")\n",
    "    for i in range(nr_res_blocks):\n",
    "        x = res_block(x, filter_size=filter_size, nonlinearity=nonlinearity, stride=1,rate=rate, gated=gated, name=name_prefix + str(i))      \n",
    "    return x    \n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def upsumpling_res_blocks(x, nr_res_blocks, filter_size, nonlinearity, rate, gated, \n",
    "                        name_prefix, a):\n",
    "    ''' Upsampling followed by a residual block '''\n",
    "    # Set Parameters and Call our Residual Block Functions\n",
    "    x = flow_architecture.transpose_conv_layer(x, 3, 2, filter_size, name_prefix)\n",
    "    for i in range(nr_res_blocks):\n",
    "        if i == 0:\n",
    "            x = res_block(x, a, filter_size=filter_size, nonlinearity=nonlinearity, rate=rate, gated=gated, name=name_prefix + str(i))\n",
    "        else:\n",
    "            x = res_block(x, filter_size=filter_size, nonlinearity=nonlinearity, rate=rate, gated=gated, name=name_prefix + str(i))\n",
    "\n",
    "    return x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def conv_res(inputs, nr_res_blocks=1, rate=1.0, nonlinearity_name='concat_elu', gated=True,depth=5):\n",
    "    \"\"\"Builds conv part of net.\n",
    "    Args:\n",
    "      inputs: input images\n",
    "      rate: dropout layer\n",
    "    \"\"\"\n",
    "    # Set Non-linearity\n",
    "    nonlinearity = set_nonlinearity(nonlinearity_name)\n",
    "    filter_size = 8\n",
    "    # Store for Concatenation of the Downsampling output with the Upsampling blocks\n",
    "    a = []\n",
    "    \n",
    "    # First Downsampling Residual Block to Convert Input Image from ( 128  ,256 ,1 ) to ( 128 , 256 , 8)\n",
    "    \n",
    "    x = downsampling_res_blocks(inputs, nr_res_blocks, filter_size, nonlinearity, rate, gated, \"resnet_1_\", False)\n",
    "    \n",
    "    # Loop Through Downsampling Blocks \n",
    "    for i in range(2,1+depth):\n",
    "        a.append(x)\n",
    "        filter_size = 2 * filter_size\n",
    "        name_prefix = \"resnet\" + str(i) + \"_\"\n",
    "        x = downsampling_res_blocks(x, nr_res_blocks, filter_size, nonlinearity, rate, gated, name_prefix, True)\n",
    "\n",
    "    # Loop Through Up-sampling Blocks \n",
    "    for i in range(1,depth):\n",
    "        filter_size = int(filter_size /2)\n",
    "        name_prefix = \"up_conv_\" + str(i)\n",
    "        x = upsumpling_res_blocks(x, nr_res_blocks, filter_size, nonlinearity, rate, gated, name_prefix, a[-i])\n",
    "    \n",
    "    # Last Convolution Layer with Activation \n",
    "    x = flow_architecture.conv_layer(x, 3, 1, 2, \"last_conv\")\n",
    "    x = tf.nn.tanh(x)\n",
    "\n",
    "    return x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def model(boundary, rate,gated=True,depth=5):\n",
    "    return conv_res(boundary, nr_res_blocks=1, rate=rate, nonlinearity_name='concat_elu', gated=gated,depth=depth)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define Dropout Rate and Set Gated = False\n",
    "\n",
    "## TODO \n",
    "## ~~ Change parameters Here ~~\n",
    "rate = 0.3\n",
    "gated = False\n",
    "depth = 5\n",
    "lr = 0.0001\n",
    "## ~~ Change parameters Here ~~\n",
    "\n",
    "# Compile the Model\n",
    "input = tf.keras.Input(shape=(128,256,1), name=\"boundary\")\n",
    "output = model(input,rate=rate,gated=gated,depth=depth)\n",
    "unet = tf.keras.Model(inputs = input, outputs=output)\n",
    "unet.compile(tf.keras.optimizers.Adam(lr), loss=loss_image)\n",
    "unet.summary()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Let's train our Model for 20 Epochs\n",
    "results = unet.fit(training_dataset, epochs=15)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Let us Plot the train History\n",
    "data_utils.plot_keras_loss(results)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Let's Test our Model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "test_loss = unet.evaluate(test_dataset, steps=3)\n",
    "print('The loss over the test dataset', test_loss)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "## Let's Test our Model\n",
    "\n",
    "test_loss = unet.evaluate(test_dataset, steps=3)\n",
    "print('The loss over the test dataset', test_loss)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "x, vxy = data_utils.load_test_data(1) # you can try different numbers between 1 and 28\n",
    "vxy_hat = unet.predict(x)\n",
    "data_utils.plot_test_result(x, vxy, vxy_hat)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Licensing\n",
    "This material is released by NVIDIA Corporation under the Creative Commons Attribution 4.0 International (CC BY 4.0)\n",
    "\n",
    "[Previous Notebook](Part4.ipynb)\n",
    "&emsp;&emsp;&emsp;&emsp;&emsp;&emsp;\n",
    "&emsp;&emsp;&emsp;&emsp;&emsp;\n",
    "&emsp;&emsp;&emsp;&emsp;&emsp;\n",
    "&emsp;&emsp;&emsp;&emsp;&emsp;\n",
    "[1](Start_Here.ipynb)\n",
    "[2](Part2.ipynb)\n",
    "[3](Part3.ipynb)\n",
    "[4](Part4.ipynb)\n",
    "[5]\n",
    "&emsp;&emsp;&emsp;&emsp;&emsp;\n",
    "&emsp;&emsp;&emsp;&emsp;&emsp;\n",
    "&emsp;&emsp;&emsp;&emsp;&emsp;\n",
    "&emsp;&emsp;&emsp;&emsp;&emsp;\n",
    "&emsp;&emsp;&emsp;&emsp;&emsp;\n",
    "\n",
    "\n",
    "&emsp;&emsp;&emsp;&emsp;&emsp;\n",
    "&emsp;&emsp;&emsp;&emsp;&emsp;\n",
    "&emsp;&emsp;&emsp;&emsp;&emsp;\n",
    "&emsp;&emsp;&emsp;&emsp;&emsp;\n",
    "&emsp;&emsp;&emsp;&emsp;&emsp;\n",
    "&emsp;&emsp;&emsp;&ensp;\n",
    "[Home Page](../Start_Here.ipynb)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "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.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
