Data_Augmentation_Conv6_CIFAR10_Example.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Data Augmentation: Conv-6 CNN LTH on CIFAR-10"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "rPBXLYKhMKTA"
   },
   "outputs": [],
   "source": [
    "import tensorflow as tf\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "import seaborn as sns\n",
    "import math\n",
    "# import tensorflow_model_optimization as tfmot\n",
    "# from tensorflow_model_optimization.sparsity import keras as sparsity\n",
    "# from tensorflow.keras import datasets, layers, models\n",
    "\n",
    "from tensorflow.keras.layers import AveragePooling2D, Conv2D, MaxPooling2D, ReLU\n",
    "from tensorflow.keras import models, layers, datasets\n",
    "from tensorflow.keras.layers import Dense, Flatten, Reshape, Input, InputLayer\n",
    "from tensorflow.keras.models import Sequential, Model\n",
    "from tensorflow.keras.initializers import RandomNormal\n",
    "\n",
    "from sklearn.metrics import accuracy_score, precision_score, recall_score\n",
    "from sklearn.cluster import DBSCAN, KMeans\n",
    "\n",
    "from tensorflow.keras.preprocessing.image import ImageDataGenerator\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 34
    },
    "colab_type": "code",
    "id": "aoIVNDPtMKTU",
    "outputId": "f515a0b1-452a-4130-bfd0-242f3ca88eec"
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'2.2.0'"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "tf.__version__"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "AUTOTUNE = tf.data.experimental.AUTOTUNE"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "CdI869wWMKTi",
    "outputId": "8e55f021-186a-4bf1-aaaa-7be7995a6a43"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "env: CUDA_DEVICE_ORDER=PCI_BUS_ID\n",
      "env: CUDA_VISIBLE_DEVICES=2\n"
     ]
    }
   ],
   "source": [
    "%env CUDA_DEVICE_ORDER=PCI_BUS_ID\n",
    "%env CUDA_VISIBLE_DEVICES = 2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "RpefuTJ2MKTt"
   },
   "outputs": [],
   "source": [
    "batch_size = 60\n",
    "num_classes = 10\n",
    "num_epochs = 100"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 52
    },
    "colab_type": "code",
    "id": "cpovn_jyMKT5",
    "outputId": "ae5ba05c-7890-4b02-8419-b778f82253e6"
   },
   "outputs": [],
   "source": [
    "# Data preprocessing and cleaning:\n",
    "# input image dimensions\n",
    "img_rows, img_cols = 32, 32\n",
    "\n",
    "# Load CIFAR-10 dataset-\n",
    "(X_train, y_train), (X_test, y_test) = tf.keras.datasets.cifar10.load_data()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 52
    },
    "colab_type": "code",
    "id": "zNbVLDY6M3tW",
    "outputId": "00d1ece3-f118-4dd4-a4fa-95e2f4918e01"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "X_train.shape = (50000, 32, 32, 3), y_train.shape = (50000, 1)\n",
      "X_test.shape = (10000, 32, 32, 3), y_test.shape = (10000, 1)\n"
     ]
    }
   ],
   "source": [
    "print(\"X_train.shape = {0}, y_train.shape = {1}\".format(X_train.shape, y_train.shape))\n",
    "print(\"X_test.shape = {0}, y_test.shape = {1}\".format(X_test.shape, y_test.shape))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 69
    },
    "colab_type": "code",
    "id": "XByeYdcwMKUE",
    "outputId": "3961609c-7718-4728-8238-326752312b06"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "'input_shape' which will be used = (32, 32, 3)\n",
      "\n"
     ]
    }
   ],
   "source": [
    "if tf.keras.backend.image_data_format() == 'channels_first':\n",
    "    X_train = X_train.reshape(X_train.shape[0], 3, img_rows, img_cols)\n",
    "    X_test = X_test.reshape(X_test.shape[0], 3, img_rows, img_cols)\n",
    "    input_shape = (3, img_rows, img_cols)\n",
    "else:\n",
    "    X_train = X_train.reshape(X_train.shape[0], img_rows, img_cols, 3)\n",
    "    X_test = X_test.reshape(X_test.shape[0], img_rows, img_cols, 3)\n",
    "    input_shape = (img_rows, img_cols, 3)\n",
    "\n",
    "print(\"\\n'input_shape' which will be used = {0}\\n\".format(input_shape))\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "pahIwSC8MKUO"
   },
   "outputs": [],
   "source": [
    "# Convert datasets to floating point types-\n",
    "X_train = X_train.astype('float32')\n",
    "X_test = X_test.astype('float32')\n",
    "\n",
    "# Normalize the training and testing datasets-\n",
    "X_train /= 255.0\n",
    "X_test /= 255.0"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "hKk5sNp6MKUX"
   },
   "outputs": [],
   "source": [
    "# convert class vectors/target to binary class matrices or one-hot encoded values-\n",
    "y_train = tf.keras.utils.to_categorical(y_train, num_classes)\n",
    "y_test = tf.keras.utils.to_categorical(y_test, num_classes)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 86
    },
    "colab_type": "code",
    "id": "_-GPw1d8MKUs",
    "outputId": "d26d829a-177a-4dc8-fdc1-c20872bcec2d"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "Dimensions of training and testing sets are:\n",
      "X_train.shape = (50000, 32, 32, 3), y_train.shape = (50000, 10)\n",
      "X_test.shape = (10000, 32, 32, 3), y_test.shape = (10000, 10)\n"
     ]
    }
   ],
   "source": [
    "print(\"\\nDimensions of training and testing sets are:\")\n",
    "print(\"X_train.shape = {0}, y_train.shape = {1}\".format(X_train.shape, y_train.shape))\n",
    "print(\"X_test.shape = {0}, y_test.shape = {1}\".format(X_test.shape, y_test.shape))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "M1kapfNVMKU3"
   },
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "OxQX5ma2MKVH"
   },
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "yx1hHreJMKVQ"
   },
   "source": [
    "### Prepare CIFAR10 dataset for _GradientTape_ training:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 46,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "F1gjoSY6MKVT"
   },
   "outputs": [],
   "source": [
    "# Create training and testing datasets-\n",
    "train_dataset = tf.data.Dataset.from_tensor_slices((X_train, y_train))\n",
    "test_dataset = tf.data.Dataset.from_tensor_slices((X_test, y_test))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 47,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_dataset_features = tf.data.Dataset.from_tensor_slices(X_train)\n",
    "train_dataset_labels = tf.data.Dataset.from_tensor_slices(y_train)\n",
    "test_dataset_features = tf.data.Dataset.from_tensor_slices(X_test)\n",
    "test_dataset_labels = tf.data.Dataset.from_tensor_slices(y_test)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 48,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "Nz5RdcIjMKVZ"
   },
   "outputs": [],
   "source": [
    "train_dataset = train_dataset.shuffle(buffer_size = 20000, reshuffle_each_iteration = True).batch(batch_size = batch_size, drop_remainder = False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 49,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "Nl62ZgFkMKVi"
   },
   "outputs": [],
   "source": [
    "test_dataset = test_dataset.batch(batch_size=batch_size, drop_remainder=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for x_t, y_t in test_dataset:\n",
    "    print(x_t.shape, y_t.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 51,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "kvdKx7gAMKVp"
   },
   "outputs": [],
   "source": [
    "# Choose an optimizer and loss function for training-\n",
    "loss_fn = tf.keras.losses.CategoricalCrossentropy()\n",
    "optimizer = tf.keras.optimizers.Adam(lr = 0.0003)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 52,
   "metadata": {},
   "outputs": [],
   "source": [
    "lr = 0.0003"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 53,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "ai3kCIYzMKVu"
   },
   "outputs": [],
   "source": [
    "# Select metrics to measure the error & accuracy of model.\n",
    "# These metrics accumulate the values over epochs and then\n",
    "# print the overall result-\n",
    "train_loss = tf.keras.metrics.Mean(name = 'train_loss')\n",
    "train_accuracy = tf.keras.metrics.CategoricalAccuracy(name = 'train_accuracy')\n",
    "\n",
    "test_loss = tf.keras.metrics.Mean(name = 'test_loss')\n",
    "test_accuracy = tf.keras.metrics.CategoricalAccuracy(name = 'test_accuracy')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "iCz54L2rMKV1"
   },
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Data Augmentation:\n",
    "\n",
    "Write a function to augment the images. Map it over the the dataset.\n",
    "This returns a dataset that augments the data on the fly."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [],
   "source": [
    "\"\"\"\n",
    "def augment(image):\n",
    "    # image,label = convert(image, label)\n",
    "    # image = tf.image.convert_image_dtype(image, tf.float32) # Cast and normalize the image to [0,1]\n",
    "\n",
    "    '''\n",
    "    image = tf.reshape(image, shape = [28 ,28, 1])\n",
    "    image = tf.image.resize_with_crop_or_pad(image, 34, 34) # Add 6 pixels of padding\n",
    "    image = tf.image.random_crop(image, size=[28, 28, 1]) # Random crop back to 28x28\n",
    "    image = tf.image.random_brightness(image, max_delta=0.5) # Random brightness\n",
    "    '''\n",
    "    image = tf.image.rotation_range = 90\n",
    "    image = tf.image.width_shift_range = 0.1\n",
    "    image = tf.image.height_shift_range = 0.1\n",
    "    image = tf.image.horizontal_flip = True\n",
    "    \n",
    "    # rotation_range=90, width_shift_range=0.1, height_shift_range=0.1,horizontal_flip=True\n",
    "\n",
    "    return image\n",
    "    \n",
    "\"\"\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "50000"
      ]
     },
     "execution_count": 20,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# num_train_examples = X_train.shape[0]\n",
    "# num_train_examples"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [],
   "source": [
    "\"\"\"\n",
    "augmented_train_batches = (\n",
    "    train_dataset_features\n",
    "    # train_dataset\n",
    "    # .take(NUM_EXAMPLES)\n",
    "    .cache()\n",
    "    .shuffle(num_train_examples // 4)\n",
    "\n",
    "    # The augmentation is added here.\n",
    "    .map(augment, num_parallel_calls=AUTOTUNE)\n",
    "    .batch(batch_size)\n",
    "    .prefetch(AUTOTUNE)\n",
    ")\n",
    "\n",
    "\n",
    "# Setup the validation dataset. This doesn't change whether or not you're using the augmentation.\n",
    "validation_batches = (\n",
    "    test_dataset_features\n",
    "    # .map(convert, num_parallel_calls=AUTOTUNE)\n",
    "    .batch(batch_size)\n",
    ")\n",
    "\n",
    "\"\"\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Data Augmentation done using _ImageDataGenerator_:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 54,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'\\nfor x, y in datagen.flow(X_train, y_train, batch_size=batch_size, shuffle=True):\\n\\tprint(x.shape, y.shape)\\n'"
      ]
     },
     "execution_count": 54,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Example of using 'tf.keras.preprocessing.image import ImageDataGenerator class's - flow(x, y)':\n",
    "\n",
    "datagen = ImageDataGenerator(\n",
    "    # featurewise_center=True,\n",
    "    # featurewise_std_normalization=True,\n",
    "    rotation_range = 90,\n",
    "    width_shift_range = 0.1,\n",
    "    height_shift_range = 0.1,\n",
    "    horizontal_flip = True\n",
    ")\n",
    "\n",
    "\n",
    "# flow():\n",
    "# Takes data & label arrays, generates batches of augmented data.\n",
    "\n",
    "# datagen.flow(X_train, y_train, batch_size=batch_size, shuffle=True)\n",
    "'''\n",
    "for x, y in datagen.flow(X_train, y_train, batch_size=batch_size, shuffle=True):\n",
    "\tprint(x.shape, y.shape)\n",
    "'''"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "NcJZj1ZfMKV8"
   },
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "-TqIXiLpMKWB"
   },
   "outputs": [],
   "source": [
    "def conv6_cnn():\n",
    "    \"\"\"\n",
    "    Function to define the architecture of a neural network model\n",
    "    following Conv-6 architecture for CIFAR-10 dataset and using\n",
    "    provided parameter which are used to prune the model.\n",
    "    \n",
    "    Conv-6 architecture-\n",
    "    64, 64, pool  -- convolutional layers\n",
    "    128, 128, pool -- convolutional layers\n",
    "    256, 256, pool -- convolutional layers\n",
    "    256, 256, 10  -- fully connected layers\n",
    "    \n",
    "    Output: Returns designed and compiled neural network model\n",
    "    \"\"\"\n",
    "    \n",
    "    l = tf.keras.layers\n",
    "    \n",
    "    model = Sequential()\n",
    "    \n",
    "    model.add(\n",
    "        Conv2D(\n",
    "            filters = 64, kernel_size = (3, 3),\n",
    "            activation='relu', kernel_initializer = tf.initializers.GlorotNormal(),\n",
    "            strides = (1, 1), padding = 'same',\n",
    "            input_shape=(32, 32, 3)\n",
    "        )    \n",
    "    )\n",
    "        \n",
    "    model.add(\n",
    "        Conv2D(\n",
    "            filters = 64, kernel_size = (3, 3),\n",
    "            activation='relu', kernel_initializer = tf.initializers.GlorotNormal(),\n",
    "            strides = (1, 1), padding = 'same'\n",
    "        )\n",
    "    )\n",
    "    \n",
    "    model.add(\n",
    "        MaxPooling2D(\n",
    "            pool_size = (2, 2),\n",
    "            strides = (2, 2)\n",
    "        )\n",
    "    )\n",
    "    \n",
    "    model.add(\n",
    "        Conv2D(\n",
    "            filters = 128, kernel_size = (3, 3),\n",
    "            activation='relu', kernel_initializer = tf.initializers.GlorotNormal(),\n",
    "            strides = (1, 1), padding = 'same'\n",
    "        )\n",
    "    )\n",
    "\n",
    "    model.add(\n",
    "        Conv2D(\n",
    "            filters = 128, kernel_size = (3, 3),\n",
    "            activation='relu', kernel_initializer = tf.initializers.GlorotNormal(),\n",
    "            strides = (1, 1), padding = 'same'\n",
    "        )\n",
    "    )\n",
    "\n",
    "    model.add(\n",
    "        MaxPooling2D(\n",
    "            pool_size = (2, 2),\n",
    "            strides = (2, 2)\n",
    "        )\n",
    "    )\n",
    "\n",
    "    model.add(\n",
    "        Conv2D(\n",
    "            filters = 256, kernel_size = (3, 3),\n",
    "            activation='relu', kernel_initializer = tf.initializers.GlorotNormal(),\n",
    "            strides = (1, 1), padding = 'same'\n",
    "        )\n",
    "    )\n",
    "\n",
    "    model.add(\n",
    "        Conv2D(\n",
    "            filters = 256, kernel_size = (3, 3),\n",
    "            activation='relu', kernel_initializer = tf.initializers.GlorotNormal(),\n",
    "            strides = (1, 1), padding = 'same'\n",
    "        )\n",
    "    )\n",
    "\n",
    "    model.add(\n",
    "        MaxPooling2D(\n",
    "            pool_size = (2, 2),\n",
    "            strides = (2, 2)\n",
    "        )\n",
    "    )\n",
    "    \n",
    "    model.add(Flatten())\n",
    "    \n",
    "    model.add(\n",
    "        Dense(\n",
    "            units = 256, activation='relu',\n",
    "            kernel_initializer = tf.initializers.GlorotNormal()\n",
    "        )\n",
    "    )\n",
    "    \n",
    "    model.add(\n",
    "        Dense(\n",
    "            units = 256, activation='relu',\n",
    "            kernel_initializer = tf.initializers.GlorotNormal()\n",
    "        )\n",
    "    )\n",
    "    \n",
    "    model.add(\n",
    "        Dense(\n",
    "            units = 10, activation='softmax'\n",
    "        )\n",
    "    )\n",
    "    \n",
    "\n",
    "    # Compile pruned CNN-\n",
    "    model.compile(\n",
    "        loss=tf.keras.losses.categorical_crossentropy,\n",
    "        # optimizer='adam',\n",
    "        optimizer=tf.keras.optimizers.Adam(lr = 0.0003),\n",
    "        metrics=['accuracy']\n",
    "    )\n",
    "    \n",
    "    \n",
    "    return model\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "n1Q8QnB2MKWJ"
   },
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Instantiate a new Conv-2 CNN model-\n",
    "orig_model = conv6_cnn()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load random weights from before-\n",
    "orig_model.load_weights(\"Conv_6_CIFAR10_Magnitude_Based_Winning_Ticket_Distribution_92.55423622890814.h5\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "Total # of non-zero parameters = 167210\n",
      "\n"
     ]
    }
   ],
   "source": [
    "# Count number of non-zero parameters-\n",
    "winning_params = 0\n",
    "\n",
    "for layer in orig_model.trainable_weights:\n",
    "    nonzeroparams = tf.math.count_nonzero(layer, axis = None).numpy()\n",
    "    # print(\"layer: {0} has {1} non-zero parameters\".format(layer.shape, nonzeroparams))\n",
    "    winning_params += nonzeroparams\n",
    "\n",
    "print(\"\\nTotal # of non-zero parameters = {0}\\n\".format(winning_params))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "Number of training weights = 2262602 and non-trainabel weights = 0.0\n",
      "\n",
      "Total number of parameters = 2262602.0\n",
      "\n"
     ]
    }
   ],
   "source": [
    "import tensorflow.keras.backend as K\n",
    "\n",
    "\n",
    "# METHOD-1: This also counts biases\n",
    "\n",
    "trainable_wts = np.sum([K.count_params(w) for w in orig_model.trainable_weights])\n",
    "non_trainable_wts = np.sum([K.count_params(w) for w in orig_model.non_trainable_weights])\n",
    "\n",
    "print(\"\\nNumber of training weights = {0} and non-trainabel weights = {1}\\n\".format(\n",
    "    trainable_wts, non_trainable_wts\n",
    "))\n",
    "print(\"Total number of parameters = {0}\\n\".format(trainable_wts + non_trainable_wts))\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "92.6098% of parameters have been pruned\n",
      "\n"
     ]
    }
   ],
   "source": [
    "print(\"\\n{0:.4f}% of parameters have been pruned\\n\".format((trainable_wts - winning_params) / trainable_wts * 100))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create mask using winning ticket-\n",
    "\n",
    "# Instantiate a new neural network model for which, the mask is to be created,\n",
    "mask_model = conv6_cnn()\n",
    "    \n",
    "# Load weights of PRUNED model-\n",
    "mask_model.set_weights(orig_model.get_weights())\n",
    "    \n",
    "# For each layer, for each weight which is 0, leave it, as is.\n",
    "# And for weights which survive the pruning,reinitialize it to ONE (1)-\n",
    "for wts in mask_model.trainable_weights:\n",
    "    wts.assign(tf.where(tf.equal(wts, 0.), 0., 1.))\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "Total # of non-zero masks = 167210\n",
      "\n"
     ]
    }
   ],
   "source": [
    "# Count number of non-zero masks-\n",
    "mask_params = 0\n",
    "\n",
    "for layer in mask_model.trainable_weights:\n",
    "    nonzeroparams = tf.math.count_nonzero(layer, axis = None).numpy()\n",
    "    # print(\"layer: {0} has {1} non-zero masks\".format(layer.shape, nonzeroparams))\n",
    "    mask_params += nonzeroparams\n",
    "\n",
    "print(\"\\nTotal # of non-zero masks = {0}\\n\".format(mask_params))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "number of non-zero parameters and masks matches!\n"
     ]
    }
   ],
   "source": [
    "if mask_params == winning_params:\n",
    "    print(\"\\nnumber of non-zero parameters and masks matches!\")\n",
    "else:\n",
    "    print(\"\\nERROR! number of non-zero parameters and masks DO NOT MATCH!\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Train winning ticket using Data Augmentation:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 61,
   "metadata": {},
   "outputs": [],
   "source": [
    "# User input parameters for Early Stopping in manual implementation-\n",
    "minimum_delta = 0.001\n",
    "patience = 3"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 62,
   "metadata": {},
   "outputs": [],
   "source": [
    "best_val_loss = 100\n",
    "loc_patience = 0"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 48,
   "metadata": {},
   "outputs": [],
   "source": [
    "# best_trained_weights = []"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 63,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Initialize a new LeNet-300-100 model-\n",
    "winning_ticket_model = conv6_cnn()\n",
    "\n",
    "# Load weights of winning ticket-\n",
    "winning_ticket_model.set_weights(orig_model.get_weights())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 64,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define 'train_one_step()' and 'test_step()' functions here-\n",
    "@tf.function\n",
    "def train_one_step(model, mask_model, optimizer, x, y):\n",
    "    '''\n",
    "    Function to compute one step of gradient descent optimization\n",
    "    '''\n",
    "    with tf.GradientTape() as tape:\n",
    "        # Make predictions using defined model-\n",
    "        y_pred = model(x)\n",
    "\n",
    "        # Compute loss-\n",
    "        loss = loss_fn(y, y_pred)\n",
    "        \n",
    "    # Compute gradients wrt defined loss and weights and biases-\n",
    "    grads = tape.gradient(loss, model.trainable_variables)\n",
    "    \n",
    "    # type(grads)\n",
    "    # list\n",
    "    \n",
    "    # List to hold element-wise multiplication between-\n",
    "    # computed gradient and masks-\n",
    "    grad_mask_mul = []\n",
    "    \n",
    "    # Perform element-wise multiplication between computed gradients and masks-\n",
    "    for grad_layer, mask in zip(grads, mask_model.trainable_weights):\n",
    "        grad_mask_mul.append(tf.math.multiply(grad_layer, mask))\n",
    "    \n",
    "    # Apply computed gradients to model's weights and biases-\n",
    "    optimizer.apply_gradients(zip(grad_mask_mul, model.trainable_variables))\n",
    "\n",
    "    # Compute accuracy-\n",
    "    train_loss(loss)\n",
    "    train_accuracy(y, y_pred)\n",
    "\n",
    "    return None\n",
    "    \n",
    "    \n",
    "@tf.function\n",
    "def test_step(model, optimizer, data, labels):\n",
    "    \"\"\"\n",
    "    Function to test model performance\n",
    "    on testing dataset\n",
    "    \"\"\"\n",
    "    \n",
    "    predictions = model(data)\n",
    "    t_loss = loss_fn(labels, predictions)\n",
    "\n",
    "    test_loss(t_loss)\n",
    "    test_accuracy(labels, predictions)\n",
    "\n",
    "    return None\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 59,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "Terminating training (datagen.flow())\n"
     ]
    }
   ],
   "source": [
    "curr_step = 0\n",
    "    \n",
    "for x, y in datagen.flow(X_train, y_train, batch_size = batch_size, shuffle = True):\n",
    "    # print(\"x.shape = {0}, y.shape = {1}\".format(x.shape, y.shape))\n",
    "    # x.shape = (60, 32, 32, 3), y.shape = (60, 10)\n",
    "    \n",
    "    train_one_step(winning_ticket_model, mask_model, optimizer, x, y)\n",
    "    # print(\"current step = \", curr_step)\n",
    "    curr_step += 1\n",
    "        \n",
    "    if curr_step >= X_train.shape[0] // batch_size:\n",
    "        print(\"\\nTerminating training (datagen.flow())\")\n",
    "        break\n",
    " "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 60,
   "metadata": {},
   "outputs": [],
   "source": [
    "for x_t, y_t in test_dataset:\n",
    "    test_step(winning_ticket_model, optimizer, x_t, y_t)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 65,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "Terminating training (datagen.flow())\n",
      "Epoch 1, Loss: 1.8008, Accuracy: 34.9520, Test Loss: 1.6015, Test Accuracy: 42.500000\n",
      "Total number of trainable parameters = 167210\n",
      "\n",
      "\n",
      "Terminating training (datagen.flow())\n",
      "Epoch 2, Loss: 1.5840, Accuracy: 42.9432, Test Loss: 1.5690, Test Accuracy: 44.250000\n",
      "Total number of trainable parameters = 167210\n",
      "\n",
      "\n",
      "Terminating training (datagen.flow())\n",
      "Epoch 3, Loss: 1.4510, Accuracy: 48.2013, Test Loss: 1.2719, Test Accuracy: 55.349998\n",
      "Total number of trainable parameters = 167210\n",
      "\n",
      "\n",
      "Terminating training (datagen.flow())\n",
      "Epoch 4, Loss: 1.3696, Accuracy: 51.1345, Test Loss: 1.2625, Test Accuracy: 56.690002\n",
      "Total number of trainable parameters = 167210\n",
      "\n",
      "\n",
      "Terminating training (datagen.flow())\n",
      "Epoch 5, Loss: 1.3086, Accuracy: 53.5774, Test Loss: 1.1230, Test Accuracy: 60.709999\n",
      "Total number of trainable parameters = 167210\n",
      "\n",
      "\n",
      "Terminating training (datagen.flow())\n",
      "Epoch 6, Loss: 1.2606, Accuracy: 55.5502, Test Loss: 1.0963, Test Accuracy: 62.150002\n",
      "Total number of trainable parameters = 167210\n",
      "\n",
      "\n",
      "Terminating training (datagen.flow())\n",
      "Epoch 7, Loss: 1.2259, Accuracy: 57.1509, Test Loss: 1.1065, Test Accuracy: 61.970001\n",
      "Total number of trainable parameters = 167210\n",
      "\n",
      "\n",
      "Terminating training (datagen.flow())\n",
      "Epoch 8, Loss: 1.1937, Accuracy: 58.0152, Test Loss: 1.1376, Test Accuracy: 61.129997\n",
      "Total number of trainable parameters = 167210\n",
      "\n",
      "\n",
      "Terminating training (datagen.flow())\n",
      "Epoch 9, Loss: 1.1714, Accuracy: 59.1297, Test Loss: 1.0177, Test Accuracy: 65.060005\n",
      "Total number of trainable parameters = 167210\n",
      "\n",
      "\n",
      "Terminating training (datagen.flow())\n",
      "Epoch 10, Loss: 1.1493, Accuracy: 59.7399, Test Loss: 1.0039, Test Accuracy: 65.669998\n",
      "Total number of trainable parameters = 167210\n",
      "\n",
      "\n",
      "Terminating training (datagen.flow())\n",
      "Epoch 11, Loss: 1.1352, Accuracy: 60.1180, Test Loss: 1.0073, Test Accuracy: 65.480003\n",
      "Total number of trainable parameters = 167210\n",
      "\n",
      "\n",
      "Terminating training (datagen.flow())\n",
      "Epoch 12, Loss: 1.1113, Accuracy: 61.2465, Test Loss: 1.0039, Test Accuracy: 65.589996\n",
      "Total number of trainable parameters = 167210\n",
      "\n",
      "\n",
      "Terminating training (datagen.flow())\n",
      "Epoch 13, Loss: 1.0941, Accuracy: 61.8607, Test Loss: 1.0154, Test Accuracy: 65.300003\n",
      "Total number of trainable parameters = 167210\n",
      "\n",
      "\n",
      "Terminating training (datagen.flow())\n",
      "Epoch 14, Loss: 1.0809, Accuracy: 62.3489, Test Loss: 0.9560, Test Accuracy: 67.549995\n",
      "Total number of trainable parameters = 167210\n",
      "\n",
      "\n",
      "Terminating training (datagen.flow())\n",
      "Epoch 15, Loss: 1.0665, Accuracy: 63.1453, Test Loss: 0.9719, Test Accuracy: 67.540001\n",
      "Total number of trainable parameters = 167210\n",
      "\n",
      "\n",
      "Terminating training (datagen.flow())\n",
      "Epoch 16, Loss: 1.0555, Accuracy: 63.3333, Test Loss: 0.9923, Test Accuracy: 66.529999\n",
      "Total number of trainable parameters = 167210\n",
      "\n",
      "\n",
      "Terminating training (datagen.flow())\n",
      "Epoch 17, Loss: 1.0473, Accuracy: 63.7315, Test Loss: 0.9861, Test Accuracy: 66.430000\n",
      "Total number of trainable parameters = 167210\n",
      "\n",
      "\n",
      "'EarlyStopping' called!\n",
      "\n"
     ]
    }
   ],
   "source": [
    "# Train model using 'GradientTape'-\n",
    "    \n",
    "# Initialize parameters for Early Stopping manual implementation-\n",
    "# best_val_loss = 100\n",
    "# loc_patience = 0\n",
    "    \n",
    "for epoch in range(num_epochs):\n",
    "    \n",
    "    if loc_patience >= patience:\n",
    "        print(\"\\n'EarlyStopping' called!\\n\")\n",
    "        break\n",
    "        \n",
    "    # Reset the metrics at the start of the next epoch\n",
    "    train_loss.reset_states()\n",
    "    train_accuracy.reset_states()\n",
    "    test_loss.reset_states()\n",
    "    test_accuracy.reset_states()\n",
    "    \n",
    "    curr_step = 0\n",
    "    \n",
    "    for x, y in datagen.flow(X_train, y_train, batch_size = batch_size, shuffle = True):\n",
    "    # for x, y in zip(augmented_train_batches, y_train):\n",
    "        train_one_step(winning_ticket_model, mask_model, optimizer, x, y)\n",
    "        # print(\"current step = \", curr_step)\n",
    "        curr_step += 1\n",
    "        \n",
    "        if curr_step >= X_train.shape[0] // batch_size:\n",
    "            print(\"\\nTerminating training (datagen.flow())\")\n",
    "            break\n",
    "\n",
    "\n",
    "    for x_t, y_t in test_dataset:\n",
    "    # for x_t, y_t in zip(validation_batches, y_test):\n",
    "        test_step(winning_ticket_model, optimizer, x_t, y_t)\n",
    "    \n",
    "    '''\n",
    "    for x, y in train_dataset:\n",
    "        train_one_step(winning_ticket_model, mask_model, optimizer, x, y)\n",
    "\n",
    "\n",
    "    for x_t, y_t in test_dataset:\n",
    "        test_step(winning_ticket_model, optimizer, x_t, y_t)\n",
    "\n",
    "    '''\n",
    "\n",
    "    \n",
    "    template = 'Epoch {0}, Loss: {1:.4f}, Accuracy: {2:.4f}, Test Loss: {3:.4f}, Test Accuracy: {4:4f}'\n",
    "    \n",
    "    '''\n",
    "    # 'i' is the index for number of pruning rounds-\n",
    "    history_main[i]['accuracy'][epoch] = train_accuracy.result() * 100\n",
    "    history_main[i]['loss'][epoch] = train_loss.result()\n",
    "    history_main[i]['val_loss'][epoch] = test_loss.result()\n",
    "    history_main[i]['val_accuracy'][epoch] = test_accuracy.result() * 100\n",
    "    ''' \n",
    "\n",
    "    print(template.format(\n",
    "        epoch + 1, train_loss.result(),\n",
    "        train_accuracy.result()*100, test_loss.result(),\n",
    "        test_accuracy.result()*100)\n",
    "         )\n",
    "    \n",
    "    # Count number of non-zero parameters in each layer and in total-\n",
    "    # print(\"layer-wise manner model, number of nonzero parameters in each layer are: \\n\")\n",
    "    model_sum_params = 0\n",
    "    \n",
    "    for layer in winning_ticket_model.trainable_weights:\n",
    "        # print(tf.math.count_nonzero(layer, axis = None).numpy())\n",
    "        model_sum_params += tf.math.count_nonzero(layer, axis = None).numpy()\n",
    "    \n",
    "    print(\"Total number of trainable parameters = {0}\\n\".format(model_sum_params))\n",
    "    \n",
    "    '''\n",
    "    if test_loss.result() < best_val_loss:\n",
    "        best_trained_weights = copy.deepcopy(winning_ticket_model.get_weights())\n",
    "        print(\"\\ntest_loss.result = {0:.4f}, best_val_loss = {1:.4f}, copied 'best weights'\\n\".format(test_loss.result(), best_val_loss))\n",
    "    '''\n",
    "    \n",
    "    # Code for manual Early Stopping:\n",
    "    if np.abs(test_loss.result() < best_val_loss) >= minimum_delta:\n",
    "        # update 'best_val_loss' variable to lowest loss encountered so far-\n",
    "        best_val_loss = test_loss.result()\n",
    "        \n",
    "        # reset 'loc_patience' variable-\n",
    "        loc_patience = 0\n",
    "        \n",
    "    else:  # there is no improvement in monitored metric 'val_loss'\n",
    "        loc_patience += 1  # number of epochs without any improvement\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 66,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Make predictions using trained model-\n",
    "y_pred = np.argmax(winning_ticket_model.predict(X_test), axis = -1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 68,
   "metadata": {},
   "outputs": [],
   "source": [
    "y_test_np = np.argmax(y_test, axis = 1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 69,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "((10000,), (10000,))"
      ]
     },
     "execution_count": 69,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "y_pred.shape, y_test_np.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 70,
   "metadata": {},
   "outputs": [],
   "source": [
    "accuracy = accuracy_score(y_test_np, y_pred)\n",
    "precision = precision_score(y_test_np, y_pred, average = 'macro')\n",
    "recall = recall_score(y_test_np, y_pred, average = 'macro')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 71,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "Conv-6 CNN (Winning Ticket Trained model) metrics:\n",
      "accuracy = 0.6643, precision = 0.6790 & recall = 0.6643\n",
      "\n"
     ]
    }
   ],
   "source": [
    "print(\"\\nConv-6 CNN (Winning Ticket Trained model) metrics:\")\n",
    "print(\"accuracy = {0:.4f}, precision = {1:.4f} & recall = {2:.4f}\\n\".format(accuracy, precision, recall))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Train winning ticket without Data Augmentation:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 72,
   "metadata": {},
   "outputs": [],
   "source": [
    "# User input parameters for Early Stopping in manual implementation-\n",
    "minimum_delta = 0.001\n",
    "patience = 3"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 73,
   "metadata": {},
   "outputs": [],
   "source": [
    "best_val_loss = 100\n",
    "loc_patience = 0"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 74,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Initialize a new LeNet-300-100 model-\n",
    "winning_ticket_model = conv6_cnn()\n",
    "\n",
    "# Load weights of winning ticket-\n",
    "winning_ticket_model.set_weights(orig_model.get_weights())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 75,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define 'train_one_step()' and 'test_step()' functions here-\n",
    "@tf.function\n",
    "def train_one_step(model, mask_model, optimizer, x, y):\n",
    "    '''\n",
    "    Function to compute one step of gradient descent optimization\n",
    "    '''\n",
    "    with tf.GradientTape() as tape:\n",
    "        # Make predictions using defined model-\n",
    "        y_pred = model(x)\n",
    "\n",
    "        # Compute loss-\n",
    "        loss = loss_fn(y, y_pred)\n",
    "        \n",
    "    # Compute gradients wrt defined loss and weights and biases-\n",
    "    grads = tape.gradient(loss, model.trainable_variables)\n",
    "    \n",
    "    # type(grads)\n",
    "    # list\n",
    "    \n",
    "    # List to hold element-wise multiplication between-\n",
    "    # computed gradient and masks-\n",
    "    grad_mask_mul = []\n",
    "    \n",
    "    # Perform element-wise multiplication between computed gradients and masks-\n",
    "    for grad_layer, mask in zip(grads, mask_model.trainable_weights):\n",
    "        grad_mask_mul.append(tf.math.multiply(grad_layer, mask))\n",
    "    \n",
    "    # Apply computed gradients to model's weights and biases-\n",
    "    optimizer.apply_gradients(zip(grad_mask_mul, model.trainable_variables))\n",
    "\n",
    "    # Compute accuracy-\n",
    "    train_loss(loss)\n",
    "    train_accuracy(y, y_pred)\n",
    "\n",
    "    return None\n",
    "    \n",
    "    \n",
    "@tf.function\n",
    "def test_step(model, optimizer, data, labels):\n",
    "    \"\"\"\n",
    "    Function to test model performance\n",
    "    on testing dataset\n",
    "    \"\"\"\n",
    "    \n",
    "    predictions = model(data)\n",
    "    t_loss = loss_fn(labels, predictions)\n",
    "\n",
    "    test_loss(t_loss)\n",
    "    test_accuracy(labels, predictions)\n",
    "\n",
    "    return None\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 77,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 1, Loss: 1.1903, Accuracy: 57.5960, Test Loss: 0.9175, Test Accuracy: 67.940002\n",
      "Total number of trainable parameters = 167210\n",
      "\n",
      "Epoch 2, Loss: 0.7849, Accuracy: 72.7580, Test Loss: 0.7680, Test Accuracy: 73.930000\n",
      "Total number of trainable parameters = 167210\n",
      "\n",
      "Epoch 3, Loss: 0.6472, Accuracy: 77.6740, Test Loss: 0.7098, Test Accuracy: 76.059998\n",
      "Total number of trainable parameters = 167210\n",
      "\n",
      "Epoch 4, Loss: 0.5587, Accuracy: 80.6160, Test Loss: 0.6894, Test Accuracy: 76.779999\n",
      "Total number of trainable parameters = 167210\n",
      "\n",
      "Epoch 5, Loss: 0.4946, Accuracy: 83.0480, Test Loss: 0.6635, Test Accuracy: 77.930000\n",
      "Total number of trainable parameters = 167210\n",
      "\n",
      "Epoch 6, Loss: 0.4456, Accuracy: 84.5440, Test Loss: 0.6789, Test Accuracy: 77.910004\n",
      "Total number of trainable parameters = 167210\n",
      "\n",
      "Epoch 7, Loss: 0.4038, Accuracy: 85.9980, Test Loss: 0.6579, Test Accuracy: 78.680000\n",
      "Total number of trainable parameters = 167210\n",
      "\n",
      "Epoch 8, Loss: 0.3704, Accuracy: 87.1760, Test Loss: 0.6524, Test Accuracy: 79.159996\n",
      "Total number of trainable parameters = 167210\n",
      "\n",
      "Epoch 9, Loss: 0.3355, Accuracy: 88.4280, Test Loss: 0.6851, Test Accuracy: 78.759995\n",
      "Total number of trainable parameters = 167210\n",
      "\n",
      "Epoch 10, Loss: 0.3133, Accuracy: 89.2780, Test Loss: 0.7128, Test Accuracy: 78.339996\n",
      "Total number of trainable parameters = 167210\n",
      "\n",
      "Epoch 11, Loss: 0.2852, Accuracy: 90.2540, Test Loss: 0.7145, Test Accuracy: 79.220001\n",
      "Total number of trainable parameters = 167210\n",
      "\n",
      "\n",
      "'EarlyStopping' called!\n",
      "\n"
     ]
    }
   ],
   "source": [
    "# Train model using 'GradientTape'-\n",
    "    \n",
    "# Initialize parameters for Early Stopping manual implementation-\n",
    "# best_val_loss = 100\n",
    "# loc_patience = 0\n",
    "    \n",
    "for epoch in range(num_epochs):\n",
    "    \n",
    "    if loc_patience >= patience:\n",
    "        print(\"\\n'EarlyStopping' called!\\n\")\n",
    "        break\n",
    "        \n",
    "    # Reset the metrics at the start of the next epoch\n",
    "    train_loss.reset_states()\n",
    "    train_accuracy.reset_states()\n",
    "    test_loss.reset_states()\n",
    "    test_accuracy.reset_states()\n",
    "    \n",
    "    # curr_step = 0\n",
    "    \n",
    "    for x, y in train_dataset:\n",
    "        train_one_step(winning_ticket_model, mask_model, optimizer, x, y)\n",
    "    \n",
    "    '''\n",
    "    # for x, y in zip(augmented_train_batches, y_train):\n",
    "    for x, y in datagen.flow(X_train, y_train, batch_size = batch_size, shuffle = True):\n",
    "        train_one_step(winning_ticket_model, mask_model, optimizer, x, y)\n",
    "        # print(\"current step = \", curr_step)\n",
    "        curr_step += 1\n",
    "        \n",
    "        if curr_step >= X_train.shape[0] // batch_size:\n",
    "            print(\"\\nTerminating training (datagen.flow())\")\n",
    "            break\n",
    "    '''\n",
    "\n",
    "    for x_t, y_t in test_dataset:\n",
    "    # for x_t, y_t in zip(validation_batches, y_test):\n",
    "        test_step(winning_ticket_model, optimizer, x_t, y_t)\n",
    "    \n",
    "    '''\n",
    "    for x, y in train_dataset:\n",
    "        train_one_step(winning_ticket_model, mask_model, optimizer, x, y)\n",
    "\n",
    "\n",
    "    for x_t, y_t in test_dataset:\n",
    "        test_step(winning_ticket_model, optimizer, x_t, y_t)\n",
    "\n",
    "    '''\n",
    "\n",
    "    \n",
    "    template = 'Epoch {0}, Loss: {1:.4f}, Accuracy: {2:.4f}, Test Loss: {3:.4f}, Test Accuracy: {4:4f}'\n",
    "    \n",
    "    '''\n",
    "    # 'i' is the index for number of pruning rounds-\n",
    "    history_main[i]['accuracy'][epoch] = train_accuracy.result() * 100\n",
    "    history_main[i]['loss'][epoch] = train_loss.result()\n",
    "    history_main[i]['val_loss'][epoch] = test_loss.result()\n",
    "    history_main[i]['val_accuracy'][epoch] = test_accuracy.result() * 100\n",
    "    ''' \n",
    "\n",
    "    print(template.format(\n",
    "        epoch + 1, train_loss.result(),\n",
    "        train_accuracy.result()*100, test_loss.result(),\n",
    "        test_accuracy.result()*100)\n",
    "         )\n",
    "    \n",
    "    # Count number of non-zero parameters in each layer and in total-\n",
    "    # print(\"layer-wise manner model, number of nonzero parameters in each layer are: \\n\")\n",
    "    model_sum_params = 0\n",
    "    \n",
    "    for layer in winning_ticket_model.trainable_weights:\n",
    "        # print(tf.math.count_nonzero(layer, axis = None).numpy())\n",
    "        model_sum_params += tf.math.count_nonzero(layer, axis = None).numpy()\n",
    "    \n",
    "    print(\"Total number of trainable parameters = {0}\\n\".format(model_sum_params))\n",
    "    \n",
    "    '''\n",
    "    if test_loss.result() < best_val_loss:\n",
    "        best_trained_weights = copy.deepcopy(winning_ticket_model.get_weights())\n",
    "        print(\"\\ntest_loss.result = {0:.4f}, best_val_loss = {1:.4f}, copied 'best weights'\\n\".format(test_loss.result(), best_val_loss))\n",
    "    '''\n",
    "    \n",
    "    # Code for manual Early Stopping:\n",
    "    if np.abs(test_loss.result() < best_val_loss) >= minimum_delta:\n",
    "        # update 'best_val_loss' variable to lowest loss encountered so far-\n",
    "        best_val_loss = test_loss.result()\n",
    "        \n",
    "        # reset 'loc_patience' variable-\n",
    "        loc_patience = 0\n",
    "        \n",
    "    else:  # there is no improvement in monitored metric 'val_loss'\n",
    "        loc_patience += 1  # number of epochs without any improvement\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "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.7.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}