H-CapsNet_Training.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "0c80bff9",
   "metadata": {},
   "source": [
    "***\n",
    "<font size=\"6\"><center><b> H-CapsNet: A Capsule Network for Hierarchical Image Classification </b></center></font>\n",
    "***"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5cc68ea2",
   "metadata": {},
   "source": [
    "# Model Description"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "52e4b202",
   "metadata": {},
   "source": [
    "**H-CapsNet model:**\n",
    "- Use Capsule network for hierarchcial classification\n",
    "- This model contains a deducated feature extraction layer and Capsule netowrk per hierarchy\n",
    "- For training use MixupData data augmentation technique\n",
    "- Model Uses Dynamic LossWeight Distribution system\n",
    "- This model is designed and evaluted using **TensorFlow 2.8.0**"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ac1aab3b",
   "metadata": {},
   "source": [
    "# Import necessary Files and Libraries"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "85c59a96",
   "metadata": {},
   "outputs": [],
   "source": [
    "import tensorflow as tf\n",
    "from tensorflow import keras\n",
    "from tensorflow.keras.preprocessing.image import ImageDataGenerator\n",
    "from tensorflow.keras.models import Sequential\n",
    "from tensorflow.keras.layers import Dense, Activation, Flatten, Dropout, BatchNormalization\n",
    "from tensorflow.keras.layers import Conv2D, MaxPooling2D\n",
    "from tensorflow.keras import regularizers, optimizers\n",
    "from tensorflow.keras import backend as K\n",
    "# Supporting Libraries:\n",
    "    #Mathplot lib for ploting graphs\n",
    "%matplotlib inline\n",
    "import matplotlib\n",
    "import matplotlib.pyplot as plt\n",
    "    # numpy and pandas\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "    #system\n",
    "import os\n",
    "import sys\n",
    "import csv\n",
    "    #import other libraries\n",
    "import math\n",
    "import random\n",
    "from datetime import datetime\n",
    "    # ML model, Dataset and evalution metrics\n",
    "from src import datasets\n",
    "from src import MLmodel\n",
    "from src import metrics\n",
    "from src import MixUp\n",
    "from src import sysenv\n",
    "    # For developind (reloades any python scripts)\n",
    "import importlib\n",
    "\n",
    "%load_ext autoreload\n",
    "%autoreload 2"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e2a610d1",
   "metadata": {},
   "source": [
    "# System information & GPU growth"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e83a195f",
   "metadata": {},
   "outputs": [],
   "source": [
    "systeminfo = sysenv.systeminfo()\n",
    "print(systeminfo)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4df9a5a5",
   "metadata": {},
   "outputs": [],
   "source": [
    "## For Using Multiple GPUs\n",
    "gpus = \"0\" ## Selecting Available gpus\n",
    "gpugrowth = sysenv.gpugrowth(gpus = gpus) ## Limiting GPUS from OS environment\n",
    "gpugrowth.memory_growth() #GPU memory growth"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "44a0baf4",
   "metadata": {},
   "source": [
    "# Training parameters"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "afd9fce3",
   "metadata": {},
   "source": [
    "## Parameters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d8b9b470",
   "metadata": {},
   "outputs": [],
   "source": [
    "train_params = {\"n_epochs\" : 100,\n",
    "                \"batch_size\": 32,\n",
    "                \"lr\": 0.001, # Initial learning rate\n",
    "                \"lr_decay\": 0.95, # Learning rate decay\n",
    "                \"decay_exe\": 9, #learning rate decay execution epoch after\n",
    "               }\n",
    "model_params = {\"P_Cap_Dim\" : 8, # Primary Capsule Dimentions\n",
    "                \"S_Cap_Dim\" : 16, # Secondary Capsule Dimention\n",
    "                \"Reconstruction_LW\" : 0.0005, # Decoder loss weight\n",
    "                \"class_loss\" : MLmodel.MarginLoss(), ## Class prediction loss\n",
    "                \"reconstruction_loss\" : 'mse'\n",
    "               }"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f8f4b411",
   "metadata": {},
   "source": [
    "## Learning rate decay"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4d2a3dcf",
   "metadata": {},
   "outputs": [],
   "source": [
    "def scheduler(epoch):\n",
    "    learning_rate_init = train_params[\"lr\"]\n",
    "    \n",
    "    if epoch > train_params[\"decay_exe\"]:\n",
    "        learning_rate_init = train_params[\"lr\"] * (train_params[\"lr_decay\"] ** (epoch-9))\n",
    "        \n",
    "    return learning_rate_init"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "35bdb8ef",
   "metadata": {},
   "source": [
    "# MNIST dataset"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "012a135d",
   "metadata": {},
   "source": [
    "## Import MNIST dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b40eba3c",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "## Import dataset\n",
    "dataset = datasets.MNIST()\n",
    "# dataset = datasets.MNIST(version = 'reduce')\n",
    "print(dataset.keys())\n",
    "dataset['tree'].show()\n",
    "\n",
    "\n",
    "input_shape = dataset['x_train'].shape[1:]\n",
    "print('INPUT SHAPE:',input_shape,'\\n')\n",
    "\n",
    "print(\"TRAIN: \\r\\n\")\n",
    "print(dataset['x_train'].shape)\n",
    "print(dataset['y_train_fine'].shape)\n",
    "print(dataset['y_train_coarse'].shape)\n",
    "\n",
    "print(\"\\nTEST: \\r\\n\")\n",
    "print(dataset['x_test'].shape)\n",
    "print(dataset['y_test_fine'].shape)\n",
    "print(dataset['y_test_coarse'].shape)\n",
    "\n",
    "fine_class = len(np.unique(np.argmax(dataset['y_test_fine'], axis=1)))\n",
    "coarse_class = len(np.unique(np.argmax(dataset['y_test_coarse'], axis=1)))\n",
    "print('\\nNumber of Classes in Label Tree:',\n",
    "      '\\nCoarse Level = ',coarse_class,\n",
    "      '\\nFine Level = ',fine_class)\n",
    "\n",
    "# Plot Random samples \n",
    "datasets.plot_sample_image(dataset['x_train'],\n",
    "                           {'coarse':dataset['y_train_coarse'],\n",
    "                            'fine':dataset['y_train_fine']})"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9852cb89",
   "metadata": {},
   "source": [
    "## Create ML Model"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "12bd67a2",
   "metadata": {},
   "source": [
    "### Model parameters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1dbf5025",
   "metadata": {},
   "outputs": [],
   "source": [
    "initial_lw = MLmodel.initial_lw({\"coarse\": coarse_class, \"fine\": fine_class})\n",
    "\n",
    "lossweight = {'coarse_lw' : K.variable(value = initial_lw['coarse'], dtype=\"float32\", name=\"coarse_lw\"),\n",
    "             'fine_lw' : K.variable(value = initial_lw['fine'], dtype=\"float32\", name=\"fine_lw\"),\n",
    "              'decoder_lw' : model_params['Reconstruction_LW']\n",
    "             }\n",
    "\n",
    "def get_compiled_model():\n",
    "    model = MLmodel.HCapsNet_2_Level_MNIST(input_shape,\n",
    "                                   coarse_class, fine_class)\n",
    "    model.compile(optimizer='adam',\n",
    "                  loss=[model_params[\"class_loss\"],\n",
    "                        model_params[\"class_loss\"],\n",
    "                        model_params[\"reconstruction_loss\"]],\n",
    "                  loss_weights=[lossweight['coarse_lw'],\n",
    "                                lossweight['fine_lw'],\n",
    "                                lossweight['decoder_lw']],\n",
    "                  metrics={'Fine_prediction_output_layer': 'accuracy',\n",
    "                           'Coarse_prediction_output_layer': 'accuracy'}\n",
    "                 )\n",
    "    return model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0c663bae",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "tf.keras.backend.clear_session() ## clear session\n",
    "\n",
    "# model = get_compiled_model()\n",
    "\n",
    "#Multiple GPU ------ Windows\n",
    "strategy = tf.distribute.MirroredStrategy(cross_device_ops=tf.distribute.HierarchicalCopyAllReduce())\n",
    "\n",
    "print(\"Number of devices: {}\".format(strategy.num_replicas_in_sync))\n",
    "\n",
    "# Open a strategy scope.\n",
    "with strategy.scope():\n",
    "    # Everything that creates variables should be under the strategy scope.\n",
    "    # In general this is only model construction & `compile()`.\n",
    "    model = get_compiled_model()\n",
    "\n",
    "### Log directory\n",
    "directory = sysenv.log_dir(dataset[\"name\"]+'/'+model.name)\n",
    "\n",
    "model.summary()\n",
    "keras.utils.plot_model(model, to_file = directory+\"/H-CapsNet.png\", show_shapes=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "18812108",
   "metadata": {},
   "source": [
    "### Callbacks"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e967ffe2",
   "metadata": {},
   "outputs": [],
   "source": [
    "tb = keras.callbacks.TensorBoard(directory+'./tb_logs')\n",
    "log = keras.callbacks.CSVLogger(directory+'/log.csv', append=True)\n",
    "\n",
    "checkpoint = keras.callbacks.ModelCheckpoint(directory+'/epoch_best.h5',\n",
    "                                             monitor='val_Fine_prediction_output_layer_accuracy',\n",
    "                                             save_best_only=True, save_weights_only=True, verbose=1\n",
    "                                            )\n",
    "\n",
    "change_lw = MLmodel.LossWeightsModifier(lossweight,\n",
    "                                        initial_lw,\n",
    "                                        directory = directory)\n",
    "\n",
    "lr_decay = keras.callbacks.LearningRateScheduler(scheduler)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "47c9be8a",
   "metadata": {},
   "source": [
    "## Training Model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a16b46f3",
   "metadata": {},
   "outputs": [],
   "source": [
    "datagen = ImageDataGenerator(width_shift_range=0.1,\n",
    "                            height_shift_range=0.1)\n",
    "\n",
    "training_generator = MixUp.MixupGenerator_2level(dataset['x_train'],\n",
    "                                                 dataset['y_train_coarse'], dataset['y_train_fine'],\n",
    "                                                 batch_size=train_params[\"batch_size\"], alpha=0.02, \n",
    "                                                 datagen=datagen\n",
    "                                                )()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "95663995",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "history = model.fit_generator(generator=training_generator,\n",
    "                                  steps_per_epoch = int(dataset['x_train'].shape[0] / train_params[\"batch_size\"]),\n",
    "                                  epochs = train_params[\"n_epochs\"],\n",
    "                                  validation_data = ([dataset['x_test'], dataset['y_test_coarse'], dataset['y_test_fine']],\n",
    "                                                     [dataset['y_test_coarse'], dataset['y_test_fine'], dataset['x_test']]),\n",
    "                                  callbacks = [tb, log, checkpoint, lr_decay, change_lw],\n",
    "                                  verbose=1)\n",
    "model_save_dir = str(directory+'/trained_model.h5')\n",
    "try:\n",
    "    model.save_weights(model_save_dir)\n",
    "    print('Trained model saved to = ', model_save_dir)\n",
    "except:\n",
    "    print('Model Wight is not saved')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f0a26c4c",
   "metadata": {},
   "source": [
    "## Model Analysis"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e255497e",
   "metadata": {},
   "outputs": [],
   "source": [
    "model_analysis = MLmodel.model_analysis(model, dataset)\n",
    "results = model_analysis.evaluate()\n",
    "predictions = model_analysis.prediction()\n",
    "\n",
    "true_label = [dataset['y_test_coarse'],dataset['y_test_fine']]\n",
    "pred_label = [predictions[0],predictions[1]]\n",
    "metrics.lvl_wise_metric(true_label,pred_label)\n",
    "\n",
    "h_measurements,consistency,exact_match = metrics.hmeasurements(true_label,\n",
    "                                       pred_label,\n",
    "                                       dataset['tree'])\n",
    "print('\\nHierarchical Precision =',h_measurements[0],\n",
    "      '\\nHierarchical Recall =', h_measurements[1],\n",
    "      '\\nHierarchical F1-Score =',h_measurements[2],\n",
    "      '\\n Consistency = ', consistency,\n",
    "      '\\n Exact Match = ', exact_match,\n",
    "     )"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ad1fd42f",
   "metadata": {},
   "source": [
    "# EMNIST Dataset"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0c7a8285",
   "metadata": {},
   "source": [
    "## Import EMNIST dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ba9b3b27",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "## Import dataset\n",
    "dataset = datasets.E_MNIST()\n",
    "# dataset = datasets.E_MNIST(version='reduce')\n",
    "print(dataset.keys())\n",
    "dataset['tree'].show()\n",
    "\n",
    "input_shape = dataset['x_train'].shape[1:]\n",
    "print('INPUT SHAPE:',input_shape,'\\n')\n",
    "\n",
    "print(\"TRAIN: \\r\\n\")\n",
    "print(dataset['x_train'].shape)\n",
    "print(dataset['y_train_fine'].shape)\n",
    "print(dataset['y_train_coarse'].shape)\n",
    "\n",
    "print(\"\\nTEST: \\r\\n\")\n",
    "print(dataset['x_test'].shape)\n",
    "print(dataset['y_test_fine'].shape)\n",
    "print(dataset['y_train_coarse'].shape)\n",
    "\n",
    "fine_class = len(np.unique(np.argmax(dataset['y_train_fine'], axis=1)))\n",
    "coarse_class = len(np.unique(np.argmax(dataset['y_train_coarse'], axis=1)))\n",
    "print('\\nNumber of Classes in Label Tree:',\n",
    "      '\\nCoarse Level = ',coarse_class,\n",
    "      '\\nFine Level = ',fine_class)\n",
    "    # Plot Random samples \n",
    "datasets.plot_sample_image(dataset['x_train'],\n",
    "                           {'coarse':dataset['y_train_coarse'],\n",
    "                            'fine':dataset['y_train_fine']})"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "95a1209a",
   "metadata": {},
   "source": [
    "## Create ML Model"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ccf19493",
   "metadata": {},
   "source": [
    "### Model parameters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "57c10aa8",
   "metadata": {},
   "outputs": [],
   "source": [
    "initial_lw = MLmodel.initial_lw({\"coarse\": coarse_class, \"fine\": fine_class})\n",
    "\n",
    "lossweight = {'coarse_lw' : K.variable(value = initial_lw['coarse'], dtype=\"float32\", name=\"coarse_lw\"),\n",
    "             'fine_lw' : K.variable(value = initial_lw['fine'], dtype=\"float32\", name=\"fine_lw\"),\n",
    "              'decoder_lw' : model_params['Reconstruction_LW']\n",
    "             }\n",
    "\n",
    "def get_compiled_model():\n",
    "    model = MLmodel.HCapsNet_2_Level_EMNIST(input_shape,\n",
    "                                   coarse_class, fine_class)\n",
    "    model.compile(optimizer='adam',\n",
    "                  loss=[model_params[\"class_loss\"],\n",
    "                        model_params[\"class_loss\"],\n",
    "                        model_params[\"reconstruction_loss\"]],\n",
    "                  loss_weights=[lossweight['coarse_lw'],\n",
    "                                lossweight['fine_lw'],\n",
    "                                lossweight['decoder_lw']],\n",
    "                  metrics={'Fine_prediction_output_layer': 'accuracy',\n",
    "                           'Coarse_prediction_output_layer': 'accuracy'}\n",
    "                 )\n",
    "    return model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c37edd8e",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "tf.keras.backend.clear_session() ## clear session\n",
    "\n",
    "# model = get_compiled_model()\n",
    "\n",
    "#Multiple GPU ------ Windows\n",
    "strategy = tf.distribute.MirroredStrategy(cross_device_ops=tf.distribute.HierarchicalCopyAllReduce())\n",
    "\n",
    "print(\"Number of devices: {}\".format(strategy.num_replicas_in_sync))\n",
    "\n",
    "# Open a strategy scope.\n",
    "with strategy.scope():\n",
    "    # Everything that creates variables should be under the strategy scope.\n",
    "    # In general this is only model construction & `compile()`.\n",
    "    model = get_compiled_model()\n",
    "\n",
    "### Log directory\n",
    "directory = sysenv.log_dir(dataset[\"name\"]+'/'+model.name)\n",
    "\n",
    "model.summary()\n",
    "keras.utils.plot_model(model, to_file = directory+\"/H-CapsNet.png\", show_shapes=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "701c299c",
   "metadata": {},
   "source": [
    "### Callbacks"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5f5ed60b",
   "metadata": {},
   "outputs": [],
   "source": [
    "tb = keras.callbacks.TensorBoard(directory+'./tb_logs')\n",
    "log = keras.callbacks.CSVLogger(directory+'/log.csv', append=True)\n",
    "\n",
    "checkpoint = keras.callbacks.ModelCheckpoint(directory+'/epoch_best.h5',\n",
    "                                             monitor='val_Fine_prediction_output_layer_accuracy',\n",
    "                                             save_best_only=True, save_weights_only=True, verbose=1\n",
    "                                            )\n",
    "\n",
    "change_lw = MLmodel.LossWeightsModifier(lossweight,\n",
    "                                        initial_lw,\n",
    "                                        directory = directory)\n",
    "lr_decay = keras.callbacks.LearningRateScheduler(scheduler)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4a2a5d1b",
   "metadata": {},
   "source": [
    "## Training Model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1db1fc82",
   "metadata": {},
   "outputs": [],
   "source": [
    "datagen = ImageDataGenerator(width_shift_range=0.1,\n",
    "                            height_shift_range=0.1)\n",
    "\n",
    "training_generator = MixUp.MixupGenerator_2level(dataset['x_train'],\n",
    "                                                 dataset['y_train_coarse'], dataset['y_train_fine'],\n",
    "                                                 batch_size=train_params[\"batch_size\"], alpha=0.02, \n",
    "                                                 datagen=datagen\n",
    "                                                )()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9221a24a",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "history = model.fit_generator(generator=training_generator,\n",
    "                                  steps_per_epoch = int(dataset['x_train'].shape[0] / train_params[\"batch_size\"]),\n",
    "                                  epochs = train_params[\"n_epochs\"],\n",
    "                                  validation_data = ([dataset['x_test'], dataset['y_test_coarse'], dataset['y_test_fine']],\n",
    "                                                     [dataset['y_test_coarse'], dataset['y_test_fine'], dataset['x_test']]),\n",
    "                                  callbacks = [tb, log, checkpoint, lr_decay, change_lw],\n",
    "                                  verbose=1)\n",
    "model_save_dir = str(directory+'/trained_model.h5')\n",
    "try:\n",
    "    model.save_weights(model_save_dir)\n",
    "    print('Trained model saved to = ', model_save_dir)\n",
    "except:\n",
    "    print('Model Wight is not saved')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6ed0ebb9",
   "metadata": {},
   "source": [
    "## Model Analysis"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7b0d29c1",
   "metadata": {},
   "outputs": [],
   "source": [
    "model_analysis = MLmodel.model_analysis(model, dataset)\n",
    "results = model_analysis.evaluate()\n",
    "predictions = model_analysis.prediction()\n",
    "\n",
    "true_label = [dataset['y_test_coarse'],dataset['y_test_fine']]\n",
    "pred_label = [predictions[0],predictions[1]]\n",
    "metrics.lvl_wise_metric(true_label,pred_label)\n",
    "\n",
    "h_measurements,consistency,exact_match = metrics.hmeasurements(true_label,\n",
    "                                       pred_label,\n",
    "                                       dataset['tree'])\n",
    "print('\\nHierarchical Precision =',h_measurements[0],\n",
    "      '\\nHierarchical Recall =', h_measurements[1],\n",
    "      '\\nHierarchical F1-Score =',h_measurements[2],\n",
    "      '\\n Consistency = ', consistency,\n",
    "      '\\n Exact Match = ', exact_match,\n",
    "     )"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ffa5377f",
   "metadata": {},
   "source": [
    "# Fashion-MNIST dataset"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "49a70a8d",
   "metadata": {},
   "source": [
    "## Import Fashion-MNIST dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "049aa68b",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "## Import dataset\n",
    "dataset = datasets.F_MNIST()\n",
    "# dataset = datasets.F_MNIST('reduce')\n",
    "print(dataset.keys())\n",
    "dataset['tree'].show()\n",
    "\n",
    "input_shape = dataset['x_train'].shape[1:]\n",
    "print('INPUT SHAPE:',input_shape,'\\n')\n",
    "\n",
    "print(\"TRAIN: \\r\\n\")\n",
    "print(dataset['x_train'].shape)\n",
    "print(dataset['y_train_fine'].shape)\n",
    "print(dataset['y_train_medium'].shape)\n",
    "print(dataset['y_train_coarse'].shape)\n",
    "\n",
    "print(\"\\nTEST: \\r\\n\")\n",
    "print(dataset['x_test'].shape)\n",
    "print(dataset['y_test_fine'].shape)\n",
    "print(dataset['y_test_medium'].shape)\n",
    "print(dataset['y_test_coarse'].shape)\n",
    "\n",
    "fine_class = len(np.unique(np.argmax(dataset['y_train_fine'], axis=1)))\n",
    "medium_class = len(np.unique(np.argmax(dataset['y_train_medium'], axis=1)))\n",
    "coarse_class = len(np.unique(np.argmax(dataset['y_train_coarse'], axis=1)))\n",
    "print('\\nNumber of Classes in Label Tree:',\n",
    "      '\\nCoarse Level = ',coarse_class,\n",
    "      '\\nMedium Level = ',medium_class,\n",
    "      '\\nFine Level = ',fine_class)\n",
    "\n",
    "datasets.plot_sample_image(dataset['x_train'],\n",
    "                           {'coarse':dataset['y_train_coarse'],\n",
    "                            'medium':dataset['y_train_medium'],\n",
    "                            'fine':dataset['y_train_fine']})"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0d77da10",
   "metadata": {},
   "source": [
    "### Model parameters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f58a5729",
   "metadata": {},
   "outputs": [],
   "source": [
    "initial_lw = MLmodel.initial_lw({\"coarse\": coarse_class,\n",
    "                                 \"medium\": medium_class,\n",
    "                                 \"fine\": fine_class})\n",
    "\n",
    "lossweight = {'coarse_lw' : K.variable(value = initial_lw['coarse'], dtype=\"float32\", name=\"coarse_lw\"),\n",
    "             'medium_lw' : K.variable(value = initial_lw['medium'], dtype=\"float32\", name=\"medium_lw\"),\n",
    "             'fine_lw' : K.variable(value = initial_lw['fine'], dtype=\"float32\", name=\"fine_lw\"),\n",
    "              'decoder_lw' : model_params['Reconstruction_LW']\n",
    "             }\n",
    "\n",
    "def get_compiled_model():\n",
    "    model = MLmodel.HCapsNet_3_Level_FMNIST(input_shape,\n",
    "                                     coarse_class, medium_class, fine_class)\n",
    "    model.compile(optimizer='adam',\n",
    "                  loss=[model_params[\"class_loss\"],\n",
    "                        model_params[\"class_loss\"],\n",
    "                        model_params[\"class_loss\"],\n",
    "                        model_params[\"reconstruction_loss\"]],\n",
    "                  loss_weights=[lossweight['coarse_lw'],\n",
    "                                lossweight['medium_lw'],\n",
    "                                lossweight['fine_lw'],\n",
    "                                lossweight['decoder_lw']],\n",
    "                  metrics={'Fine_prediction_output_layer': 'accuracy',\n",
    "                           'Medium_prediction_output_layer': 'accuracy',\n",
    "                           'Coarse_prediction_output_layer': 'accuracy'}\n",
    "                 )\n",
    "    return model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d1c9b2ce",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "tf.keras.backend.clear_session() ## clear session\n",
    "\n",
    "# model = get_compiled_model()\n",
    "\n",
    "#Multiple GPU ------ Windows\n",
    "strategy = tf.distribute.MirroredStrategy(cross_device_ops=tf.distribute.HierarchicalCopyAllReduce())\n",
    "\n",
    "print(\"Number of devices: {}\".format(strategy.num_replicas_in_sync))\n",
    "\n",
    "# Open a strategy scope.\n",
    "with strategy.scope():\n",
    "    # Everything that creates variables should be under the strategy scope.\n",
    "    # In general this is only model construction & `compile()`.\n",
    "    model = get_compiled_model()\n",
    "\n",
    "### Log directory\n",
    "directory = sysenv.log_dir(dataset[\"name\"]+'/'+model.name)\n",
    "\n",
    "model.summary()\n",
    "keras.utils.plot_model(model, to_file = directory+\"/H-CapsNet.png\", show_shapes=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3b6dd366",
   "metadata": {},
   "source": [
    "### Callbacks"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a5f9628c",
   "metadata": {},
   "outputs": [],
   "source": [
    "tb = keras.callbacks.TensorBoard(directory+'./tb_logs')\n",
    "log = keras.callbacks.CSVLogger(directory+'/log.csv', append=True)\n",
    "\n",
    "checkpoint = keras.callbacks.ModelCheckpoint(directory+'/epoch_best.h5',\n",
    "                                             monitor='val_Fine_prediction_output_layer_accuracy',\n",
    "                                             save_best_only=True, save_weights_only=True, verbose=1\n",
    "                                            )\n",
    "\n",
    "change_lw = MLmodel.LossWeightsModifier(lossweight,\n",
    "                                        initial_lw,\n",
    "                                        directory = directory)\n",
    "lr_decay = keras.callbacks.LearningRateScheduler(scheduler)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "88203484",
   "metadata": {},
   "source": [
    "## Training Model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0d52911d",
   "metadata": {},
   "outputs": [],
   "source": [
    "datagen = ImageDataGenerator(width_shift_range=0.1,\n",
    "                            height_shift_range=0.1)\n",
    "\n",
    "training_generator = MixUp.MixupGenerator_3level(dataset['x_train'],\n",
    "                                                 dataset['y_train_coarse'], \n",
    "                                                 dataset['y_train_medium'],\n",
    "                                                 dataset['y_train_fine'],\n",
    "                                                 batch_size=train_params[\"batch_size\"],\n",
    "                                                 alpha=0.2, \n",
    "                                                 datagen=datagen\n",
    "                                                )()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ebcd4475",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "history = model.fit_generator(generator=training_generator,\n",
    "                              steps_per_epoch = int(dataset['x_train'].shape[0] / train_params[\"batch_size\"]),\n",
    "                              epochs = train_params[\"n_epochs\"],\n",
    "                              validation_data = ([dataset['x_test'],\n",
    "                                                  dataset['y_test_coarse'],\n",
    "                                                  dataset['y_test_medium'],\n",
    "                                                  dataset['y_test_fine']],\n",
    "                                                 [dataset['y_test_coarse'],\n",
    "                                                  dataset['y_test_medium'],\n",
    "                                                  dataset['y_test_fine'],\n",
    "                                                  dataset['x_test']]),\n",
    "                                  callbacks = [tb,\n",
    "                                               log,\n",
    "                                               checkpoint,\n",
    "                                               lr_decay,\n",
    "                                               change_lw],\n",
    "                                  verbose=1)\n",
    "model_save_dir = str(directory+'/trained_model.h5')\n",
    "try:\n",
    "    model.save_weights(model_save_dir)\n",
    "    print('Trained model saved to = ', model_save_dir)\n",
    "except:\n",
    "    print('Model Wight is not saved')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f23fef32",
   "metadata": {},
   "source": [
    "## Model Analysis"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "51db3a4a",
   "metadata": {},
   "outputs": [],
   "source": [
    "model_analysis = MLmodel.model_analysis(model, dataset)\n",
    "\n",
    "results = model_analysis.evaluate()\n",
    "\n",
    "predictions = model_analysis.prediction()\n",
    "\n",
    "true_label = [dataset['y_test_coarse'],dataset['y_test_medium'],dataset['y_test_fine']]\n",
    "pred_label = [predictions[0],predictions[1],predictions[2]]\n",
    "metrics.lvl_wise_metric(true_label,pred_label)\n",
    "\n",
    "h_measurements,consistency,exact_match = metrics.hmeasurements(true_label,\n",
    "                                       pred_label,\n",
    "                                       dataset['tree'])\n",
    "print('\\nHierarchical Precision =',h_measurements[0],\n",
    "      '\\nHierarchical Recall =', h_measurements[1],\n",
    "      '\\nHierarchical F1-Score =',h_measurements[2],\n",
    "      '\\n Consistency = ', consistency,\n",
    "      '\\n Exact Match = ', exact_match,\n",
    "     )"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "94110b27",
   "metadata": {},
   "source": [
    "# CIFAR-10 Dataset"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6c386a4b",
   "metadata": {},
   "source": [
    "## Import CIFAR-10 dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bd673379",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "## Import dataset\n",
    "dataset = datasets.CIFAR10()\n",
    "# dataset = datasets.CIFAR10('reduce')\n",
    "print(dataset.keys())\n",
    "dataset['tree'].show()\n",
    "\n",
    "input_shape = dataset['x_train'].shape[1:]\n",
    "print('INPUT SHAPE:',input_shape,'\\n')\n",
    "\n",
    "print(\"TRAIN: \\r\\n\")\n",
    "print(dataset['x_train'].shape)\n",
    "print(dataset['y_train_fine'].shape)\n",
    "print(dataset['y_train_medium'].shape)\n",
    "print(dataset['y_train_coarse'].shape)\n",
    "\n",
    "print(\"\\nTEST: \\r\\n\")\n",
    "print(dataset['x_test'].shape)\n",
    "print(dataset['y_test_fine'].shape)\n",
    "print(dataset['y_test_medium'].shape)\n",
    "print(dataset['y_test_coarse'].shape)\n",
    "\n",
    "fine_class = len(np.unique(np.argmax(dataset['y_train_fine'], axis=1)))\n",
    "medium_class = len(np.unique(np.argmax(dataset['y_train_medium'], axis=1)))\n",
    "coarse_class = len(np.unique(np.argmax(dataset['y_train_coarse'], axis=1)))\n",
    "print('\\nNumber of Classes in Label Tree:',\n",
    "      '\\nCoarse Level = ',coarse_class,\n",
    "      '\\nMedium Level = ',medium_class,\n",
    "      '\\nFine Level = ',fine_class)\n",
    "datasets.plot_sample_image(dataset['x_train'],\n",
    "                           {'coarse':dataset['y_train_coarse'],\n",
    "                            'medium':dataset['y_train_medium'],\n",
    "                            'fine':dataset['y_train_fine']})"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1a4535dc",
   "metadata": {},
   "source": [
    "### Model parameters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "53f6ff95",
   "metadata": {},
   "outputs": [],
   "source": [
    "initial_lw = MLmodel.initial_lw({\"coarse\": coarse_class,\n",
    "                                 \"medium\": medium_class,\n",
    "                                 \"fine\": fine_class})\n",
    "\n",
    "lossweight = {'coarse_lw' : K.variable(value = initial_lw['coarse'], dtype=\"float32\", name=\"coarse_lw\"),\n",
    "             'medium_lw' : K.variable(value = initial_lw['medium'], dtype=\"float32\", name=\"medium_lw\"),\n",
    "             'fine_lw' : K.variable(value = initial_lw['fine'], dtype=\"float32\", name=\"fine_lw\"),\n",
    "              'decoder_lw' : model_params['Reconstruction_LW']\n",
    "             }\n",
    "\n",
    "def get_compiled_model():\n",
    "    model = MLmodel.HCapsNet_3_Level(input_shape,\n",
    "                                     coarse_class, medium_class, fine_class)\n",
    "    model.compile(optimizer='adam',\n",
    "                  loss=[model_params[\"class_loss\"],\n",
    "                        model_params[\"class_loss\"],\n",
    "                        model_params[\"class_loss\"],\n",
    "                        model_params[\"reconstruction_loss\"]],\n",
    "                  loss_weights=[lossweight['coarse_lw'],\n",
    "                                lossweight['medium_lw'],\n",
    "                                lossweight['fine_lw'],\n",
    "                                lossweight['decoder_lw']],\n",
    "                  metrics={'Fine_prediction_output_layer': 'accuracy',\n",
    "                           'Medium_prediction_output_layer': 'accuracy',\n",
    "                           'Coarse_prediction_output_layer': 'accuracy'}\n",
    "                 )\n",
    "    return model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "980911c6",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "tf.keras.backend.clear_session() ## clear session\n",
    "\n",
    "# model = get_compiled_model()\n",
    "\n",
    "#Multiple GPU ------ Windows\n",
    "strategy = tf.distribute.MirroredStrategy(cross_device_ops=tf.distribute.HierarchicalCopyAllReduce())\n",
    "\n",
    "print(\"Number of devices: {}\".format(strategy.num_replicas_in_sync))\n",
    "\n",
    "# Open a strategy scope.\n",
    "with strategy.scope():\n",
    "    # Everything that creates variables should be under the strategy scope.\n",
    "    # In general this is only model construction & `compile()`.\n",
    "    model = get_compiled_model()\n",
    "\n",
    "### Log directory\n",
    "directory = sysenv.log_dir(dataset[\"name\"]+'/'+model.name)\n",
    "\n",
    "model.summary()\n",
    "keras.utils.plot_model(model, to_file = directory+\"/H-CapsNet.png\", show_shapes=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8985d20d",
   "metadata": {},
   "source": [
    "### Callbacks"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8e52eeaf",
   "metadata": {},
   "outputs": [],
   "source": [
    "tb = keras.callbacks.TensorBoard(directory+'./tb_logs')\n",
    "log = keras.callbacks.CSVLogger(directory+'/log.csv', append=True)\n",
    "\n",
    "checkpoint = keras.callbacks.ModelCheckpoint(directory+'/epoch_best.h5',\n",
    "                                             monitor='val_Fine_prediction_output_layer_accuracy',\n",
    "                                             save_best_only=True, save_weights_only=True, verbose=1\n",
    "                                            )\n",
    "\n",
    "change_lw = MLmodel.LossWeightsModifier(lossweight,\n",
    "                                        initial_lw,\n",
    "                                        directory = directory)\n",
    "\n",
    "lr_decay = keras.callbacks.LearningRateScheduler(scheduler)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "faaadb54",
   "metadata": {},
   "source": [
    "## Training Model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "95fa75af",
   "metadata": {},
   "outputs": [],
   "source": [
    "datagen = ImageDataGenerator(width_shift_range=0.1,\n",
    "                            height_shift_range=0.1)\n",
    "\n",
    "training_generator = MixUp.MixupGenerator_3level(dataset['x_train'],\n",
    "                                                 dataset['y_train_coarse'], \n",
    "                                                 dataset['y_train_medium'],\n",
    "                                                 dataset['y_train_fine'],\n",
    "                                                 batch_size=train_params[\"batch_size\"],\n",
    "                                                 alpha=0.2, \n",
    "                                                 datagen=datagen\n",
    "                                                )()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ab98dea9",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "history = model.fit_generator(generator=training_generator,\n",
    "                              steps_per_epoch = int(dataset['x_train'].shape[0] / train_params[\"batch_size\"]),\n",
    "                              epochs = train_params[\"n_epochs\"],\n",
    "                              validation_data = ([dataset['x_test'],\n",
    "                                                  dataset['y_test_coarse'],\n",
    "                                                  dataset['y_test_medium'],\n",
    "                                                  dataset['y_test_fine']],\n",
    "                                                 [dataset['y_test_coarse'],\n",
    "                                                  dataset['y_test_medium'],\n",
    "                                                  dataset['y_test_fine'],\n",
    "                                                  dataset['x_test']]),\n",
    "                                  callbacks = [tb,\n",
    "                                               log,\n",
    "                                               checkpoint,\n",
    "                                               lr_decay,\n",
    "                                               change_lw],\n",
    "                                  verbose=1)\n",
    "model_save_dir = str(directory+'/trained_model.h5')\n",
    "try:\n",
    "    model.save_weights(model_save_dir)\n",
    "    print('Trained model saved to = ', model_save_dir)\n",
    "except:\n",
    "    print('Model Wight is not saved')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "eda68f5c",
   "metadata": {},
   "source": [
    "## Model Analysis"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "50226fb7",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "model_analysis = MLmodel.model_analysis(model, dataset)\n",
    "\n",
    "results = model_analysis.evaluate()\n",
    "\n",
    "predictions = model_analysis.prediction()\n",
    "\n",
    "true_label = [dataset['y_test_coarse'],dataset['y_test_medium'],dataset['y_test_fine']]\n",
    "pred_label = [predictions[0],predictions[1],predictions[2]]\n",
    "metrics.lvl_wise_metric(true_label,pred_label)\n",
    "\n",
    "h_measurements,consistency,exact_match = metrics.hmeasurements(true_label,\n",
    "                                       pred_label,\n",
    "                                       dataset['tree'])\n",
    "print('\\nHierarchical Precision =',h_measurements[0],\n",
    "      '\\nHierarchical Recall =', h_measurements[1],\n",
    "      '\\nHierarchical F1-Score =',h_measurements[2],\n",
    "      '\\n Consistency = ', consistency,\n",
    "      '\\n Exact Match = ', exact_match,\n",
    "     )"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "892271fd",
   "metadata": {},
   "source": [
    "# CIFAR-100 Dataset"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5d93e8b4",
   "metadata": {},
   "source": [
    "## Import CIFAR-100 dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d6674cb0",
   "metadata": {},
   "outputs": [],
   "source": [
    "## Import dataset\n",
    "dataset = datasets.CIFAR100()\n",
    "# dataset = datasets.CIFAR100('reduce')\n",
    "print(dataset.keys())\n",
    "dataset['tree'].show()\n",
    "\n",
    "input_shape = dataset['x_train'].shape[1:]\n",
    "print('INPUT SHAPE:',input_shape,'\\n')\n",
    "\n",
    "print(\"TRAIN: \\r\\n\")\n",
    "print(dataset['x_train'].shape)\n",
    "print(dataset['y_train_fine'].shape)\n",
    "print(dataset['y_train_medium'].shape)\n",
    "print(dataset['y_train_coarse'].shape)\n",
    "\n",
    "print(\"\\nTEST: \\r\\n\")\n",
    "print(dataset['x_test'].shape)\n",
    "print(dataset['y_test_fine'].shape)\n",
    "print(dataset['y_test_medium'].shape)\n",
    "print(dataset['y_test_coarse'].shape)\n",
    "\n",
    "fine_class = len(np.unique(np.argmax(dataset['y_train_fine'], axis=1)))\n",
    "medium_class = len(np.unique(np.argmax(dataset['y_train_medium'], axis=1)))\n",
    "coarse_class = len(np.unique(np.argmax(dataset['y_train_coarse'], axis=1)))\n",
    "print('\\nNumber of Classes in Label Tree:',\n",
    "      '\\nCoarse Level = ',coarse_class,\n",
    "      '\\nMedium Level = ',medium_class,\n",
    "      '\\nFine Level = ',fine_class)\n",
    "\n",
    "datasets.plot_sample_image(dataset['x_train'],\n",
    "                           {'coarse':dataset['y_train_coarse'],\n",
    "                            'medium':dataset['y_train_medium'],\n",
    "                            'fine':dataset['y_train_fine']})"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "44313807",
   "metadata": {},
   "source": [
    "### Model parameters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dcb47c62",
   "metadata": {},
   "outputs": [],
   "source": [
    "initial_lw = MLmodel.initial_lw({\"coarse\": coarse_class,\n",
    "                                 \"medium\": medium_class,\n",
    "                                 \"fine\": fine_class})\n",
    "\n",
    "lossweight = {'coarse_lw' : K.variable(value = initial_lw['coarse'], dtype=\"float32\", name=\"coarse_lw\"),\n",
    "             'medium_lw' : K.variable(value = initial_lw['medium'], dtype=\"float32\", name=\"medium_lw\"),\n",
    "             'fine_lw' : K.variable(value = initial_lw['fine'], dtype=\"float32\", name=\"fine_lw\"),\n",
    "              'decoder_lw' : model_params['Reconstruction_LW']\n",
    "             }\n",
    "\n",
    "def get_compiled_model():\n",
    "    model = MLmodel.HCapsNet_3_Level(input_shape,\n",
    "                                     coarse_class, medium_class, fine_class)\n",
    "    model.compile(optimizer='adam',\n",
    "                  loss=[model_params[\"class_loss\"],\n",
    "                        model_params[\"class_loss\"],\n",
    "                        model_params[\"class_loss\"],\n",
    "                        model_params[\"reconstruction_loss\"]],\n",
    "                  loss_weights=[lossweight['coarse_lw'],\n",
    "                                lossweight['medium_lw'],\n",
    "                                lossweight['fine_lw'],\n",
    "                                lossweight['decoder_lw']],\n",
    "                  metrics={'Fine_prediction_output_layer': 'accuracy',\n",
    "                           'Medium_prediction_output_layer': 'accuracy',\n",
    "                           'Coarse_prediction_output_layer': 'accuracy'}\n",
    "                 )\n",
    "    return model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6fca43d3",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "tf.keras.backend.clear_session() ## clear session\n",
    "\n",
    "# model = get_compiled_model()\n",
    "\n",
    "#Multiple GPU ------ Windows\n",
    "strategy = tf.distribute.MirroredStrategy(cross_device_ops=tf.distribute.HierarchicalCopyAllReduce())\n",
    "\n",
    "print(\"Number of devices: {}\".format(strategy.num_replicas_in_sync))\n",
    "\n",
    "# Open a strategy scope.\n",
    "with strategy.scope():\n",
    "    # Everything that creates variables should be under the strategy scope.\n",
    "    # In general this is only model construction & `compile()`.\n",
    "    model = get_compiled_model()\n",
    "\n",
    "### Log directory\n",
    "directory = sysenv.log_dir(dataset[\"name\"]+'/'+model.name)\n",
    "\n",
    "model.summary()\n",
    "keras.utils.plot_model(model, to_file = directory+\"/H-CapsNet.png\", show_shapes=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c10290ce",
   "metadata": {},
   "source": [
    "### Callbacks"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "916ac92f",
   "metadata": {},
   "outputs": [],
   "source": [
    "tb = keras.callbacks.TensorBoard(directory+'./tb_logs')\n",
    "\n",
    "log = keras.callbacks.CSVLogger(directory+'/log.csv', append=True)\n",
    "\n",
    "checkpoint = keras.callbacks.ModelCheckpoint(directory+'/epoch_best.h5',\n",
    "                                             monitor='val_Fine_prediction_output_layer_accuracy',\n",
    "                                             save_best_only=True, save_weights_only=True, verbose=1\n",
    "                                            )\n",
    "\n",
    "change_lw = MLmodel.LossWeightsModifier(lossweight,\n",
    "                                        initial_lw,\n",
    "                                        directory = directory)\n",
    "\n",
    "lr_decay = keras.callbacks.LearningRateScheduler(scheduler)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ba88a88b",
   "metadata": {},
   "source": [
    "## Training Model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fd442a3e",
   "metadata": {},
   "outputs": [],
   "source": [
    "datagen = ImageDataGenerator(width_shift_range=0.1,\n",
    "                            height_shift_range=0.1)\n",
    "\n",
    "training_generator = MixUp.MixupGenerator_3level(dataset['x_train'],\n",
    "                                                 dataset['y_train_coarse'], \n",
    "                                                 dataset['y_train_medium'],\n",
    "                                                 dataset['y_train_fine'],\n",
    "                                                 batch_size=train_params[\"batch_size\"],\n",
    "                                                 alpha=0.2, \n",
    "                                                 datagen=datagen\n",
    "                                                )()"
   ]
  },
  {
   "cell_type": "raw",
   "id": "34b142fe",
   "metadata": {
    "scrolled": true
   },
   "source": [
    "history = model.fit(training_generator,\n",
    "                    steps_per_epoch = int(dataset['x_train'].shape[0] / train_params[\"batch_size\"]),\n",
    "                    epochs = train_params[\"n_epochs\"],\n",
    "                    validation_data = ([dataset['x_test'],\n",
    "                                        dataset['y_test_coarse'],\n",
    "                                        dataset['y_test_medium'],\n",
    "                                        dataset['y_test_fine']],\n",
    "                                        [dataset['y_test_coarse'],\n",
    "                                        dataset['y_test_medium'],\n",
    "                                        dataset['y_test_fine'],\n",
    "                                        dataset['x_test']]),\n",
    "                    validation_batch_size = train_params[\"batch_size\"],\n",
    "                    callbacks = [tb,\n",
    "                                log,\n",
    "                                checkpoint,\n",
    "                                lr_decay,\n",
    "                                change_lw],\n",
    "                    verbose=1)\n",
    "model_save_dir = str(directory+'/trained_model.h5')\n",
    "try:\n",
    "    model.save_weights(model_save_dir)\n",
    "    print('Trained model saved to = ', model_save_dir)\n",
    "except:\n",
    "    print('Model Wight is not saved')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "38159923",
   "metadata": {},
   "outputs": [],
   "source": [
    "history = model.fit_generator(generator=training_generator,\n",
    "                              steps_per_epoch = int(dataset['x_train'].shape[0] / train_params[\"batch_size\"]),\n",
    "                              epochs = train_params[\"n_epochs\"],\n",
    "                              validation_data = ([dataset['x_test'],\n",
    "                                                  dataset['y_test_coarse'],\n",
    "                                                  dataset['y_test_medium'],\n",
    "                                                  dataset['y_test_fine']],\n",
    "                                                 [dataset['y_test_coarse'],\n",
    "                                                  dataset['y_test_medium'],\n",
    "                                                  dataset['y_test_fine'],\n",
    "                                                  dataset['x_test']]),\n",
    "                                  callbacks = [tb,\n",
    "                                               log,\n",
    "                                               checkpoint,\n",
    "                                               lr_decay,\n",
    "                                               change_lw],\n",
    "                                  verbose=1)\n",
    "model_save_dir = str(directory+'/trained_model.h5')\n",
    "try:\n",
    "    model.save_weights(model_save_dir)\n",
    "    print('Trained model saved to = ', model_save_dir)\n",
    "except:\n",
    "    print('Model Wight is not saved')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d3340b1f",
   "metadata": {},
   "source": [
    "## Model Analysis"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4d2bab9d",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "model_analysis = MLmodel.model_analysis(model, dataset)\n",
    "\n",
    "results = model_analysis.evaluate()\n",
    "\n",
    "predictions = model_analysis.prediction()\n",
    "\n",
    "true_label = [dataset['y_test_coarse'],dataset['y_test_medium'],dataset['y_test_fine']]\n",
    "pred_label = [predictions[0],predictions[1],predictions[2]]\n",
    "metrics.lvl_wise_metric(true_label,pred_label)\n",
    "\n",
    "h_measurements,consistency,exact_match = metrics.hmeasurements(true_label,\n",
    "                                       pred_label,\n",
    "                                       dataset['tree'])\n",
    "print('\\nHierarchical Precision =',h_measurements[0],\n",
    "      '\\nHierarchical Recall =', h_measurements[1],\n",
    "      '\\nHierarchical F1-Score =',h_measurements[2],\n",
    "      '\\n Consistency = ', consistency,\n",
    "      '\\n Exact Match = ', exact_match,\n",
    "     )"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d1582f28",
   "metadata": {},
   "source": [
    "# Marine-tree Dataset"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "972d8884",
   "metadata": {},
   "source": [
    "## Import Dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f5230444",
   "metadata": {},
   "outputs": [],
   "source": [
    "dataset = datasets.get_Marine_dataset(output_level='level_depth_3',\n",
    "                                      dataset_path ='D:\\Datasets\\Marine_tree',\n",
    "                                      image_size=(64,64),\n",
    "                                      batch_size=train_params['batch_size'],\n",
    "                                      subtype='Combined',\n",
    "                                      data_normalizing ='normalize',\n",
    "                                      class_encoding = 'One_Hot_Encoder',\n",
    "                                      data_augmantation = 'mixup'\n",
    "                                     )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ea37ca18",
   "metadata": {},
   "outputs": [],
   "source": [
    "for x,y in dataset.train_dataset.take(1):\n",
    "    for i in range(len(x)):\n",
    "        print('Example = ', i)\n",
    "        plt.imshow(x[i])\n",
    "        plt.show()\n",
    "        print('Coarse =', {k:v for k,v in enumerate(y[0][i].numpy()) if v != 0}) # coarse lables\n",
    "        print('Medium =', {k:v for k,v in enumerate(y[1][i].numpy()) if v != 0}) # medium lables\n",
    "        print('Fine   =', {k:v for k,v in enumerate(y[2][i].numpy()) if v != 0}) # fine lables"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "77f6bd54",
   "metadata": {},
   "source": [
    "## Model parameters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5a9918a8",
   "metadata": {},
   "outputs": [],
   "source": [
    "coarse_class, medium_class, fine_class = dataset.num_classes\n",
    "\n",
    "initial_lw = MLmodel.initial_lw({\"coarse\": coarse_class,\n",
    "                                 \"medium\": medium_class,\n",
    "                                 \"fine\": fine_class})\n",
    "\n",
    "lossweight = {'coarse_lw' : K.variable(value = initial_lw['coarse'], dtype=\"float32\", name=\"coarse_lw\"),\n",
    "             'medium_lw' : K.variable(value = initial_lw['medium'], dtype=\"float32\", name=\"medium_lw\"),\n",
    "             'fine_lw' : K.variable(value = initial_lw['fine'], dtype=\"float32\", name=\"fine_lw\"),\n",
    "              'decoder_lw' : model_params['Reconstruction_LW']\n",
    "             }"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c3212255",
   "metadata": {},
   "source": [
    "## Get Model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e93687e2",
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_compiled_model():\n",
    "    model = MLmodel.HCapsNet_3_Level(input_shape = dataset.image_size,\n",
    "                                        no_coarse_class  = coarse_class,\n",
    "                                        no_medium_class  = medium_class,\n",
    "                                        no_fine_class    = fine_class\n",
    "                                       )\n",
    "    model.compile(optimizer='adam',\n",
    "                  loss=[model_params[\"class_loss\"],\n",
    "                        model_params[\"class_loss\"],\n",
    "                        model_params[\"class_loss\"],\n",
    "                        model_params[\"reconstruction_loss\"]],\n",
    "                  loss_weights=[lossweight['coarse_lw'],\n",
    "                                lossweight['medium_lw'],\n",
    "                                lossweight['fine_lw'],\n",
    "                                lossweight['decoder_lw']],\n",
    "                  metrics={'Fine_prediction_output_layer': 'accuracy',\n",
    "                           'Medium_prediction_output_layer': 'accuracy',\n",
    "                           'Coarse_prediction_output_layer': 'accuracy'}\n",
    "                 )\n",
    "    return model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "58462ae1",
   "metadata": {},
   "outputs": [],
   "source": [
    "tf.keras.backend.clear_session() ## clear session"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "793fa892",
   "metadata": {},
   "outputs": [],
   "source": [
    "model = get_compiled_model()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d388a39a",
   "metadata": {},
   "source": [
    "## Log Directory"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b8b1c403",
   "metadata": {},
   "outputs": [],
   "source": [
    "### Log directory\n",
    "directory = sysenv.log_dir(dataset.name+'/'+model.name)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9c6c5843",
   "metadata": {},
   "source": [
    "## Model Summary"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "aa82300f",
   "metadata": {},
   "outputs": [],
   "source": [
    "model.summary()\n",
    "keras.utils.plot_model(model, to_file = directory+\"/H-CapsNet_Marine_tree.png\", show_shapes=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e1d55b70",
   "metadata": {},
   "source": [
    "## Callbacks"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "787d06e8",
   "metadata": {},
   "outputs": [],
   "source": [
    "tb = keras.callbacks.TensorBoard(directory+'/tb_logs')\n",
    "log = keras.callbacks.CSVLogger(directory+'/log.csv', append=True)\n",
    "\n",
    "checkpoint = keras.callbacks.ModelCheckpoint(directory+'/epoch_best.h5',\n",
    "                                             monitor='val_Fine_prediction_output_layer_accuracy',\n",
    "                                             save_best_only=True, save_weights_only=True, verbose=1\n",
    "                                            )\n",
    "\n",
    "change_lw = MLmodel.LossWeightsModifier(lossweight,\n",
    "                                        initial_lw,\n",
    "                                        directory = directory)\n",
    "\n",
    "lr_decay = keras.callbacks.LearningRateScheduler(scheduler)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5237f7f4",
   "metadata": {},
   "source": [
    "## Dataset Training Pipeline"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "69fc08b0",
   "metadata": {},
   "outputs": [],
   "source": [
    "def pipeline_multi_input_output(image, label):\n",
    "    label_0 = label[0]\n",
    "    label_1 = label[1]\n",
    "    label_2 = label[2]\n",
    "    return (image, label_0, label_1, label_2), (label_0, label_1, label_2, image)\n",
    "\n",
    "### matchin X,Y with model input\n",
    "training_dataset_match = dataset.train_dataset.map(pipeline_multi_input_output) ## Mixup dataset\n",
    "val_dataset_match = dataset.val_dataset.map(pipeline_multi_input_output) ## Val Dataset\n",
    "test_dataset_match = dataset.test_dataset.map(pipeline_multi_input_output) ## test Dataset"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b6d35186",
   "metadata": {},
   "source": [
    "## Training Model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a7dece34",
   "metadata": {},
   "outputs": [],
   "source": [
    "history = model.fit(training_dataset_match,\n",
    "                    batch_size = train_params[\"batch_size\"],\n",
    "                    epochs = train_params[\"n_epochs\"],\n",
    "                    validation_data = val_dataset_match,\n",
    "                    callbacks = [tb,\n",
    "                                 log,\n",
    "                                 checkpoint,\n",
    "                                 lr_decay,\n",
    "                                 change_lw],\n",
    "                    verbose=1)\n",
    "model_save_dir = str(directory+'/trained_model.h5')\n",
    "try:\n",
    "    model.save_weights(model_save_dir)\n",
    "    print('Trained model saved to = ', model_save_dir)\n",
    "except:\n",
    "    print('Model Wight is not saved')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "570de1b5",
   "metadata": {},
   "outputs": [],
   "source": [
    "try:\n",
    "    history_dict = history.history\n",
    "    plotter = tfdocs.plots.HistoryPlotter()\n",
    "    plotter.plot({\"Coarse\": history}, metric = \"Coarse_prediction_output_layer_accuracy\")\n",
    "    plotter.plot({\"Medium\": history}, metric = \"Medium_prediction_output_layer_accuracy\")\n",
    "    plotter.plot({\"Fine\": history}, metric = \"Fine_prediction_output_layer_accuracy\")\n",
    "    plt.title(\"Model Accuracy\")\n",
    "    plt.ylim([0,1])\n",
    "except:\n",
    "    print('Not available')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a23264fc",
   "metadata": {},
   "source": [
    "## Model Analysis"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7961b036",
   "metadata": {},
   "outputs": [],
   "source": [
    "results = model.evaluate(test_dataset_match)\n",
    "for n in range(len(results)):\n",
    "    print(str(n+1)+'.',model.metrics_names[n], '==>', results[n])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "16c0b7f6",
   "metadata": {},
   "outputs": [],
   "source": [
    "def predict_from_pipeline(model, dataset):\n",
    "    y_pred_c = []\n",
    "    y_pred_m = []\n",
    "    y_pred_f = []\n",
    "    \n",
    "    y_true_c = []\n",
    "    y_true_m = []\n",
    "    y_true_f = []\n",
    "    for x, y in dataset:\n",
    "        batch_pred = model.predict(x)\n",
    "        \n",
    "        y_true_c.extend(y[0].numpy().tolist())\n",
    "        y_true_m.extend(y[1].numpy().tolist())\n",
    "        y_true_f.extend(y[2].numpy().tolist())\n",
    "        \n",
    "        y_pred_c.extend(batch_pred[0].tolist())\n",
    "        y_pred_m.extend(batch_pred[1].tolist())\n",
    "        y_pred_f.extend(batch_pred[2].tolist())\n",
    "        \n",
    "    return np.array(y_true_c), np.array(y_true_m), np.array(y_true_f), np.array(y_pred_c), np.array(y_pred_m), np.array(y_pred_f)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f6c1dbad",
   "metadata": {},
   "outputs": [],
   "source": [
    "y_true_c, y_true_m, y_true_f, y_pred_c, y_pred_m, y_pred_f = predict_from_pipeline(model, test_dataset_match)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e36e2d07",
   "metadata": {},
   "outputs": [],
   "source": [
    "metrics.lvl_wise_metric([y_true_c, y_true_m, y_true_f],\n",
    "                        [y_pred_c, y_pred_m, y_pred_f])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cfd220f4",
   "metadata": {},
   "outputs": [],
   "source": [
    "h_measurements,consistency,exact_match = metrics.hmeasurements([y_true_c, y_true_m, y_true_f],\n",
    "                                                               [y_pred_c, y_pred_m, y_pred_f],\n",
    "                                                               dataset.get_tree())\n",
    "print('\\nHierarchical Precision =',h_measurements[0],\n",
    "      '\\nHierarchical Recall =', h_measurements[1],\n",
    "      '\\nHierarchical F1-Score =',h_measurements[2],\n",
    "      '\\nConsistency = ', consistency,\n",
    "      '\\nExact Match = ', exact_match,\n",
    "     )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8e0ee8e8",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.8.12"
  },
  "toc": {
   "base_numbering": 1,
   "nav_menu": {},
   "number_sections": true,
   "sideBar": true,
   "skip_h1_title": false,
   "title_cell": "Table of Contents",
   "title_sidebar": "Contents",
   "toc_cell": false,
   "toc_position": {
    "height": "calc(100% - 180px)",
    "left": "10px",
    "top": "150px",
    "width": "203px"
   },
   "toc_section_display": true,
   "toc_window_display": true
  },
  "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": 5
}