{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "source": [
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "import tensorflow as tf\n",
    "from tensorflow import keras\n",
    "%matplotlib inline\n",
    "tf.random.set_seed(1)\n",
    "np.random.seed(1)"
   ],
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "Init Plugin\n",
      "Init Graph Optimizer\n",
      "Init Kernel\n"
     ]
    }
   ],
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-10-02T03:24:22.854773Z",
     "start_time": "2021-10-02T03:24:21.767376Z"
    },
    "id": "_G5CWvypQMXr",
    "scrolled": true
   }
  },
  {
   "cell_type": "markdown",
   "source": [
    "## Create data"
   ],
   "metadata": {
    "id": "eikJH7uSQMXu"
   }
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "source": [
    "def my_func(x):\n",
    "    val  = tf.exp(-200*(x-0.25)**2) - tf.exp(-200*(x-0.75)**2)\n",
    "    return val\n",
    "\n",
    "N       = 100\n",
    "x_train = tf.reshape(tf.linspace(0,1,N),(-1,1)) # Want each row to be a sample (since input_dim=1)\n",
    "y_train = my_func(x_train)\n",
    "\n",
    "print(x_train.shape,x_train.dtype)\n",
    "print(y_train.shape,y_train.dtype)\n",
    "\n",
    "plt.figure()\n",
    "plt.plot(x_train,y_train,'-o')\n",
    "plt.show()"
   ],
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "Metal device set to: Apple M1\n",
      "\n",
      "systemMemory: 16.00 GB\n",
      "maxCacheSize: 5.33 GB\n",
      "\n",
      "(100, 1) <dtype: 'float64'>\n",
      "(100, 1) <dtype: 'float64'>\n"
     ]
    },
    {
     "output_type": "stream",
     "name": "stderr",
     "text": [
      "2021-10-01 20:41:32.590745: I tensorflow/core/common_runtime/pluggable_device/pluggable_device_factory.cc:305] Could not identify NUMA node of platform GPU ID 0, defaulting to 0. Your kernel may not have been built with NUMA support.\n",
      "2021-10-01 20:41:32.591032: I tensorflow/core/common_runtime/pluggable_device/pluggable_device_factory.cc:271] Created TensorFlow device (/job:localhost/replica:0/task:0/device:GPU:0 with 0 MB memory) -> physical PluggableDevice (device: 0, name: METAL, pci bus id: <undefined>)\n"
     ]
    },
    {
     "output_type": "display_data",
     "data": {
      "image/png": "[encoded data, 12608 characters]",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     }
    }
   ],
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-10-02T03:24:23.157087Z",
     "start_time": "2021-10-02T03:24:22.999584Z"
    },
    "id": "BeW6_bSVQMXx",
    "scrolled": false
   }
  },
  {
   "cell_type": "markdown",
   "source": [
    "## Define basic network architecture"
   ],
   "metadata": {
    "id": "Zt8VCsvnQMXz"
   }
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "source": [
    "def MLP(Input_Dim=1,Output_Dim=1,Width=15,Depth=8):\n",
    "\n",
    "    Reg_Func      = keras.regularizers.l2\n",
    "    Reg_Param     = 1e-5\n",
    "    Act_Func      = tf.math.tanh\n",
    "\n",
    "\n",
    "    assert Depth > 1, 'Depth of generator must be greater than 1'\n",
    "    \n",
    "    model = tf.keras.Sequential()\n",
    "    \n",
    "    model.add(keras.layers.Dense(Width, input_shape=(Input_Dim,), activation=Act_Func,\n",
    "                                     kernel_initializer='RandomNormal', bias_initializer='RandomNormal',\n",
    "                                     kernel_regularizer=Reg_Func(Reg_Param)))\n",
    "\n",
    "    # Adding remaining hidden layers\n",
    "    if(Depth > 2):\n",
    "        for l in range(Depth - 2):\n",
    "            #model.add(keras.layers.BatchNormalization())\n",
    "            model.add(keras.layers.Dense(Width, activation=Act_Func,\n",
    "                                         kernel_initializer='RandomNormal', bias_initializer='RandomNormal',\n",
    "                                         kernel_regularizer=Reg_Func(Reg_Param)))\n",
    "            \n",
    "    # Adding output layer\n",
    "    model.add(keras.layers.Dense(Output_Dim, activation=None,\n",
    "                                 kernel_initializer='RandomNormal', bias_initializer='RandomNormal',\n",
    "                                 kernel_regularizer=Reg_Func(Reg_Param)))\n",
    "\n",
    "    return model"
   ],
   "outputs": [],
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-10-02T03:24:29.117196Z",
     "start_time": "2021-10-02T03:24:29.038337Z"
    },
    "id": "fe7ZzCh4QMXz",
    "scrolled": true
   }
  },
  {
   "cell_type": "markdown",
   "source": [
    "## Define loss function"
   ],
   "metadata": {
    "id": "_LERV1EKQMX0"
   }
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "source": [
    "mse = keras.losses.MeanSquaredError() # Inbuilt loss\n",
    "\n",
    "# Define custom loss function\n",
    "def my_mse(real,gen):\n",
    "    \n",
    "    real = tf.cast(real, gen.dtype) # make sure both tensors have the same type\n",
    "    loss = tf.reduce_mean(tf.math.squared_difference(real, gen))\n",
    "\n",
    "    return loss\n"
   ],
   "outputs": [],
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-10-02T03:24:31.714329Z",
     "start_time": "2021-10-02T03:24:31.706765Z"
    },
    "id": "01Tl5E_uQMX0",
    "scrolled": true
   }
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "source": [
    "y1 = tf.constant([1.0,2.0,3.0])\n",
    "y2 = tf.constant([3.0,-4.0,-1.0])\n",
    "\n",
    "print(mse(y1,y2))\n",
    "print(my_mse(y1,y2))"
   ],
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "tf.Tensor(18.666668, shape=(), dtype=float32)\n",
      "tf.Tensor(18.666668, shape=(), dtype=float32)\n"
     ]
    }
   ],
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-10-02T03:24:32.624207Z",
     "start_time": "2021-10-02T03:24:32.518779Z"
    },
    "id": "f_Gc_-7Xa7S0"
   }
  },
  {
   "cell_type": "markdown",
   "source": [
    "## Have a look at details here\n",
    "### https://www.tensorflow.org/guide/keras/writing_a_training_loop_from_scratch"
   ],
   "metadata": {
    "id": "3H2hrl6hQMX1"
   }
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "source": [
    "tf.config.list_physical_devices(\n",
    "    device_type=None\n",
    ")"
   ],
   "outputs": [
    {
     "output_type": "execute_result",
     "data": {
      "text/plain": [
       "[PhysicalDevice(name='/physical_device:CPU:0', device_type='CPU'),\n",
       " PhysicalDevice(name='/physical_device:GPU:0', device_type='GPU')]"
      ]
     },
     "metadata": {},
     "execution_count": 6
    }
   ],
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-10-02T03:24:36.834392Z",
     "start_time": "2021-10-02T03:24:36.826185Z"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "source": [
    "tf.keras.backend.clear_session()\n",
    "model      = MLP(1,1,15,8)\n",
    "optimizer  = keras.optimizers.Adam(learning_rate=1e-3) \n",
    "max_epoch  = 3000\n",
    "with tf.device(\"/device:CPU:0\"):\n",
    "    for epoch in range(1,max_epoch+1):\n",
    "        with tf.GradientTape() as tape:\n",
    "            gen_out  = model(x_train)\n",
    "            loss_val = mse(y_train, gen_out) # Using in-built mse\n",
    "            # Adding regularization losses\n",
    "            loss_val += sum(model.losses)\n",
    "\n",
    "        grads = tape.gradient(loss_val, model.trainable_variables)\n",
    "\n",
    "        optimizer.apply_gradients(zip(grads, model.trainable_variables)) # zip used to create an iterator over the tuples\n",
    "\n",
    "        if epoch % 100 == 0 or epoch==max_epoch:\n",
    "            print(f\"Epoch: {epoch}, loss: {loss_val:.2e}\")  "
   ],
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "Epoch: 100, loss: 1.72e-01\n",
      "Epoch: 200, loss: 1.04e-01\n",
      "Epoch: 300, loss: 1.03e-01\n",
      "Epoch: 400, loss: 9.82e-02\n",
      "Epoch: 500, loss: 6.35e-02\n",
      "Epoch: 600, loss: 6.05e-02\n",
      "Epoch: 700, loss: 5.98e-02\n",
      "Epoch: 800, loss: 5.94e-02\n",
      "Epoch: 900, loss: 5.92e-02\n",
      "Epoch: 1000, loss: 5.90e-02\n",
      "Epoch: 1100, loss: 5.89e-02\n",
      "Epoch: 1200, loss: 5.88e-02\n",
      "Epoch: 1300, loss: 5.86e-02\n",
      "Epoch: 1400, loss: 5.85e-02\n",
      "Epoch: 1500, loss: 5.82e-02\n",
      "Epoch: 1600, loss: 5.80e-02\n",
      "Epoch: 1700, loss: 5.75e-02\n",
      "Epoch: 1800, loss: 5.61e-02\n",
      "Epoch: 1900, loss: 1.53e-02\n",
      "Epoch: 2000, loss: 1.27e-03\n",
      "Epoch: 2100, loss: 1.13e-03\n",
      "Epoch: 2200, loss: 1.08e-03\n",
      "Epoch: 2300, loss: 1.04e-03\n",
      "Epoch: 2400, loss: 1.56e-03\n",
      "Epoch: 2500, loss: 9.69e-04\n",
      "Epoch: 2600, loss: 9.37e-04\n",
      "Epoch: 2700, loss: 9.10e-04\n",
      "Epoch: 2800, loss: 8.80e-04\n",
      "Epoch: 2900, loss: 8.55e-04\n",
      "Epoch: 3000, loss: 8.34e-04\n"
     ]
    }
   ],
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-10-02T03:25:33.979691Z",
     "start_time": "2021-10-02T03:24:50.029220Z"
    },
    "id": "1cualRc2RNJ6"
   }
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "source": [
    "plt.figure(figsize=(10,10))\n",
    "plt.plot(x_train,model.predict(x_train),'-',label='NN',linewidth=4)\n",
    "plt.plot(x_train,y_train,'--',label='True',linewidth=4)\n",
    "plt.legend(fontsize=15)\n",
    "plt.show()"
   ],
   "outputs": [
    {
     "output_type": "stream",
     "name": "stderr",
     "text": [
      "2021-10-01 20:42:39.962368: I tensorflow/compiler/mlir/mlir_graph_optimization_pass.cc:176] None of the MLIR Optimization Passes are enabled (registered 2)\n",
      "2021-10-01 20:42:39.964168: W tensorflow/core/platform/profile_utils/cpu_utils.cc:128] Failed to get CPU frequency: 0 Hz\n",
      "2021-10-01 20:42:39.998600: I tensorflow/core/grappler/optimizers/custom_graph_optimizer_registry.cc:112] Plugin optimizer for device_type GPU is enabled.\n"
     ]
    },
    {
     "output_type": "display_data",
     "data": {
      "image/png": "[encoded data, 37772 characters]",
      "text/plain": [
       "<Figure size 720x720 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     }
    }
   ],
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-10-02T03:25:51.800408Z",
     "start_time": "2021-10-02T03:25:50.365398Z"
    },
    "id": "y5tyN2oPQMX3",
    "scrolled": false
   }
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "source": [
    "tf.keras.backend.clear_session()\n",
    "model      = MLP(1,1,15,8)\n",
    "optimizer  = keras.optimizers.Adam(learning_rate=1e-3) \n",
    "max_epoch  = 3000\n",
    "\n",
    "for epoch in range(1,max_epoch+1):\n",
    "    \n",
    "    with tf.GradientTape() as tape:\n",
    "            \n",
    "        gen_out  = model(x_train)\n",
    "        loss_val = my_mse(y_train, gen_out) # Using custom mse\n",
    "\n",
    "        # Adding regularization losses\n",
    "        loss_val += sum(model.losses)\n",
    "            \n",
    "    grads = tape.gradient(loss_val, model.trainable_variables)\n",
    "    \n",
    "    optimizer.apply_gradients(zip(grads, model.trainable_variables)) # zip used to create an iterator over the tuples\n",
    "\n",
    "    if epoch % 100 == 0 or epoch==max_epoch:\n",
    "        print(f\"Epoch: {epoch}, loss: {loss_val:.2e}\")  "
   ],
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "Epoch: 100, loss: 1.75e-01\n",
      "Epoch: 200, loss: 1.03e-01\n",
      "Epoch: 300, loss: 1.01e-01\n",
      "Epoch: 400, loss: 6.66e-02\n",
      "Epoch: 500, loss: 6.12e-02\n",
      "Epoch: 600, loss: 5.98e-02\n",
      "Epoch: 700, loss: 5.94e-02\n",
      "Epoch: 800, loss: 5.92e-02\n",
      "Epoch: 900, loss: 5.90e-02\n",
      "Epoch: 1000, loss: 5.89e-02\n",
      "Epoch: 1100, loss: 5.88e-02\n",
      "Epoch: 1200, loss: 5.87e-02\n",
      "Epoch: 1300, loss: 5.87e-02\n",
      "Epoch: 1400, loss: 5.86e-02\n",
      "Epoch: 1500, loss: 5.87e-02\n",
      "Epoch: 1600, loss: 5.84e-02\n",
      "Epoch: 1700, loss: 5.82e-02\n",
      "Epoch: 1800, loss: 5.93e-02\n",
      "Epoch: 1900, loss: 5.79e-02\n",
      "Epoch: 2000, loss: 5.76e-02\n",
      "Epoch: 2100, loss: 5.73e-02\n",
      "Epoch: 2200, loss: 5.67e-02\n",
      "Epoch: 2300, loss: 5.51e-02\n",
      "Epoch: 2400, loss: 6.83e-03\n",
      "Epoch: 2500, loss: 1.33e-03\n",
      "Epoch: 2600, loss: 1.06e-03\n",
      "Epoch: 2700, loss: 1.00e-03\n",
      "Epoch: 2800, loss: 9.54e-04\n",
      "Epoch: 2900, loss: 9.15e-04\n",
      "Epoch: 3000, loss: 8.82e-04\n"
     ]
    }
   ],
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-10-02T02:06:39.441482Z",
     "start_time": "2021-10-02T02:04:46.414431Z"
    },
    "id": "GRpQP8GuQMX4",
    "scrolled": true
   }
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "source": [
    "plt.figure(figsize=(10,10))\n",
    "plt.plot(x_train,model.predict(x_train),'-',label='NN',linewidth=4)\n",
    "plt.plot(x_train,y_train,'--',label='True',linewidth=4)\n",
    "plt.legend(fontsize=15)\n",
    "plt.show()"
   ],
   "outputs": [
    {
     "output_type": "stream",
     "name": "stderr",
     "text": [
      "2021-10-01 19:06:39.495984: I tensorflow/core/grappler/optimizers/custom_graph_optimizer_registry.cc:112] Plugin optimizer for device_type GPU is enabled.\n"
     ]
    },
    {
     "output_type": "display_data",
     "data": {
      "image/png": "[encoded data, 38540 characters]",
      "text/plain": [
       "<Figure size 720x720 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     }
    }
   ],
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-10-02T02:06:39.666449Z",
     "start_time": "2021-10-02T02:06:39.443819Z"
    },
    "id": "B6S8GHU5fs1B"
   }
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "source": [],
   "outputs": [],
   "metadata": {
    "id": "rI_V9lF3f0Rv"
   }
  }
 ],
 "metadata": {
  "colab": {
   "collapsed_sections": [],
   "name": "Custom_training_demo.ipynb",
   "provenance": []
  },
  "interpreter": {
   "hash": "859cadc33cf9c93d9354f622d0bfdc75f7d75fc23a9b22a18dfc646072acf83e"
  },
  "kernelspec": {
   "name": "python3",
   "display_name": "Python 3.8.12 64-bit ('tf': conda)"
  },
  "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.8.12"
  },
  "varInspector": {
   "cols": {
    "lenName": 16,
    "lenType": 16,
    "lenVar": 40
   },
   "kernels_config": {
    "python": {
     "delete_cmd_postfix": "",
     "delete_cmd_prefix": "del ",
     "library": "var_list.py",
     "varRefreshCmd": "print(var_dic_list())"
    },
    "r": {
     "delete_cmd_postfix": ") ",
     "delete_cmd_prefix": "rm(",
     "library": "var_list.r",
     "varRefreshCmd": "cat(var_dic_list()) "
    }
   },
   "types_to_exclude": [
    "module",
    "function",
    "builtin_function_or_method",
    "instance",
    "_Feature"
   ],
   "window_display": false
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}