Check_Winning_Ticket.ipynb

{
 "cells": [
  {
   "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": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "l = tf.keras.layers"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "def nn_model():\n",
    "    \"\"\"\n",
    "    Function to create LeNet 300-100-10\n",
    "    model for MNIST classification\n",
    "    \"\"\"\n",
    "\n",
    "    model = Sequential()\n",
    "\n",
    "    model.add(l.InputLayer(input_shape=(784, )))\n",
    "\n",
    "    model.add(Flatten())\n",
    "\n",
    "    model.add(Dense(units = 300, activation='relu', kernel_initializer=tf.initializers.GlorotUniform()))\n",
    "\n",
    "    # model.add(l.Dropout(0.2))\n",
    "\n",
    "    model.add(Dense(units = 100, activation='relu', kernel_initializer=tf.initializers.GlorotUniform()))\n",
    "\n",
    "    # model.add(l.Dropout(0.1))\n",
    "\n",
    "    model.add(Dense(units = num_classes, activation='softmax'))\n",
    "    \n",
    "    return model\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "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": 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": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "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": [
    "# Initialize model-\n",
    "model = pruned_nn(pruning_params_unpruned)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load winning ticket weights-\n",
    "model.load_weights(\"Winning_Ticket_Weights_Experimental.h5\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Strip the pruning wrappers from pruned model-\n",
    "model_stripped = sparsity.strip_pruning(model)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "In Winning Ticket, number of nonzero parameters in each layer are: \n",
      "\n",
      "49627\n",
      "0\n",
      "6330\n",
      "0\n",
      "211\n",
      "0\n",
      "\n",
      "Total number of trainable parameters = 56168\n",
      "\n"
     ]
    }
   ],
   "source": [
    "print(\"\\nIn Winning Ticket, number of nonzero parameters in each layer are: \\n\")\n",
    "\n",
    "model_sum_params = 0\n",
    "\n",
    "for layer in model_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(\"\\nTotal number of trainable parameters = {0}\\n\".format(model_sum_params))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "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": 21,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load weights of GradientTape trained and PRUNED model-\n",
    "# mask_model.load_weights(\"Pruned_Weights.h5\")\n",
    "mask_model.load_weights(\"Winning_Ticket_Weights_Experimental.h5\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Strip the model of its pruning parameters-\n",
    "mask_model_stripped = sparsity.strip_pruning(mask_model)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "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": [
    "### Prepare dataset for _GradientTape_:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "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": 25,
   "metadata": {},
   "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": 26,
   "metadata": {},
   "outputs": [],
   "source": [
    "test_dataset = test_dataset.batch(batch_size=batch_size, drop_remainder=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "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": 28,
   "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.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": 29,
   "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": 30,
   "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": 31,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Dictionary to hold scalar metrics-\n",
    "history = {}\n",
    "\n",
    "history['accuracy'] = np.zeros(num_epochs)\n",
    "history['val_accuracy'] = np.zeros(num_epochs)\n",
    "history['loss'] = np.zeros(num_epochs)\n",
    "history['val_loss'] = np.zeros(num_epochs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "metadata": {},
   "outputs": [],
   "source": [
    "# User input-\n",
    "minimum_delta = 0.001\n",
    "patience = 3"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "metadata": {},
   "outputs": [],
   "source": [
    "best_val_loss = 1\n",
    "loc_patience = 0"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 1, Loss: 0.1759, Accuracy: 98.9948, Test Loss: 0.0746, Test Accuracy: 99.534950\n",
      "Total number of trainable parameters = 56168\n",
      "\n",
      "Epoch 2, Loss: 0.0457, Accuracy: 99.7279, Test Loss: 0.0566, Test Accuracy: 99.618980\n",
      "Total number of trainable parameters = 56168\n",
      "\n",
      "Epoch 3, Loss: 0.0246, Accuracy: 99.8562, Test Loss: 0.0575, Test Accuracy: 99.628975\n",
      "Total number of trainable parameters = 56168\n",
      "\n",
      "Epoch 4, Loss: 0.0148, Accuracy: 99.9216, Test Loss: 0.0589, Test Accuracy: 99.656952\n",
      "Total number of trainable parameters = 56168\n",
      "\n",
      "Epoch 5, Loss: 0.0088, Accuracy: 99.9549, Test Loss: 0.0689, Test Accuracy: 99.633965\n",
      "Total number of trainable parameters = 56168\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_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_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",
    "    history['accuracy'][epoch] = train_accuracy.result()\n",
    "    history['loss'][epoch] = train_loss.result()\n",
    "    history['val_loss'][epoch] = test_loss.result()\n",
    "    history['val_accuracy'][epoch] = test_accuracy.result()\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_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": 35,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Resize numpy arrays according to the epoch when 'EarlyStopping' was called-\n",
    "for metrics in history.keys():\n",
    "    history[metrics] = np.resize(history[metrics], new_shape=epoch)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[0.06900685519112594, 0.9815]"
      ]
     },
     "execution_count": 36,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model.evaluate(X_test, y_test, verbose=0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "metadata": {},
   "outputs": [],
   "source": [
    "y_pred = model.predict_classes(X_test)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 41,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "Accuracy of Winning Ticket (5-round) = 0.9815\n",
      "\n"
     ]
    }
   ],
   "source": [
    "accuracy = accuracy_score(np.argmax(y_test, axis = 1), y_pred)\n",
    "\n",
    "print(\"\\nAccuracy of Winning Ticket (5-round) = {0:.4f}\\n\".format(accuracy))"
   ]
  },
  {
   "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.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}