Recreating_Error.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Executing _@tf.function_ annotated function second time gives error:\n",
    "The following code uses TensorFlow 2.0 along with Python 3.7.5, GradientTape and _tensorflow_model_optimization_ for model optimization (model pruning) for MNIST dataset classification"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import tensorflow as tf\n",
    "import numpy as np\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",
    "import matplotlib.pyplot as plt\n",
    "from tensorflow.keras.layers import AveragePooling2D, Conv2D\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",
    "# import math\n",
    "from sklearn.metrics import accuracy_score, precision_score, recall_score"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'2.0.0'"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "tf.__version__"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "batch_size = 32\n",
    "num_classes = 10\n",
    "num_epochs = 50"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Data preprocessing and cleadning:\n",
    "# input image dimensions\n",
    "img_rows, img_cols = 28, 28\n",
    "\n",
    "# Load MNIST dataset-\n",
    "(X_train, y_train), (X_test, y_test) = tf.keras.datasets.mnist.load_data()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "'input_shape' which will be used = (28, 28, 1)\n",
      "\n"
     ]
    }
   ],
   "source": [
    "if tf.keras.backend.image_data_format() == 'channels_first':\n",
    "    X_train = X_train.reshape(X_train.shape[0], 1, img_rows, img_cols)\n",
    "    X_test = X_test.reshape(X_test.shape[0], 1, img_rows, img_cols)\n",
    "    input_shape = (1, img_rows, img_cols)\n",
    "else:\n",
    "    X_train = X_train.reshape(X_train.shape[0], img_rows, img_cols, 1)\n",
    "    X_test = X_test.reshape(X_test.shape[0], img_rows, img_cols, 1)\n",
    "    input_shape = (img_rows, img_cols, 1)\n",
    "\n",
    "print(\"\\n'input_shape' which will be used = {0}\\n\".format(input_shape))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "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": 7,
   "metadata": {},
   "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": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Reshape training and testing sets-\n",
    "X_train = X_train.reshape(X_train.shape[0], 784)\n",
    "X_test = X_test.reshape(X_test.shape[0], 784)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "Dimensions of training and testing sets are:\n",
      "X_train.shape = (60000, 784), y_train = (60000, 10)\n",
      "X_test.shape = (10000, 784), y_test = (10000, 10)\n"
     ]
    }
   ],
   "source": [
    "print(\"\\nDimensions of training and testing sets are:\")\n",
    "print(\"X_train.shape = {0}, y_train = {1}\".format(X_train.shape, y_train.shape))\n",
    "print(\"X_test.shape = {0}, y_test = {1}\".format(X_test.shape, y_test.shape))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "'end_step parameter' for this dataset =  93750\n"
     ]
    }
   ],
   "source": [
    "# The model is first trained without any pruning for 'num_epochs' epochs-\n",
    "epochs = num_epochs\n",
    "\n",
    "num_train_samples = X_train.shape[0]\n",
    "\n",
    "end_step = np.ceil(1.0 * num_train_samples / batch_size).astype(np.int32) * epochs\n",
    "\n",
    "print(\"'end_step parameter' for this dataset =  {0}\".format(end_step))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Specify the parameters to be used for layer-wise pruning, NO PRUNING is done here:\n",
    "pruning_params_unpruned = {\n",
    "    'pruning_schedule': sparsity.ConstantSparsity(\n",
    "        target_sparsity=0.0, begin_step=0,\n",
    "        end_step = 0, frequency=100\n",
    "    )\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "l = tf.keras.layers"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "def pruned_nn(pruning_params):\n",
    "    \"\"\"\n",
    "    Function to define the architecture of a neural network model\n",
    "    following 300 100 architecture for MNIST dataset and using\n",
    "    provided parameter which are used to prune the model.\n",
    "    \n",
    "    Input: 'pruning_params' Python 3 dictionary containing parameters which are used for pruning\n",
    "    Output: Returns designed and compiled neural network model\n",
    "    \"\"\"\n",
    "    \n",
    "    pruned_model = Sequential()\n",
    "    pruned_model.add(l.InputLayer(input_shape=(784, )))\n",
    "    pruned_model.add(Flatten())\n",
    "    pruned_model.add(sparsity.prune_low_magnitude(\n",
    "        Dense(units = 300, activation='relu', kernel_initializer=tf.initializers.GlorotUniform()),\n",
    "        **pruning_params))\n",
    "    # pruned_model.add(l.Dropout(0.2))\n",
    "    pruned_model.add(sparsity.prune_low_magnitude(\n",
    "        Dense(units = 100, activation='relu', kernel_initializer=tf.initializers.GlorotUniform()),\n",
    "        **pruning_params))\n",
    "    # pruned_model.add(l.Dropout(0.1))\n",
    "    pruned_model.add(sparsity.prune_low_magnitude(\n",
    "        Dense(units = num_classes, activation='softmax'),\n",
    "        **pruning_params))\n",
    "    \n",
    "    # Compile pruned CNN-\n",
    "    pruned_model.compile(\n",
    "        loss=tf.keras.losses.categorical_crossentropy,\n",
    "        # optimizer='adam',\n",
    "        optimizer=tf.keras.optimizers.Adam(lr = 0.001),\n",
    "        metrics=['accuracy'])\n",
    "    \n",
    "    return pruned_model\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Add a pruning step callback to peg the pruning step to the optimizer's\n",
    "# step. Also add a callback to add pruning summaries to tensorboard\n",
    "callbacks = [\n",
    "             sparsity.UpdatePruningStep(),\n",
    "             # sparsity.PruningSummaries(log_dir = logdir, profile_batch=0),\n",
    "             tf.keras.callbacks.EarlyStopping(\n",
    "                 monitor='val_loss', patience = 3,\n",
    "                 min_delta=0.001\n",
    "             )\n",
    "]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "WARNING:tensorflow:From /opt/conda/lib/python3.7/site-packages/tensorflow_model_optimization/python/core/sparsity/keras/pruning_wrapper.py:183: Layer.add_variable (from tensorflow.python.keras.engine.base_layer) is deprecated and will be removed in a future version.\n",
      "Instructions for updating:\n",
      "Please use `layer.add_weight` method instead.\n"
     ]
    }
   ],
   "source": [
    "# Instantiate NN-\n",
    "orig_model = pruned_nn(pruning_params_unpruned)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Save randomly initialized weights-\n",
    "orig_model.save_weights(\"Random_Weights-Error_Recreation.h5\", overwrite = True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Train on 60000 samples, validate on 10000 samples\n",
      "Epoch 1/50\n",
      "60000/60000 [==============================] - 11s 184us/sample - loss: 0.2040 - accuracy: 0.9386 - val_loss: 0.1077 - val_accuracy: 0.9654\n",
      "Epoch 2/50\n",
      "60000/60000 [==============================] - 9s 155us/sample - loss: 0.0833 - accuracy: 0.9746 - val_loss: 0.0851 - val_accuracy: 0.9725\n",
      "Epoch 3/50\n",
      "60000/60000 [==============================] - 9s 156us/sample - loss: 0.0580 - accuracy: 0.9813 - val_loss: 0.0909 - val_accuracy: 0.9714\n",
      "Epoch 4/50\n",
      "60000/60000 [==============================] - 9s 152us/sample - loss: 0.0428 - accuracy: 0.9856 - val_loss: 0.0729 - val_accuracy: 0.9766\n",
      "Epoch 5/50\n",
      "60000/60000 [==============================] - 9s 154us/sample - loss: 0.0340 - accuracy: 0.9889 - val_loss: 0.0764 - val_accuracy: 0.9788\n",
      "Epoch 6/50\n",
      "60000/60000 [==============================] - 9s 153us/sample - loss: 0.0281 - accuracy: 0.9906 - val_loss: 0.0655 - val_accuracy: 0.9816\n",
      "Epoch 7/50\n",
      "60000/60000 [==============================] - 9s 154us/sample - loss: 0.0232 - accuracy: 0.9922 - val_loss: 0.0753 - val_accuracy: 0.9802\n",
      "Epoch 8/50\n",
      "60000/60000 [==============================] - 9s 152us/sample - loss: 0.0190 - accuracy: 0.9940 - val_loss: 0.0858 - val_accuracy: 0.9790\n",
      "Epoch 9/50\n",
      "60000/60000 [==============================] - 9s 153us/sample - loss: 0.0184 - accuracy: 0.9942 - val_loss: 0.0926 - val_accuracy: 0.9791\n"
     ]
    }
   ],
   "source": [
    "# Train model untile convergence-\n",
    "# Train unpruned Neural Network-\n",
    "history_orig = orig_model.fit(\n",
    "    x = X_train, y = y_train,\n",
    "    batch_size = batch_size,\n",
    "    epochs = epochs,\n",
    "    verbose = 1,\n",
    "    callbacks = callbacks,\n",
    "    validation_data = (X_test, y_test),\n",
    "    shuffle = True\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Strip model of pruning parameters-\n",
    "orig_model_stripped = sparsity.strip_pruning(orig_model)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Train parameters of neural network AFTER training UNPRUNED model-\n",
    "orig_model.save_weights(\"Trained_Weights-Error_Recreation.h5\", overwrite = True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "In original unpruned model, number of nonzero parameters in each layer are: \n",
      "\n",
      "235200\n",
      "300\n",
      "30000\n",
      "100\n",
      "1000\n",
      "10\n",
      "\n",
      "Total number of trainable parameters = 266610\n",
      "\n"
     ]
    }
   ],
   "source": [
    "# Count number of non-zero parameters in trained neural network-\n",
    "\n",
    "print(\"\\nIn original unpruned model, number of nonzero parameters in each layer are: \\n\")\n",
    "\n",
    "orig_sum_params = 0\n",
    "\n",
    "for layer in orig_model.trainable_weights:\n",
    "    print(tf.math.count_nonzero(layer, axis = None).numpy())\n",
    "    orig_sum_params += tf.math.count_nonzero(layer, axis = None).numpy()\n",
    "\n",
    "print(\"\\nTotal number of trainable parameters = {0}\\n\".format(orig_sum_params))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Specify the parameters to be used for layer-wise pruning, NO PRUNING is done here:\n",
    "pruning_params_constantsparsity = {\n",
    "    'pruning_schedule': sparsity.ConstantSparsity(\n",
    "        target_sparsity=0.2674, begin_step=100,\n",
    "        end_step = end_step, frequency=100\n",
    "    )\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Instantiate a Nueal Network model to be pruned using parameters from above-\n",
    "pruned_model = pruned_nn(pruning_params_constantsparsity)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load weights from original trained and unpruned model-\n",
    "pruned_model.load_weights(\"Trained_Weights-Error_Recreation.h5\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Train on 60000 samples, validate on 10000 samples\n",
      "Epoch 1/50\n",
      "60000/60000 [==============================] - 11s 179us/sample - loss: 0.0129 - accuracy: 0.9959 - val_loss: 0.0943 - val_accuracy: 0.9790\n",
      "Epoch 2/50\n",
      "60000/60000 [==============================] - 10s 162us/sample - loss: 0.0082 - accuracy: 0.9971 - val_loss: 0.0805 - val_accuracy: 0.9832\n",
      "Epoch 3/50\n",
      "60000/60000 [==============================] - 10s 162us/sample - loss: 0.0093 - accuracy: 0.9970 - val_loss: 0.1076 - val_accuracy: 0.9787\n",
      "Epoch 4/50\n",
      "60000/60000 [==============================] - 10s 160us/sample - loss: 0.0070 - accuracy: 0.9979 - val_loss: 0.1010 - val_accuracy: 0.9814\n",
      "Epoch 5/50\n",
      "60000/60000 [==============================] - 9s 158us/sample - loss: 0.0066 - accuracy: 0.9979 - val_loss: 0.0911 - val_accuracy: 0.9838\n"
     ]
    }
   ],
   "source": [
    "# Train pruned NN-\n",
    "history_pruned = pruned_model.fit(\n",
    "    x = X_train, y = y_train,\n",
    "    batch_size = batch_size,\n",
    "    epochs = epochs,\n",
    "    verbose = 1,\n",
    "    callbacks = callbacks,\n",
    "    validation_data = (X_test, y_test),\n",
    "    shuffle = True\n",
    ")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Strip the pruning wrappers from pruned model-\n",
    "pruned_model_stripped = sparsity.strip_pruning(pruned_model)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "In pruned model, number of nonzero parameters in each layer are: \n",
      "\n",
      "172308\n",
      "300\n",
      "21978\n",
      "100\n",
      "733\n",
      "10\n",
      "\n",
      "Total number of trainable parameters = 195429\n",
      "\n"
     ]
    }
   ],
   "source": [
    "print(\"\\nIn pruned model, number of nonzero parameters in each layer are: \\n\")\n",
    "\n",
    "pruned_sum_params = 0\n",
    "\n",
    "for layer in pruned_model.trainable_weights:\n",
    "# for layer in pruned_model_stripped.trainable_weights:\n",
    "    print(tf.math.count_nonzero(layer, axis = None).numpy())\n",
    "    pruned_sum_params += tf.math.count_nonzero(layer, axis = None).numpy()\n",
    "\n",
    "print(\"\\nTotal number of trainable parameters = {0}\\n\".format(pruned_sum_params))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "% of weights pruned = 26.70%\n",
      "\n"
     ]
    }
   ],
   "source": [
    "print(\"\\n% of weights pruned = {0:.2f}%\\n\".format(\n",
    "    ((orig_sum_params - pruned_sum_params) / orig_sum_params) * 100\n",
    "))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Save weights of PRUNED and Trained model BEFORE stripping-\n",
    "pruned_model.save_weights(\"Pruned_Weights-Error_Recreation.h5\", overwrite = True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Create mask for maintaining the sparsity of winning tickets:\n",
    "In order for pruned model to maintain it's sparsity, a mask is created which will be used by _GradientTape_ subsequently to train models.\n",
    "\n",
    "The mask is created as follows-\n",
    "1. Weights surviving the pruning are initialized to one (1)\n",
    "1. Weights which are pruned are initialized to zero (0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Instantiate a new neural network model for which, the mask is to be created,\n",
    "# according to the paper-\n",
    "mask_model = pruned_nn(pruning_params_unpruned)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load weights of PRUNED model-\n",
    "mask_model.load_weights(\"Pruned_Weights-Error_Recreation.h5\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Strip the model of its pruning parameters-\n",
    "mask_model_stripped = sparsity.strip_pruning(mask_model)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 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",
    "\n",
    "for wts in mask_model_stripped.trainable_weights:\n",
    "    wts.assign(tf.where(tf.equal(wts, 0.), 0., 1.))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Reset remaining parameters:\n",
    "Reset the remaining parameters in the _pruned_ model to their random weights when the model was initially created.\n",
    "\n",
    "In order to extract the winning ticket from the pruned neural network, __reset__ the weights of the surviving parts of the _pruned_ neural network to their _original randomly initialized and unpruned_ weights which were received before the training of the neural network model began (from above)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Instantiate a new neural network model for which, the weights are to be extracted, according to the paper-\n",
    "winning_ticket_model = pruned_nn(pruning_params_unpruned)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load weights of PRUNED model-\n",
    "winning_ticket_model.load_weights(\"Pruned_Weights-Error_Recreation.h5\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Strip the model of its pruning parameters-\n",
    "winning_ticket_model_stripped = sparsity.strip_pruning(winning_ticket_model)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "metadata": {},
   "outputs": [],
   "source": [
    "# For each layer, for each weight which is 0, leave it, as is. And for weights which survive the pruning,\n",
    "# reinitialize it to the value, the model received BEFORE it was trained and pruned-\n",
    "for orig_wts, pruned_wts in zip(orig_model_stripped.trainable_weights, winning_ticket_model_stripped.trainable_weights):\n",
    "    pruned_wts.assign(tf.where(tf.equal(pruned_wts, 0), pruned_wts, orig_wts))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 39,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Save the weights (with pruning parameters) extracted to a file-\n",
    "winning_ticket_model.save_weights(\"Winning_Ticket_Weights-Error_Recreation.h5\", overwrite=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 40,
   "metadata": {},
   "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": 41,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_dataset = train_dataset.shuffle(\n",
    "    buffer_size = 20000, reshuffle_each_iteration = True).batch(batch_size = batch_size, drop_remainder = False)\n",
    "\n",
    "test_dataset = test_dataset.batch(batch_size=batch_size, drop_remainder=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 42,
   "metadata": {},
   "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.001)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 43,
   "metadata": {},
   "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.BinaryAccuracy(name = 'train_accuracy')\n",
    "\n",
    "test_loss = tf.keras.metrics.Mean(name = 'test_loss')\n",
    "test_accuracy = tf.keras.metrics.BinaryAccuracy(name = 'train_accuracy')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 44,
   "metadata": {},
   "outputs": [],
   "source": [
    "@tf.function\n",
    "def train_one_step(model, mask_model, optimizer, x, y):\n",
    "    '''\n",
    "    def train_step(data, labels):\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"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 45,
   "metadata": {},
   "outputs": [],
   "source": [
    "@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"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Retrain Pruned model:\n",
    "The pruned model from above is retrained while _maintaining sparsity_ using mask from _mask_model_stripped_"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 46,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Instantiate a model\n",
    "model_gt = pruned_nn(pruning_params_unpruned)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 47,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load winning ticket (from above-)\n",
    "model_gt.load_weights(\"Winning_Ticket_Weights-Error_Recreation.h5\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 48,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Strip model of pruning parameters-\n",
    "model_gt_stripped = sparsity.strip_pruning(model_gt)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 49,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "number of trainable parameters in original model = 266610\n",
      "number of trainable parameters in pruned model = 195429\n",
      "\n"
     ]
    }
   ],
   "source": [
    "print(\"\\nnumber of trainable parameters in original model = {0}\".format(orig_sum_params))\n",
    "print(\"number of trainable parameters in pruned model = {0}\\n\".format(pruned_sum_params))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 50,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define variables for manual Early Stopping-\n",
    "best_val_loss = 1\n",
    "loc_patience = 0"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 51,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define variables for manual Early Stopping-\n",
    "patience = 3\n",
    "minimum_delta = 0.001"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 52,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 1, Loss: 0.0110, Accuracy: 99.9307, Test Loss: 0.0909, Test Accuracy: 99.612961\n",
      "Total number of trainable parameters = 195429\n",
      "\n",
      "Epoch 2, Loss: 0.0082, Accuracy: 99.9489, Test Loss: 0.0937, Test Accuracy: 99.623940\n",
      "Total number of trainable parameters = 195429\n",
      "\n",
      "Epoch 3, Loss: 0.0092, Accuracy: 99.9385, Test Loss: 0.0900, Test Accuracy: 99.623955\n",
      "Total number of trainable parameters = 195429\n",
      "\n",
      "Epoch 4, Loss: 0.0073, Accuracy: 99.9530, Test Loss: 0.1048, Test Accuracy: 99.608978\n",
      "Total number of trainable parameters = 195429\n",
      "\n",
      "Epoch 5, Loss: 0.0071, Accuracy: 99.9548, Test Loss: 0.0984, Test Accuracy: 99.631973\n",
      "Total number of trainable parameters = 195429\n",
      "\n",
      "Epoch 6, Loss: 0.0064, Accuracy: 99.9566, Test Loss: 0.1090, Test Accuracy: 99.591972\n",
      "Total number of trainable parameters = 195429\n",
      "\n",
      "\n",
      "'EarlyStopping' called!\n",
      "\n"
     ]
    }
   ],
   "source": [
    "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",
    "    for x, y in train_dataset:\n",
    "        # train_step(x, y)\n",
    "        train_one_step(model_gt_stripped, mask_model_stripped, optimizer, x, y)\n",
    "\n",
    "    for x_t, y_t in test_dataset:\n",
    "        # test_step(x_t, y_t)\n",
    "        test_step(model_gt_stripped, optimizer, x_t, y_t)\n",
    "\n",
    "    template = 'Epoch {0}, Loss: {1:.4f}, Accuracy: {2:.4f}, Test Loss: {3:.4f}, Test Accuracy: {4:4f}'\n",
    "    \n",
    "    print(template.format(epoch + 1, \n",
    "                              train_loss.result(), train_accuracy.result()*100,\n",
    "                              test_loss.result(), test_accuracy.result()*100))\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",
    "\n",
    "    model_sum_params = 0\n",
    "    \n",
    "    for layer in model_gt_stripped.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",
    "    # 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",
    "    \n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 53,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Save weights of winning ticket trained with GradientTape\n",
    "# WITH pruning parameter, so that it can be used to prune-\n",
    "model_gt.save_weights(\"Trained_Weights-Error_Recreation.h5\", overwrite=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Prune the trained model further:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 54,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Specify the parameters to be used for layer-wise pruning-\n",
    "pruning_params_constantsparsity = {\n",
    "    'pruning_schedule': sparsity.ConstantSparsity(\n",
    "        target_sparsity=0.4633, begin_step=100,\n",
    "        end_step = end_step, frequency=100\n",
    "    )\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 55,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Instantiate a Neural Network model to be pruned-\n",
    "pruned_model = pruned_nn(pruning_params_constantsparsity)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 56,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load weights of winning ticket from previous round to be PRUNED-\n",
    "pruned_model.load_weights(\"Trained_Weights-Error_Recreation.h5\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 57,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Pruning parameters of (GradientTape) trained model\n",
      "\n",
      "Train on 60000 samples, validate on 10000 samples\n",
      "Epoch 1/50\n",
      "60000/60000 [==============================] - 11s 179us/sample - loss: 0.0050 - accuracy: 0.9984 - val_loss: 0.0992 - val_accuracy: 0.9819\n",
      "Epoch 2/50\n",
      "60000/60000 [==============================] - 10s 161us/sample - loss: 0.0031 - accuracy: 0.9990 - val_loss: 0.1003 - val_accuracy: 0.9823\n",
      "Epoch 3/50\n",
      "60000/60000 [==============================] - 10s 160us/sample - loss: 0.0032 - accuracy: 0.9990 - val_loss: 0.1167 - val_accuracy: 0.9793\n",
      "Epoch 4/50\n",
      "60000/60000 [==============================] - 10s 162us/sample - loss: 0.0041 - accuracy: 0.9985 - val_loss: 0.1192 - val_accuracy: 0.9807\n"
     ]
    }
   ],
   "source": [
    "print(\"Pruning parameters of (GradientTape) trained model\\n\")\n",
    "\n",
    "# Train pruned Neural Network-\n",
    "history_pruned_gt = pruned_model.fit(\n",
    "    x = X_train, y = y_train,\n",
    "    batch_size = batch_size,\n",
    "    epochs = epochs,\n",
    "    verbose = 1,\n",
    "    callbacks = callbacks,\n",
    "    validation_data = (X_test, y_test),\n",
    "    shuffle = True\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 58,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Strip pruning wrapper off-\n",
    "pruned_model_stripped = sparsity.strip_pruning(pruned_model)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 59,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "Pruned model, # of non-zero trainable parameters = 143280\n"
     ]
    }
   ],
   "source": [
    "# Count number of non-zero parameters [Round - 2]-\n",
    "pruned_sum_params2 = 0\n",
    "    \n",
    "for layer in pruned_model_stripped.trainable_weights:\n",
    "    pruned_sum_params2 += tf.math.count_nonzero(layer, axis = None).numpy()\n",
    "\n",
    "print(\"\\nPruned model, # of non-zero trainable parameters = {0}\".format(pruned_sum_params2))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 60,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "% of weights pruned away = 46.26%\n",
      "\n"
     ]
    }
   ],
   "source": [
    "# Sanity-check: confirm that 46.33% of the weights are actually pruned away from the network-\n",
    "print(\"\\n% of weights pruned away = {0:.2f}%\\n\".format( \\\n",
    "    (orig_sum_params - pruned_sum_params2) / orig_sum_params * 100))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 61,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Save weights of PRUNED and Trained model BEFORE stripping-\n",
    "pruned_model.save_weights(\"Pruned_Weights-Error_Recreation.h5\", overwrite = True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Create mask to maintain sparsity of pruned model:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 62,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Instantiate a new neural network model for which, the mask is to be created,\n",
    "# according to the paper-\n",
    "mask_model = pruned_nn(pruning_params_unpruned)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 63,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load weights of PRUNED model-\n",
    "mask_model.load_weights(\"Pruned_Weights-Error_Recreation.h5\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 64,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Strip the model of its pruning parameters-\n",
    "mask_model_stripped = sparsity.strip_pruning(mask_model)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 65,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 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",
    "\n",
    "for wts in mask_model_stripped.trainable_weights:\n",
    "    wts.assign(tf.where(tf.equal(wts, 0.), 0., 1.))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Reset weights:\n",
    "Weights which are zero are left as it is, but for non-zero weights, reset them to random weights when the model was initialized"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 66,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Instantiate a new neural network model for which, the weights are to be extracted, according to the paper-\n",
    "winning_ticket_model = pruned_nn(pruning_params_unpruned)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 67,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load weights of PRUNED model-\n",
    "winning_ticket_model.load_weights(\"Pruned_Weights-Error_Recreation.h5\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 68,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Strip the model of its pruning parameters-\n",
    "winning_ticket_model_stripped = sparsity.strip_pruning(winning_ticket_model)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 69,
   "metadata": {},
   "outputs": [],
   "source": [
    "# For each layer, for each weight which is 0, leave it, as is. And for weights which survive the pruning,\n",
    "# reinitialize it to the value, the model received BEFORE it was trained and pruned-\n",
    "for orig_wts, pruned_wts in zip(orig_model_stripped.trainable_weights, winning_ticket_model_stripped.trainable_weights):\n",
    "    pruned_wts.assign(tf.where(tf.equal(pruned_wts, 0), pruned_wts, orig_wts))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 70,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Save the weights (with pruning parameters) extracted to a file-\n",
    "winning_ticket_model.save_weights(\"Winning_Ticket_Weights-Error_Recreation.h5\", overwrite=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Train pruned model (sparsity = 46.26%) using _GradientTape_"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 71,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Instantiate a model\n",
    "model_gt = pruned_nn(pruning_params_unpruned)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 72,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load winning ticket (from above-)\n",
    "model_gt.load_weights(\"Winning_Ticket_Weights-Error_Recreation.h5\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 73,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Strip model of pruning parameters-\n",
    "model_gt_stripped = sparsity.strip_pruning(model_gt)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 74,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "number of trainable parameters in original model = 266610\n",
      "number of trainable parameters [in Round - 2] in pruned model = 143280\n",
      "\n"
     ]
    }
   ],
   "source": [
    "print(\"\\nnumber of trainable parameters in original model = {0}\".format(orig_sum_params))\n",
    "print(\"number of trainable parameters [in Round - 2] in pruned model = {0}\\n\".format(pruned_sum_params2))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 75,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define variables for manual Early Stopping-\n",
    "best_val_loss = 1\n",
    "loc_patience = 0"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 76,
   "metadata": {},
   "outputs": [
    {
     "ename": "ValueError",
     "evalue": "in converted code:\n\n    <ipython-input-44-d0ca499a4063>:29 train_one_step  *\n        optimizer.apply_gradients(zip(grad_mask_mul, model.trainable_variables))\n    /home/majumdar/.local/lib/python3.7/site-packages/tensorflow_core/python/keras/optimizer_v2/optimizer_v2.py:435 apply_gradients\n        self._create_slots(var_list)\n    /home/majumdar/.local/lib/python3.7/site-packages/tensorflow_core/python/keras/optimizer_v2/adam.py:146 _create_slots\n        self.add_slot(var, 'm')\n    /home/majumdar/.local/lib/python3.7/site-packages/tensorflow_core/python/keras/optimizer_v2/optimizer_v2.py:587 add_slot\n        initial_value=initial_value)\n    /home/majumdar/.local/lib/python3.7/site-packages/tensorflow_core/python/ops/variables.py:260 __call__\n        return cls._variable_v2_call(*args, **kwargs)\n    /home/majumdar/.local/lib/python3.7/site-packages/tensorflow_core/python/ops/variables.py:254 _variable_v2_call\n        shape=shape)\n    /home/majumdar/.local/lib/python3.7/site-packages/tensorflow_core/python/ops/variables.py:65 getter\n        return captured_getter(captured_previous, **kwargs)\n    /home/majumdar/.local/lib/python3.7/site-packages/tensorflow_core/python/eager/def_function.py:413 invalid_creator_scope\n        \"tf.function-decorated function tried to create \"\n\n    ValueError: tf.function-decorated function tried to create variables on non-first call.\n",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mValueError\u001b[0m                                Traceback (most recent call last)",
      "\u001b[0;32m<ipython-input-76-9827a84843f1>\u001b[0m in \u001b[0;36m<module>\u001b[0;34m\u001b[0m\n\u001b[1;32m     13\u001b[0m     \u001b[0;32mfor\u001b[0m \u001b[0mx\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0my\u001b[0m \u001b[0;32min\u001b[0m \u001b[0mtrain_dataset\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     14\u001b[0m         \u001b[0;31m# train_step(x, y)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 15\u001b[0;31m         \u001b[0mtrain_one_step\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mmodel_gt_stripped\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mmask_model_stripped\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0moptimizer\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mx\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0my\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     16\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     17\u001b[0m     \u001b[0;32mfor\u001b[0m \u001b[0mx_t\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0my_t\u001b[0m \u001b[0;32min\u001b[0m \u001b[0mtest_dataset\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/.local/lib/python3.7/site-packages/tensorflow_core/python/eager/def_function.py\u001b[0m in \u001b[0;36m__call__\u001b[0;34m(self, *args, **kwds)\u001b[0m\n\u001b[1;32m    455\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    456\u001b[0m     \u001b[0mtracing_count\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_get_tracing_count\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 457\u001b[0;31m     \u001b[0mresult\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_call\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwds\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    458\u001b[0m     \u001b[0;32mif\u001b[0m \u001b[0mtracing_count\u001b[0m \u001b[0;34m==\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_get_tracing_count\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    459\u001b[0m       \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_call_counter\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mcalled_without_tracing\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/.local/lib/python3.7/site-packages/tensorflow_core/python/eager/def_function.py\u001b[0m in \u001b[0;36m_call\u001b[0;34m(self, *args, **kwds)\u001b[0m\n\u001b[1;32m    485\u001b[0m       \u001b[0;31m# In this case we have created variables on the first call, so we run the\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    486\u001b[0m       \u001b[0;31m# defunned version which is guaranteed to never create variables.\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 487\u001b[0;31m       \u001b[0;32mreturn\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_stateless_fn\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwds\u001b[0m\u001b[0;34m)\u001b[0m  \u001b[0;31m# pylint: disable=not-callable\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    488\u001b[0m     \u001b[0;32melif\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_stateful_fn\u001b[0m \u001b[0;32mis\u001b[0m \u001b[0;32mnot\u001b[0m \u001b[0;32mNone\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    489\u001b[0m       \u001b[0;31m# Release the lock early so that multiple threads can perform the call\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/.local/lib/python3.7/site-packages/tensorflow_core/python/eager/function.py\u001b[0m in \u001b[0;36m__call__\u001b[0;34m(self, *args, **kwargs)\u001b[0m\n\u001b[1;32m   1820\u001b[0m   \u001b[0;32mdef\u001b[0m \u001b[0m__call__\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1821\u001b[0m     \u001b[0;34m\"\"\"Calls a graph function specialized to the inputs.\"\"\"\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m-> 1822\u001b[0;31m     \u001b[0mgraph_function\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mkwargs\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_maybe_define_function\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m   1823\u001b[0m     \u001b[0;32mreturn\u001b[0m \u001b[0mgraph_function\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_filtered_call\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m  \u001b[0;31m# pylint: disable=protected-access\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1824\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/.local/lib/python3.7/site-packages/tensorflow_core/python/eager/function.py\u001b[0m in \u001b[0;36m_maybe_define_function\u001b[0;34m(self, args, kwargs)\u001b[0m\n\u001b[1;32m   2148\u001b[0m         \u001b[0mgraph_function\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_function_cache\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mprimary\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mget\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mcache_key\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;32mNone\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   2149\u001b[0m         \u001b[0;32mif\u001b[0m \u001b[0mgraph_function\u001b[0m \u001b[0;32mis\u001b[0m \u001b[0;32mNone\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m-> 2150\u001b[0;31m           \u001b[0mgraph_function\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_create_graph_function\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m   2151\u001b[0m           \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_function_cache\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mprimary\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0mcache_key\u001b[0m\u001b[0;34m]\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mgraph_function\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   2152\u001b[0m         \u001b[0;32mreturn\u001b[0m \u001b[0mgraph_function\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mkwargs\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/.local/lib/python3.7/site-packages/tensorflow_core/python/eager/function.py\u001b[0m in \u001b[0;36m_create_graph_function\u001b[0;34m(self, args, kwargs, override_flat_arg_shapes)\u001b[0m\n\u001b[1;32m   2039\u001b[0m             \u001b[0marg_names\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0marg_names\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   2040\u001b[0m             \u001b[0moverride_flat_arg_shapes\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0moverride_flat_arg_shapes\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m-> 2041\u001b[0;31m             capture_by_value=self._capture_by_value),\n\u001b[0m\u001b[1;32m   2042\u001b[0m         \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_function_attributes\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   2043\u001b[0m         \u001b[0;31m# Tell the ConcreteFunction to clean up its graph once it goes out of\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/.local/lib/python3.7/site-packages/tensorflow_core/python/framework/func_graph.py\u001b[0m in \u001b[0;36mfunc_graph_from_py_func\u001b[0;34m(name, python_func, args, kwargs, signature, func_graph, autograph, autograph_options, add_control_dependencies, arg_names, op_return_value, collections, capture_by_value, override_flat_arg_shapes)\u001b[0m\n\u001b[1;32m    913\u001b[0m                                           converted_func)\n\u001b[1;32m    914\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 915\u001b[0;31m       \u001b[0mfunc_outputs\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mpython_func\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0mfunc_args\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mfunc_kwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    916\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    917\u001b[0m       \u001b[0;31m# invariant: `func_outputs` contains only Tensors, CompositeTensors,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/.local/lib/python3.7/site-packages/tensorflow_core/python/eager/def_function.py\u001b[0m in \u001b[0;36mwrapped_fn\u001b[0;34m(*args, **kwds)\u001b[0m\n\u001b[1;32m    356\u001b[0m         \u001b[0;31m# __wrapped__ allows AutoGraph to swap in a converted function. We give\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    357\u001b[0m         \u001b[0;31m# the function a weak reference to itself to avoid a reference cycle.\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 358\u001b[0;31m         \u001b[0;32mreturn\u001b[0m \u001b[0mweak_wrapped_fn\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m__wrapped__\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwds\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    359\u001b[0m     \u001b[0mweak_wrapped_fn\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mweakref\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mref\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mwrapped_fn\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    360\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/.local/lib/python3.7/site-packages/tensorflow_core/python/framework/func_graph.py\u001b[0m in \u001b[0;36mwrapper\u001b[0;34m(*args, **kwargs)\u001b[0m\n\u001b[1;32m    903\u001b[0m           \u001b[0;32mexcept\u001b[0m \u001b[0mException\u001b[0m \u001b[0;32mas\u001b[0m \u001b[0me\u001b[0m\u001b[0;34m:\u001b[0m  \u001b[0;31m# pylint:disable=broad-except\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    904\u001b[0m             \u001b[0;32mif\u001b[0m \u001b[0mhasattr\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0me\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m\"ag_error_metadata\"\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 905\u001b[0;31m               \u001b[0;32mraise\u001b[0m \u001b[0me\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mag_error_metadata\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mto_exception\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0me\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    906\u001b[0m             \u001b[0;32melse\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    907\u001b[0m               \u001b[0;32mraise\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;31mValueError\u001b[0m: in converted code:\n\n    <ipython-input-44-d0ca499a4063>:29 train_one_step  *\n        optimizer.apply_gradients(zip(grad_mask_mul, model.trainable_variables))\n    /home/majumdar/.local/lib/python3.7/site-packages/tensorflow_core/python/keras/optimizer_v2/optimizer_v2.py:435 apply_gradients\n        self._create_slots(var_list)\n    /home/majumdar/.local/lib/python3.7/site-packages/tensorflow_core/python/keras/optimizer_v2/adam.py:146 _create_slots\n        self.add_slot(var, 'm')\n    /home/majumdar/.local/lib/python3.7/site-packages/tensorflow_core/python/keras/optimizer_v2/optimizer_v2.py:587 add_slot\n        initial_value=initial_value)\n    /home/majumdar/.local/lib/python3.7/site-packages/tensorflow_core/python/ops/variables.py:260 __call__\n        return cls._variable_v2_call(*args, **kwargs)\n    /home/majumdar/.local/lib/python3.7/site-packages/tensorflow_core/python/ops/variables.py:254 _variable_v2_call\n        shape=shape)\n    /home/majumdar/.local/lib/python3.7/site-packages/tensorflow_core/python/ops/variables.py:65 getter\n        return captured_getter(captured_previous, **kwargs)\n    /home/majumdar/.local/lib/python3.7/site-packages/tensorflow_core/python/eager/def_function.py:413 invalid_creator_scope\n        \"tf.function-decorated function tried to create \"\n\n    ValueError: tf.function-decorated function tried to create variables on non-first call.\n"
     ]
    }
   ],
   "source": [
    "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",
    "    for x, y in train_dataset:\n",
    "        # train_step(x, y)\n",
    "        train_one_step(model_gt_stripped, mask_model_stripped, optimizer, x, y)\n",
    "\n",
    "    for x_t, y_t in test_dataset:\n",
    "        # test_step(x_t, y_t)\n",
    "        test_step(model_gt_stripped, optimizer, x_t, y_t)\n",
    "\n",
    "    template = 'Epoch {0}, Loss: {1:.4f}, Accuracy: {2:.4f}, Test Loss: {3:.4f}, Test Accuracy: {4:4f}'\n",
    "    \n",
    "    print(template.format(epoch + 1, \n",
    "                              train_loss.result(), train_accuracy.result()*100,\n",
    "                              test_loss.result(), test_accuracy.result()*100))\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",
    "\n",
    "    model_sum_params = 0\n",
    "    \n",
    "    for layer in model_gt_stripped.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",
    "    # 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": []
  }
 ],
 "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
}