final_demo.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "source": [
        "# Training the Model\n",
        "We develop a space optimized CNN for image classification through synchronous distributed training, weight pruning, quantization, and data augmentation in Vertex AI on Google Cloud Platform. This is tested on the Cifar100 image dataset. In this notebook, you can try different hyperparameters for this model and attempt to beat 63% test accuracy on this dataset. Or, users can use this supermodel on a new dataset by uploading their own data under the \"Data\" section."
      ],
      "metadata": {
        "id": "nbU_ZUcnHQt3"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Installs"
      ],
      "metadata": {
        "id": "IGc3rZpINiUY"
      }
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "UmCFZ9uR6lyN",
        "outputId": "67d18fc5-6fe6-47fc-83d1-97d707f93652"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Collecting tensorflow-model-optimization\n",
            "  Downloading tensorflow_model_optimization-0.7.3-py2.py3-none-any.whl (238 kB)\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m238.9/238.9 kB\u001b[0m \u001b[31m7.0 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25hRequirement already satisfied: numpy~=1.14 in /opt/conda/lib/python3.7/site-packages (from tensorflow-model-optimization) (1.21.6)\n",
            "Requirement already satisfied: six~=1.10 in /opt/conda/lib/python3.7/site-packages (from tensorflow-model-optimization) (1.16.0)\n",
            "Requirement already satisfied: dm-tree~=0.1.1 in /opt/conda/lib/python3.7/site-packages (from tensorflow-model-optimization) (0.1.7)\n",
            "Installing collected packages: tensorflow-model-optimization\n",
            "Successfully installed tensorflow-model-optimization-0.7.3\n"
          ]
        }
      ],
      "source": [
        "!pip install tensorflow-model-optimization"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Cqsw0IP65J9B"
      },
      "outputs": [],
      "source": [
        "from __future__ import print_function\n",
        "import keras\n",
        "import tempfile\n",
        "from keras.layers import Dense, Conv2D, BatchNormalization, Activation\n",
        "from keras.layers import AveragePooling2D, Input, Flatten\n",
        "from keras.optimizers import Adam\n",
        "from keras.callbacks import ModelCheckpoint, LearningRateScheduler\n",
        "from keras.callbacks import ReduceLROnPlateau\n",
        "from keras.callbacks import CSVLogger  \n",
        "from keras.preprocessing.image import ImageDataGenerator\n",
        "from keras.regularizers import l2\n",
        "from keras import backend as K\n",
        "from keras.models import Model\n",
        "from keras.datasets import cifar10\n",
        "from tensorflow_model_optimization.python.core.sparsity.keras import pruning_callbacks\n",
        "import tensorflow_model_optimization as tfmot\n",
        "import numpy as np\n",
        "import time\n",
        "import os\n",
        "import pickle\n",
        "from keras.datasets import cifar100\n",
        "import tensorflow as tf\n",
        "import multiprocessing\n",
        "import os\n",
        "import random\n",
        "import zipfile\n",
        "%load_ext tensorboard"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Model Setup"
      ],
      "metadata": {
        "id": "_VTppvqyH7io"
      }
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "f3u6BwEZ5Ngo"
      },
      "outputs": [],
      "source": [
        "# timing per epoch iteration\n",
        "class TimeHistory(keras.callbacks.Callback):\n",
        "    def on_train_begin(self, logs={}):\n",
        "        self.times = []\n",
        "\n",
        "    def on_epoch_begin(self, batch, logs={}):\n",
        "        self.epoch_time_start = time.time()\n",
        "\n",
        "    def on_epoch_end(self, batch, logs={}):\n",
        "        self.times.append(time.time() - self.epoch_time_start)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "OJNDDDzn5Qyy",
        "outputId": "155abd49-1dab-4d05-a33d-9aa8be71cf9c"
      },
      "outputs": [
        {
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "2022-12-17 05:32:31.905911: I tensorflow/stream_executor/cuda/cuda_gpu_executor.cc:980] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero\n",
            "2022-12-17 05:32:32.841486: I tensorflow/stream_executor/cuda/cuda_gpu_executor.cc:980] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero\n",
            "2022-12-17 05:32:32.841970: I tensorflow/stream_executor/cuda/cuda_gpu_executor.cc:980] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero\n",
            "2022-12-17 05:32:33.024132: I tensorflow/core/platform/cpu_feature_guard.cc:193] This TensorFlow binary is optimized with oneAPI Deep Neural Network Library (oneDNN) to use the following CPU instructions in performance-critical operations:  AVX2 FMA\n",
            "To enable them in other operations, rebuild TensorFlow with the appropriate compiler flags.\n",
            "2022-12-17 05:32:33.049520: I tensorflow/stream_executor/cuda/cuda_gpu_executor.cc:980] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero\n",
            "2022-12-17 05:32:33.049908: I tensorflow/stream_executor/cuda/cuda_gpu_executor.cc:980] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero\n",
            "2022-12-17 05:32:33.050101: I tensorflow/stream_executor/cuda/cuda_gpu_executor.cc:980] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero\n",
            "2022-12-17 05:32:40.377643: I tensorflow/stream_executor/cuda/cuda_gpu_executor.cc:980] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero\n",
            "2022-12-17 05:32:40.399836: I tensorflow/stream_executor/cuda/cuda_gpu_executor.cc:980] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero\n",
            "2022-12-17 05:32:40.400280: I tensorflow/stream_executor/cuda/cuda_gpu_executor.cc:980] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero\n",
            "2022-12-17 05:32:40.400484: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1616] Created device /job:localhost/replica:0/task:0/device:GPU:0 with 14626 MB memory:  -> device: 0, name: Tesla V100-SXM2-16GB, pci bus id: 0000:00:04.0, compute capability: 7.0\n"
          ]
        }
      ],
      "source": [
        "# using data sharding, one divide strategy for single GPU distributed training. \n",
        "# the OneDeviceStrategy places variables and computation on a single specified device.\n",
        "strategy = tf.distribute.OneDeviceStrategy(device=\"/gpu:0\")  #   tf.distribute.experimental.CentralStorageStrategy()   #tf.distribute.MirroredStrategy()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "8xlQL90G5SzL"
      },
      "outputs": [],
      "source": [
        "# define builder function for resnet layers\n",
        "def resnet_layer(inputs,\n",
        "                 num_filters=16,\n",
        "                 kernel_size=3,\n",
        "                 strides=1,\n",
        "                 activation='relu',\n",
        "                 batch_normalization=True,\n",
        "                 conv_first=True):\n",
        "    \"\"\"2D Convolution-Batch Normalization-Activation stack builder\n",
        "\n",
        "    # Arguments\n",
        "        inputs (tensor): input tensor from input image or previous layer\n",
        "        num_filters (int): Conv2D number of filters\n",
        "        kernel_size (int): Conv2D square kernel dimensions\n",
        "        strides (int): Conv2D square stride dimensions\n",
        "        activation (string): activation name\n",
        "        batch_normalization (bool): whether to include batch normalization\n",
        "        conv_first (bool): conv-bn-activation (True) or\n",
        "            bn-activation-conv (False)\n",
        "\n",
        "    # Returns\n",
        "        x (tensor): tensor as input to the next layer\n",
        "    \"\"\"\n",
        "    conv = Conv2D(num_filters,\n",
        "                  kernel_size=kernel_size,\n",
        "                  strides=strides,\n",
        "                  padding='same',\n",
        "                  kernel_initializer='he_normal',\n",
        "                  kernel_regularizer=l2(1e-4))\n",
        "\n",
        "    x = inputs\n",
        "    if conv_first:\n",
        "        x = conv(x)\n",
        "        if batch_normalization:\n",
        "            x = BatchNormalization()(x)\n",
        "        if activation is not None:\n",
        "            x = Activation(activation)(x)\n",
        "    else:\n",
        "        if batch_normalization:\n",
        "            x = BatchNormalization()(x)\n",
        "        if activation is not None:\n",
        "            x = Activation(activation)(x)\n",
        "        x = conv(x)\n",
        "    return x"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "q0qeB2wJ5nsX"
      },
      "outputs": [],
      "source": [
        "# actual resnet code\n",
        "def resnet_v1(input_shape, depth, num_classes=10):\n",
        "    \"\"\"ResNet Version 1 Model builder [a]\n",
        "\n",
        "    Stacks of 2 x (3 x 3) Conv2D-BN-ReLU\n",
        "    Last ReLU is after the shortcut connection.\n",
        "    At the beginning of each stage, the feature map size is halved (downsampled)\n",
        "    by a convolutional layer with strides=2, while the number of filters is\n",
        "    doubled. Within each stage, the layers have the same number filters and the\n",
        "    same number of filters.\n",
        "    Features maps sizes:\n",
        "    stage 0: 32x32, 16\n",
        "    stage 1: 16x16, 32\n",
        "    stage 2:  8x8,  64\n",
        "    The Number of parameters is approx the same as Table 6 of [a]:\n",
        "    ResNet20 0.27M\n",
        "    ResNet32 0.46M\n",
        "    ResNet44 0.66M\n",
        "    ResNet56 0.85M\n",
        "    ResNet110 1.7M\n",
        "\n",
        "    # Arguments\n",
        "        input_shape (tensor): shape of input image tensor\n",
        "        depth (int): number of core convolutional layers\n",
        "        num_classes (int): number of classes (CIFAR10 has 10)\n",
        "\n",
        "    # Returns\n",
        "        model (Model): Keras model instance\n",
        "    \"\"\"\n",
        "    if (depth - 2) % 6 != 0:\n",
        "        raise ValueError('depth should be 6n+2 (eg 20, 32, 44 in [a])')\n",
        "    # Start model definition.\n",
        "    num_filters = 16\n",
        "    num_res_blocks = int((depth - 2) / 6)\n",
        "\n",
        "    inputs = Input(shape=input_shape)\n",
        "    x = resnet_layer(inputs=inputs)\n",
        "    # Instantiate the stack of residual units\n",
        "    for stack in range(3):\n",
        "        for res_block in range(num_res_blocks):\n",
        "            strides = 1\n",
        "            if stack > 0 and res_block == 0:  # first layer but not first stack\n",
        "                strides = 2  # downsample\n",
        "            y = resnet_layer(inputs=x,\n",
        "                             num_filters=num_filters,\n",
        "                             strides=strides)\n",
        "            y = resnet_layer(inputs=y,\n",
        "                             num_filters=num_filters,\n",
        "                             activation=None)\n",
        "            if stack > 0 and res_block == 0:  # first layer but not first stack\n",
        "                # linear projection residual shortcut connection to match\n",
        "                # changed dims\n",
        "                x = resnet_layer(inputs=x,\n",
        "                                 num_filters=num_filters,\n",
        "                                 kernel_size=1,\n",
        "                                 strides=strides,\n",
        "                                 activation=None,\n",
        "                                 batch_normalization=False)\n",
        "            x = keras.layers.add([x, y])\n",
        "            x = Activation('relu')(x)\n",
        "        num_filters *= 2\n",
        "\n",
        "    # Add classifier on top.\n",
        "    # v1 does not use BN after last shortcut connection-ReLU\n",
        "    x = AveragePooling2D(pool_size=8)(x)\n",
        "    y = Flatten()(x)\n",
        "    outputs = Dense(num_classes,\n",
        "                    activation='softmax',\n",
        "                    kernel_initializer='he_normal')(y)\n",
        "    # Instantiate model.\n",
        "    model = Model(inputs=inputs, outputs=outputs)\n",
        "    return model"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "KI--slgF5-D0"
      },
      "outputs": [],
      "source": [
        "# Learning rate scheduler \n",
        "def lr_schedule(epoch):\n",
        "    \"\"\"Learning Rate Schedule\n",
        "\n",
        "    Learning rate is scheduled to be reduced after 80, 120, 160, 180 epochs.\n",
        "    Called automatically every epoch as part of callbacks during training.\n",
        "\n",
        "    # Arguments\n",
        "        epoch (int): The number of epochs\n",
        "\n",
        "    # Returns\n",
        "        lr (float32): learning rate\n",
        "    \"\"\"\n",
        "    lr = 1e-3\n",
        "    if epoch > 180:\n",
        "        lr *= 0.5e-3\n",
        "    elif epoch > 160:\n",
        "        lr *= 1e-3\n",
        "    elif epoch > 120:\n",
        "        lr *= 1e-2\n",
        "    elif epoch > 80:\n",
        "        lr *= 1e-1\n",
        "    print('Learning rate: ', lr)\n",
        "    return lr"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "jtXh5ajC6AKC"
      },
      "outputs": [],
      "source": [
        "# tflite conversion\n",
        "def tflite_conv(model,path,quant=False):\n",
        "    converter = tf.lite.TFLiteConverter.from_keras_model(model)\n",
        "    if quant == True:\n",
        "        converter.optimizations = [tf.lite.Optimize.DEFAULT]\n",
        "    pruned_tflite_model = converter.convert()\n",
        "    with open(path, 'wb') as f:\n",
        "      f.write(pruned_tflite_model)\n",
        "    print('Saved pruned TFLite model to:',path)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "GQ4jU6PM6CJV"
      },
      "outputs": [],
      "source": [
        "# returns size of gzipped model, in bytes.\n",
        "def get_gzipped_model_size(file):\n",
        "    _, zipped_file = tempfile.mkstemp('.zip')\n",
        "    with zipfile.ZipFile(zipped_file, 'w', compression=zipfile.ZIP_DEFLATED) as f:\n",
        "        f.write(file)\n",
        "\n",
        "    return os.path.getsize(zipped_file)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "TqbzgNzq6EQB",
        "outputId": "b09f407c-f067-4af3-cf3c-e79a3efa8d50"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "1"
            ]
          },
          "execution_count": 12,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "# number of model replicas\n",
        "strategy.num_replicas_in_sync"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "J7YmlsR26M8G"
      },
      "outputs": [],
      "source": [
        "# actual trainer function\n",
        "def resnet_training(x_train,y_train,x_test, y_test,layers,frequency,\\\n",
        "                    initial_sparsity,final_sparsity,gpu=\"T4\",\\\n",
        "                    const=True, poly=False,file_name='cifar10_test1',\\\n",
        "                    num_classes=10,begin_step=0,end_step='default'):\n",
        "\n",
        "    # Default parameters\n",
        "    batch_size = 128 * strategy.num_replicas_in_sync \n",
        "    epochs = 300\n",
        "    n = layers\n",
        "\n",
        "    # Computed depth from supplied model parameter n\n",
        "    depth = n * 6 + 2\n",
        "\n",
        "    # Input image dimensions.\n",
        "    input_shape = x_train.shape[1:]\n",
        "\n",
        "    # Subtracting pixel mean improves accuracy\n",
        "    x_train_mean = np.mean(x_train, axis=0)\n",
        "    x_train -= x_train_mean\n",
        "    x_test -= x_train_mean\n",
        "\n",
        "    print('x_train shape:', x_train.shape)\n",
        "    print(x_train.shape[0], 'train samples')\n",
        "    print(x_test.shape[0], 'test samples')\n",
        "    print('y_train shape:', y_train.shape)\n",
        "\n",
        "\n",
        "    num_images = x_train.shape[0] \n",
        "\n",
        "    if end_step == 'default':\n",
        "        end_step = np.ceil(num_images / batch_size).astype(np.int32) * epochs\n",
        "\n",
        "    #Cifar100 supermodel hyperparameters: polnomialDecay, initial_sparsity=0.50, final_sparsity=0.6, frequency=100, begin and end step defaults\n",
        "    if poly:\n",
        "        pruning_params = {\n",
        "              'pruning_schedule': tfmot.sparsity.keras.PolynomialDecay(initial_sparsity=initial_sparsity,\n",
        "                                                                       final_sparsity=final_sparsity,\n",
        "                                                                       begin_step=begin_step,\n",
        "                                                                       end_step=end_step,\n",
        "                                                                      frequency=frequency)\n",
        "        }\n",
        "    if const:\n",
        "            pruning_params = {\n",
        "              'pruning_schedule': tfmot.sparsity.keras.ConstantSparsity(target_sparsity=final_sparsity,\n",
        "                                                                       begin_step=begin_step,\n",
        "                                                                       end_step=end_step,\n",
        "                                                                       frequency=frequency)\n",
        "        }\n",
        "\n",
        "    prune_low_magnitude = tfmot.sparsity.keras.prune_low_magnitude\n",
        "    with strategy.scope():\n",
        "        model = resnet_v1(input_shape=input_shape, depth=depth,num_classes=num_classes)\n",
        "        model = prune_low_magnitude(model, **pruning_params)    \n",
        "\n",
        "        model.compile(loss= 'categorical_crossentropy',    \n",
        "                  optimizer=Adam(lr=lr_schedule(0)),\n",
        "                  metrics=['accuracy'])\n",
        "\n",
        "    model.summary()\n",
        "\n",
        "    # Prepare model model saving directory.\n",
        "    save_dir = os.path.join(os.getcwd(), 'saved_models')\n",
        "    if not os.path.isdir(save_dir):\n",
        "        os.makedirs(save_dir)\n",
        "\n",
        "    model_run=file_name+'_'+str(depth)+'_'+gpu\n",
        "    model_path=model_run+'.h5'\n",
        "    print(\"RRrrrrrR\",model_path)\n",
        "    \n",
        "    filepath = os.path.join(save_dir, model_path)\n",
        "    print(\"RRR\",filepath)\n",
        "    \n",
        "    print(model_run)\n",
        "\n",
        "    # Prepare callbacks for model saving and for learning rate adjustment.\n",
        "    checkpoint = ModelCheckpoint(filepath=filepath,\n",
        "                                 monitor='val_acc',\n",
        "                                 verbose=1,\n",
        "                                 save_best_only=True)\n",
        "\n",
        "    lr_scheduler = LearningRateScheduler(lr_schedule)\n",
        "\n",
        "    lr_reducer = ReduceLROnPlateau(factor=np.sqrt(0.1),\n",
        "                                   cooldown=0,\n",
        "                                   patience=5,\n",
        "                                   min_lr=0.5e-6)\n",
        "    logdir = tempfile.mkdtemp()\n",
        "    time_callback = TimeHistory()\n",
        "    logname='/home/jupyter/final_proj/log_'+model_run+'.csv'\n",
        "    csv_logger = CSVLogger(logname, append=True, separator=';')\n",
        "    callbacks = [checkpoint, lr_reducer, lr_scheduler, csv_logger, time_callback, pruning_callbacks.UpdatePruningStep(), tfmot.sparsity.keras.PruningSummaries(log_dir=logdir)]\n",
        "\n",
        "    st = time.time()\n",
        "    ######################## no augmentation################\n",
        "    # model.fit(x_train, y_train,\n",
        "    #           batch_size=batch_size,\n",
        "    #           epochs=epochs,\n",
        "    #           validation_data=(x_test, y_test),\n",
        "    #           shuffle=True,\n",
        "    #           callbacks=callbacks)\n",
        "    ######################## no augmentation################\n",
        "\n",
        "    ########################augmentation################\n",
        "    datagen = ImageDataGenerator(\n",
        "          # set input mean to 0 over the dataset\n",
        "          featurewise_center=False,\n",
        "          # set each sample mean to 0\n",
        "          samplewise_center=False,\n",
        "          # divide inputs by std of dataset\n",
        "          featurewise_std_normalization=False,\n",
        "          # divide each input by its std\n",
        "          samplewise_std_normalization=False,\n",
        "          # apply ZCA whitening\n",
        "          zca_whitening=False,\n",
        "          # epsilon for ZCA whitening\n",
        "          zca_epsilon=1e-06,\n",
        "          # randomly rotate images in the range (deg 0 to 180)\n",
        "          rotation_range=0,\n",
        "          # randomly shift images horizontally\n",
        "          width_shift_range=0.1,\n",
        "          # randomly shift images vertically\n",
        "          height_shift_range=0.1,\n",
        "          # set range for random shear\n",
        "          shear_range=0.,\n",
        "          # set range for random zoom\n",
        "          zoom_range=0.,\n",
        "          # set range for random channel shifts\n",
        "          channel_shift_range=0.,\n",
        "          # set mode for filling points outside the input boundaries\n",
        "          fill_mode='nearest',\n",
        "          # value used for fill_mode = \"constant\"\n",
        "          cval=0.,\n",
        "          # randomly flip images\n",
        "          horizontal_flip=True,\n",
        "          # randomly flip images\n",
        "          vertical_flip=False,\n",
        "          # set rescaling factor (applied before any other transformation)\n",
        "          rescale=None,\n",
        "          # set function that will be applied on each input\n",
        "          preprocessing_function=None,\n",
        "          # image data format, either \"channels_first\" or \"channels_last\"\n",
        "          data_format=None,\n",
        "          # fraction of images reserved for validation (strictly between 0 and 1)\n",
        "          validation_split=0.0)\n",
        "          \n",
        "    datagen.fit(x_train)\n",
        "\n",
        "    # Fit the model on the batches generated by datagen.flow().\n",
        "    model.fit_generator(datagen.flow(x_train, y_train, batch_size=batch_size),\n",
        "                        validation_data=(x_test, y_test),\n",
        "                        epochs=epochs, verbose=1, \n",
        "                        callbacks=callbacks)\n",
        "     \n",
        "     ########################augmentation################\n",
        "\n",
        "    %tensorboard --logdir={logdir} #visualize input through tensorboard\n",
        "    training_time = time.time() - st\n",
        "\n",
        "    print(f\"overall training time is {training_time}\")\n",
        "    epoch_times = time_callback.times\n",
        "    print(f\"each epoch training time is {epoch_times}\")\n",
        "\n",
        "    # Score trained model.\n",
        "    scores = model.evaluate(x_test, y_test, verbose=1)\n",
        "    print('Test loss:', scores[0])\n",
        "    print('Test accuracy:', scores[1])\n",
        "    print(\"tt\",save_dir+'/'+model_run)\n",
        "    \n",
        "    #save standard model \n",
        "    model.save(save_dir+'/'+model_path)\n",
        "    \n",
        "    #saving data \n",
        "    with open(model_run+'.pickle', 'wb') as handle:\n",
        "        pickle.dump([training_time,epoch_times,scores[0],scores[1]], handle)\n",
        "\n",
        "    #save pruned model\n",
        "    model_for_export = tfmot.sparsity.keras.strip_pruning(model)\n",
        "    pruned_keras_file = save_dir+'/PRUNE_'+model_path\n",
        "    keras.models.save_model(model_for_export, pruned_keras_file, include_optimizer=False)\n",
        "    print('Saved pruned Keras model to:', pruned_keras_file)\n",
        "    \n",
        "    #convert to tflite+ save\n",
        "    pruned_tflite_file=save_dir+'/lite_PRUNE_'+model_path\n",
        "    tflite_conv(model,pruned_tflite_file)\n",
        "    \n",
        "    #save to tf lite + qaunt\n",
        "    pruned_tflite_quant_file=save_dir+'/lite_quant_PRUNE_'+model_path\n",
        "    tflite_conv(model,pruned_tflite_quant_file,quant=True)\n",
        "    \n",
        "    pruned=get_gzipped_model_size(pruned_keras_file)\n",
        "    tflite_pruned=get_gzipped_model_size(pruned_tflite_file)\n",
        "    tflite_quant_pruned=get_gzipped_model_size(pruned_tflite_quant_file)\n",
        "    \n",
        "    \n",
        "    print(\"Size of gzipped pruned Keras model: %.2f bytes\" % (pruned))\n",
        "    print(\"Size of gzipped pruned TFlite model: %.2f bytes\" % (tflite_pruned))\n",
        "    print(\"Size of gzipped pruned and quantized TFlite model: %.2f bytes\" % (tflite_quant_pruned))\n",
        "    \n",
        "    with open(model_run+'.pickle', 'wb') as handle: #save model information for analysis\n",
        "        pickle.dump([training_time,epoch_times,scores[0],scores[1],pruned,tflite_pruned,tflite_quant_pruned], handle)\n",
        "\n",
        "    return(model,model_path)"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Data: loading and preparation"
      ],
      "metadata": {
        "id": "8KizmH9JNndK"
      }
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "XRxSJSxR6QKd",
        "outputId": "7041ab59-ee67-4782-df51-29c5f0f61897"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Downloading data from https://www.cs.toronto.edu/~kriz/cifar-100-python.tar.gz\n",
            "169001437/169001437 [==============================] - 2s 0us/step\n"
          ]
        }
      ],
      "source": [
        "# Load the CIFAR100 data.\n",
        "(X_train, Y_train), (X_test, Y_test) = cifar100.load_data()\n",
        "# Normalize data.\n",
        "X_train = X_train.astype('float32') / 255\n",
        "X_test = X_test.astype('float32 ') / 255\n",
        "# Convert class vectors to binary class matrices.\n",
        "Y_train = keras.utils.to_categorical(Y_train)\n",
        "Y_test = keras.utils.to_categorical(Y_test)"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Executable Code \n",
        "### Exploring the supermodel: <br>\n",
        "Try exploring hyperparamters:<br>\n",
        "-Frequency of pruning<br>\n",
        "-Initial/final sparsity percentage of weights <br>\n",
        "-Constant or polynomial decaying pruning schedule (const = True, poly = True)<br>\n",
        "-Layers: *note this is not the actual number of layers <br>\n",
        "\n"
      ],
      "metadata": {
        "id": "-ty7RAaZH_kN"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# Model parameter\n",
        "# ------------------\n",
        "#           |        | \n",
        "# Model     |\"layers\"| \n",
        "# -----------------\n",
        "# ResNet20  | 3      |\n",
        "# ResNet32  | 5      | \n",
        "# ResNet44  | 7      | \n",
        "# ResNet56  | 9      | \n",
        "# ResNet110 |18      | \n",
        "# ResNet164 |27      | \n",
        "# ResNet1001|111     | \n",
        "# -------------------"
      ],
      "metadata": {
        "id": "CozL4cnyQIr7"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Run the code below to execute the top performing super model:"
      ],
      "metadata": {
        "id": "uIWcj23wIOdA"
      }
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "mVVvxo1P6T2o"
      },
      "outputs": [],
      "source": [
        "layers = [3, 7]\n",
        "names = [\"twenty\", \"fortyfour\"]\n",
        "for x in range(len(layers)):\n",
        "    name=\"resnet_super_\"+names[x]\n",
        "    model1 = resnet_training(X_train,Y_train,X_test, Y_test,layers=layers[x],frequency=100,\\\n",
        "                    initial_sparsity = 0.5, final_sparsity=0.6,gpu=\"V100\",\\\n",
        "                    const=False, poly=True,file_name=name,\\\n",
        "                    num_classes=100,begin_step=0,end_step='default')"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "yC-PpggqHfTm"
      },
      "outputs": [],
      "source": []
    }
  ],
  "metadata": {
    "colab": {
      "provenance": []
    },
    "kernelspec": {
      "display_name": "TensorFlow 2 (Local)",
      "language": "python",
      "name": "local-tf2"
    },
    "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.12"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}