Model_optimization.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/Zain-Pirzada/Diabetic-Reinopathy-Kaggle-Competition/blob/main/Model_optimization.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 5,
      "id": "ab3f7119",
      "metadata": {
        "id": "ab3f7119"
      },
      "outputs": [],
      "source": [
        "import tensorflow as tf"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 7,
      "id": "5ea0ad5f",
      "metadata": {
        "id": "5ea0ad5f",
        "pycharm": {
          "is_executing": true
        }
      },
      "outputs": [],
      "source": [
        "from keras.applications import InceptionResNetV2\n",
        "from keras.layers.core import Lambda\n",
        "from keras.applications.mobilenet_v2 import MobileNetV2\n",
        "from keras.applications.inception_v3 import InceptionV3 \n",
        "from keras.applications.densenet import DenseNet121\n",
        "from keras.applications.vgg16 import VGG16\n",
        "from tensorflow.keras.applications.resnet50 import ResNet50\n",
        "from keras.applications.xception import Xception\n",
        "from keras.applications.mobilenet import MobileNet\n",
        "from keras.layers import *\n",
        "from keras.models import Model\n",
        "from keras.callbacks import ModelCheckpoint, ReduceLROnPlateau\n",
        "from keras.preprocessing.image import ImageDataGenerator\n",
        "from keras.optimizers import *\n",
        "# from keras.utils.multi_gpu_utils import multi_gpu_model\n",
        "from matplotlib import pyplot as plt\n",
        "from keras.models import load_model\n",
        "import keras.backend as K\n",
        "import os\n",
        "import tensorflow as tf\n",
        "\n",
        "\n",
        "import numpy as np\n",
        "import pandas as pd\n",
        "from PIL import Image\n",
        "import cv2 \n",
        "\n",
        "import sys\n",
        "import os.path"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 8,
      "id": "3660b2c7",
      "metadata": {
        "id": "3660b2c7"
      },
      "outputs": [],
      "source": [
        "\n",
        "def Global_attention_block(inputs):\n",
        "    shape=K.int_shape(inputs)\n",
        "    x=AveragePooling2D(pool_size=(shape[1],shape[2])) (inputs)\n",
        "    x=Conv2D(shape[3],1, padding='same') (x)\n",
        "    x=Activation('relu') (x)\n",
        "    x=Conv2D(shape[3],1, padding='same') (x)\n",
        "    x=Activation('sigmoid') (x)\n",
        "    C_A=Multiply()([x,inputs])\n",
        "    \n",
        "    x=Lambda(lambda x: K.mean(x,axis=-1,keepdims=True))  (C_A)\n",
        "    x=Activation('sigmoid') (x)\n",
        "    S_A=Multiply()([x,C_A])\n",
        "    return S_A"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 9,
      "id": "2f08d5ab",
      "metadata": {
        "id": "2f08d5ab"
      },
      "outputs": [],
      "source": [
        "def Category_attention_block(inputs,classes,k):\n",
        "    shape=K.int_shape(inputs)\n",
        "    F=Conv2D(k*classes,1, padding='same') (inputs)\n",
        "    F=BatchNormalization() (F)\n",
        "    F1=Activation('relu') (F)\n",
        "    \n",
        "    F2=F1\n",
        "    x=GlobalMaxPool2D()(F2)\n",
        "    \n",
        "    x=Reshape((classes,k)) (x)\n",
        "    S=Lambda(lambda x: K.mean(x,axis=-1,keepdims=False))  (x)\n",
        "    \n",
        "    x=Reshape((shape[1],shape[2],classes,k)) (F1)\n",
        "    x=Lambda(lambda x: K.mean(x,axis=-1,keepdims=False))  (x)\n",
        "    x=Multiply()([S,x])\n",
        "    M=Lambda(lambda x: K.mean(x,axis=-1,keepdims=True))  (x)\n",
        "    \n",
        "    semantic=Multiply()([inputs,M])\n",
        "    return semantic\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 10,
      "id": "700cfcb9",
      "metadata": {
        "id": "700cfcb9"
      },
      "outputs": [],
      "source": [
        "def smooth_curve(points, factor=0.9):\n",
        "    smoothed_points = []\n",
        "    for point in points:\n",
        "        if smoothed_points:\n",
        "            previous = smoothed_points[-1]\n",
        "            smoothed_points.append(previous * factor + point * (1 - factor))\n",
        "        else:\n",
        "            smoothed_points.append(point)\n",
        "    return smoothed_points    \n",
        "   "
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 11,
      "id": "b0b0dd5d",
      "metadata": {
        "id": "b0b0dd5d"
      },
      "outputs": [],
      "source": [
        "def plotmodel(history,name):\n",
        "    \n",
        "    acc = history.history['acc']\n",
        "    val_acc = history.history['val_acc']\n",
        "    loss = history.history['loss']\n",
        "    val_loss = history.history['val_loss']\n",
        "    epochs = range(1, len(acc) + 1) \n",
        "    \n",
        "    plt.figure(1)                  \n",
        "    plt.plot(epochs,smooth_curve(acc))\n",
        "    plt.plot(epochs,smooth_curve(val_acc))\n",
        "    plt.ylabel('acc')\n",
        "    plt.xlabel('epoch')\n",
        "    plt.legend(['train_acc', 'val_acc'], loc='upper left')\n",
        "    plt.savefig('acc_'+name+'.png')\n",
        "    \n",
        "    plt.figure(2)\n",
        "    plt.plot(epochs,smooth_curve(loss))\n",
        "    plt.plot(epochs,smooth_curve(val_loss))\n",
        "    plt.ylabel('loss')\n",
        "    plt.xlabel('epoch')\n",
        "    plt.legend(['train_loss', 'val_loss'], loc='upper right')\n",
        "    plt.savefig('loss_'+name+'.png')"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 12,
      "id": "6539459d",
      "metadata": {
        "id": "6539459d"
      },
      "outputs": [],
      "source": [
        "def label_smooth(y_true, y_pred):\n",
        "    y_true=((1-0.1)*y_true+0.05)\n",
        "    return K.categorical_crossentropy(y_true, y_pred) \n",
        " "
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 63,
      "id": "6910ffda",
      "metadata": {
        "id": "6910ffda"
      },
      "outputs": [],
      "source": [
        "def get_base_model(model_name,image_size):\n",
        "    if model_name =='vgg16':\n",
        "        base_model=VGG16              (include_top=False,weights='imagenet',input_shape=(image_size,image_size,3))\n",
        "    if model_name =='resnet50':\n",
        "        base_model=ResNet50           (include_top=False,weights='imagenet',input_shape=(image_size,image_size,3))\n",
        "    if model_name =='xception':\n",
        "        base_model=Xception           (include_top=False, weights='imagenet',input_shape=(image_size,image_size,3))\n",
        "    if model_name =='densenet121':\n",
        "        base_model=DenseNet121       (include_top=False, weights='imagenet',input_shape=(image_size,image_size,3))\n",
        "    if model_name =='mobilenet0.75':\n",
        "        base_model=MobileNet         (include_top=False,weights='imagenet',alpha=0.75,input_shape=(image_size,image_size,3))\n",
        "    if model_name =='mobilenet1.0':\n",
        "        base_model=MobileNet         (include_top=False,weights='imagenet',alpha=1.0,input_shape=(image_size,image_size,3))\n",
        "    if model_name =='mobilenetv2':\n",
        "        base_model=MobileNetV2      (include_top=False,weights='imagenet',alpha=1.0,input_shape=(image_size,image_size,3))\n",
        "    if model_name =='inceptionv3':   \n",
        "        base_model=InceptionV3       (include_top=False,weights='imagenet',input_shape=(image_size,image_size,3))\n",
        "    if model_name =='inceptionv2':\n",
        "        base_model=InceptionResNetV2 (include_top=False, weights='imagenet',input_shape=(image_size,image_size,3))\n",
        "      \n",
        "    return base_model\n",
        "\n",
        "    "
      ]
    },
    {
      "cell_type": "code",
      "source": [],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 165
        },
        "id": "Ljy6-0VLL8pr",
        "outputId": "642ddbb4-b08f-4e94-f39d-09e590f7d94a"
      },
      "id": "Ljy6-0VLL8pr",
      "execution_count": 71,
      "outputs": [
        {
          "output_type": "error",
          "ename": "NameError",
          "evalue": "ignored",
          "traceback": [
            "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
            "\u001b[0;31mNameError\u001b[0m                                 Traceback (most recent call last)",
            "\u001b[0;32m<ipython-input-71-26be24aabd97>\u001b[0m in \u001b[0;36m<module>\u001b[0;34m\u001b[0m\n\u001b[0;32m----> 1\u001b[0;31m \u001b[0mpruned_model\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0msparsity\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mprune_low_magnitude\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mbase_model\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mpruning_params\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m",
            "\u001b[0;31mNameError\u001b[0m: name 'pruning_params' is not defined"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 14,
      "id": "5810ac36",
      "metadata": {
        "id": "5810ac36"
      },
      "outputs": [],
      "source": [
        "def clahe(cliplimit,gridsize,path):\n",
        "    for img in os.listdir(path):\n",
        "        imgarray = Image.open(os.path.join(path, img))\n",
        "    lab = cv2.cvtColor(img_array, cv2.COLOR_BGR2Lab)\n",
        "\n",
        "        # -----Splitting the LAB image to different channels-------------------------\n",
        "    l, a, b = cv2.split(lab)\n",
        "\n",
        "        # -----Applying CLAHE to L-channel-------------------------------------------\n",
        "    clahe = cv2.createCLAHE(clipLimit=cliplimit, tileGridSize=(gridsize, gridsize))\n",
        "    cl = clahe.apply(l)\n",
        "\n",
        "        # -----Merge the CLAHE enhanced L-channel with the a and b channel-----------\n",
        "    limg = cv2.merge((cl, a, b))\n",
        "\n",
        "        # -----Converting image from LAB Color model to RGB model--------------------\n",
        "    final = cv2.cvtColor(limg, cv2.COLOR_LAB2BGR)\n",
        "    return final"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "pip install -q tensorflow-model-optimization"
      ],
      "metadata": {
        "id": "jnzpib8pGS6i"
      },
      "id": "jnzpib8pGS6i",
      "execution_count": 18,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "from tensorflow_model_optimization.sparsity import keras as sparsity\n"
      ],
      "metadata": {
        "id": "iYbmYcSJGMMz"
      },
      "id": "iYbmYcSJGMMz",
      "execution_count": 70,
      "outputs": []
    },
    {
      "cell_type": "code",
      "execution_count": 65,
      "id": "882ac25c",
      "metadata": {
        "id": "882ac25c"
      },
      "outputs": [],
      "source": [
        "def train_model(model,dataset,image_size,batch_size,save_name,lr1,lr2,Epochs1,Epochs2):\n",
        "    \n",
        "    #dataParam={'messidor': [957,243,2,'./data/messidor/train','./data/messidor/test'],\n",
        "              #  'kaggle': [30000,5126,5,'./data/kaggle/train','./data/kaggle/valid'],\n",
        "              #  'DDR':   [9851,2503,5,'./data/DDR/train','./data/DDR/valid']} #6119\n",
        "    dataParam = {'DR': [1000,1000,5,'//content/drive/MyDrive/CABnet-master/data/train','//content/drive/MyDrive/CABnet-master/data/val']}\n",
        "\n",
        "    train_num,valid_num,classes,train_dir,test_dir = dataParam[dataset]\n",
        "    \n",
        "    train=ImageDataGenerator(horizontal_flip=True,vertical_flip=True,rotation_range=90)          \n",
        "    valid = ImageDataGenerator()\n",
        "    train_data=train.flow_from_directory(train_dir,\n",
        "                                         target_size=(image_size,image_size),\n",
        "                                         shuffle = True,\n",
        "                                         batch_size=batch_size)\n",
        "    valid_data=valid.flow_from_directory(test_dir,\n",
        "                                         target_size=(image_size,image_size),\n",
        "                                         shuffle = False,\n",
        "                                         batch_size=batch_size)\n",
        "\n",
        "    lr_decay=ReduceLROnPlateau(monitor='val_loss', factor=0.8, patience=3, verbose=1)\n",
        "    #save_model=ModelCheckpoint('new/'+save_name+'{epoch:02d}.h5', monitor='val_loss',period=10)\n",
        "    # Define the pruning parameters\n",
        "    pruning_params = {\n",
        "    'pruning_schedule': sparsity.PolynomialDecay(initial_sparsity=0.50,\n",
        "                                                  final_sparsity=0.90,\n",
        "                                                  begin_step=0,\n",
        "                                                  end_step=5,\n",
        "                                                  frequency=100)\n",
        "        }\n",
        "\n",
        "\n",
        "    \n",
        "    for layer in base_model.layers:\n",
        "        layer.trainable = False   \n",
        "\n",
        "    #  Apply pruning to the model\n",
        "\n",
        "    pruned_model.compile(optimizer=Adam(lr=lr1,decay=0.00001),loss=loss_fun,metrics=['acc'])\n",
        "    pruned_model.fit_generator(train_data,\n",
        "                        steps_per_epoch=train_num/batch_size,\n",
        "                        validation_data=valid_data,\n",
        "                        validation_steps=valid_num/batch_size,\n",
        "                        epochs=Epochs1, \n",
        "                        workers=2,\n",
        "                        callbacks=[lr_decay])   \n",
        "    \n",
        "    for layer in base_model.layers:\n",
        "        layer.trainable = True\n",
        "    \n",
        "    \n",
        "    pruned_model.compile(optimizer=Adam(lr=lr2,decay=0.00001),loss=loss_fun,metrics=['acc'])\n",
        "    pruned_model.fit(train_data,\n",
        "                        steps_per_epoch=train_num/batch_size,\n",
        "                        validation_data=valid_data,\n",
        "                        validation_steps=valid_num/batch_size,\n",
        "                        epochs=Epochs2,\n",
        "                        workers=2,\n",
        "                        callbacks=[lr_decay])\n",
        "    \n",
        "    \n",
        "    \n",
        "    \n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 66,
      "id": "abe17c12",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "abe17c12",
        "outputId": "54f65b4c-e5ac-48f1-8de9-babbcd22ca6c"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Model: \"model_13\"\n",
            "__________________________________________________________________________________________________\n",
            " Layer (type)                   Output Shape         Param #     Connected to                     \n",
            "==================================================================================================\n",
            " input_14 (InputLayer)          [(None, 224, 224, 3  0           []                               \n",
            "                                )]                                                                \n",
            "                                                                                                  \n",
            " block1_conv1 (Conv2D)          (None, 224, 224, 64  1792        ['input_14[0][0]']               \n",
            "                                )                                                                 \n",
            "                                                                                                  \n",
            " block1_conv2 (Conv2D)          (None, 224, 224, 64  36928       ['block1_conv1[0][0]']           \n",
            "                                )                                                                 \n",
            "                                                                                                  \n",
            " block1_pool (MaxPooling2D)     (None, 112, 112, 64  0           ['block1_conv2[0][0]']           \n",
            "                                )                                                                 \n",
            "                                                                                                  \n",
            " block2_conv1 (Conv2D)          (None, 112, 112, 12  73856       ['block1_pool[0][0]']            \n",
            "                                8)                                                                \n",
            "                                                                                                  \n",
            " block2_conv2 (Conv2D)          (None, 112, 112, 12  147584      ['block2_conv1[0][0]']           \n",
            "                                8)                                                                \n",
            "                                                                                                  \n",
            " block2_pool (MaxPooling2D)     (None, 56, 56, 128)  0           ['block2_conv2[0][0]']           \n",
            "                                                                                                  \n",
            " block3_conv1 (Conv2D)          (None, 56, 56, 256)  295168      ['block2_pool[0][0]']            \n",
            "                                                                                                  \n",
            " block3_conv2 (Conv2D)          (None, 56, 56, 256)  590080      ['block3_conv1[0][0]']           \n",
            "                                                                                                  \n",
            " block3_conv3 (Conv2D)          (None, 56, 56, 256)  590080      ['block3_conv2[0][0]']           \n",
            "                                                                                                  \n",
            " block3_pool (MaxPooling2D)     (None, 28, 28, 256)  0           ['block3_conv3[0][0]']           \n",
            "                                                                                                  \n",
            " block4_conv1 (Conv2D)          (None, 28, 28, 512)  1180160     ['block3_pool[0][0]']            \n",
            "                                                                                                  \n",
            " block4_conv2 (Conv2D)          (None, 28, 28, 512)  2359808     ['block4_conv1[0][0]']           \n",
            "                                                                                                  \n",
            " block4_conv3 (Conv2D)          (None, 28, 28, 512)  2359808     ['block4_conv2[0][0]']           \n",
            "                                                                                                  \n",
            " block4_pool (MaxPooling2D)     (None, 14, 14, 512)  0           ['block4_conv3[0][0]']           \n",
            "                                                                                                  \n",
            " block5_conv1 (Conv2D)          (None, 14, 14, 512)  2359808     ['block4_pool[0][0]']            \n",
            "                                                                                                  \n",
            " block5_conv2 (Conv2D)          (None, 14, 14, 512)  2359808     ['block5_conv1[0][0]']           \n",
            "                                                                                                  \n",
            " block5_conv3 (Conv2D)          (None, 14, 14, 512)  2359808     ['block5_conv2[0][0]']           \n",
            "                                                                                                  \n",
            " block5_pool (MaxPooling2D)     (None, 7, 7, 512)    0           ['block5_conv3[0][0]']           \n",
            "                                                                                                  \n",
            " average_pooling2d_13 (AverageP  (None, 1, 1, 512)   0           ['block5_pool[0][0]']            \n",
            " ooling2D)                                                                                        \n",
            "                                                                                                  \n",
            " conv2d_39 (Conv2D)             (None, 1, 1, 512)    262656      ['average_pooling2d_13[0][0]']   \n",
            "                                                                                                  \n",
            " activation_52 (Activation)     (None, 1, 1, 512)    0           ['conv2d_39[0][0]']              \n",
            "                                                                                                  \n",
            " conv2d_40 (Conv2D)             (None, 1, 1, 512)    262656      ['activation_52[0][0]']          \n",
            "                                                                                                  \n",
            " activation_53 (Activation)     (None, 1, 1, 512)    0           ['conv2d_40[0][0]']              \n",
            "                                                                                                  \n",
            " multiply_52 (Multiply)         (None, 7, 7, 512)    0           ['activation_53[0][0]',          \n",
            "                                                                  'block5_pool[0][0]']            \n",
            "                                                                                                  \n",
            " lambda_52 (Lambda)             (None, 7, 7, 1)      0           ['multiply_52[0][0]']            \n",
            "                                                                                                  \n",
            " activation_54 (Activation)     (None, 7, 7, 1)      0           ['lambda_52[0][0]']              \n",
            "                                                                                                  \n",
            " multiply_53 (Multiply)         (None, 7, 7, 512)    0           ['activation_54[0][0]',          \n",
            "                                                                  'multiply_52[0][0]']            \n",
            "                                                                                                  \n",
            " conv2d_41 (Conv2D)             (None, 7, 7, 25)     12825       ['multiply_53[0][0]']            \n",
            "                                                                                                  \n",
            " batch_normalization_13 (BatchN  (None, 7, 7, 25)    100         ['conv2d_41[0][0]']              \n",
            " ormalization)                                                                                    \n",
            "                                                                                                  \n",
            " activation_55 (Activation)     (None, 7, 7, 25)     0           ['batch_normalization_13[0][0]'] \n",
            "                                                                                                  \n",
            " global_max_pooling2d_13 (Globa  (None, 25)          0           ['activation_55[0][0]']          \n",
            " lMaxPooling2D)                                                                                   \n",
            "                                                                                                  \n",
            " reshape_26 (Reshape)           (None, 5, 5)         0           ['global_max_pooling2d_13[0][0]']\n",
            "                                                                                                  \n",
            " reshape_27 (Reshape)           (None, 7, 7, 5, 5)   0           ['activation_55[0][0]']          \n",
            "                                                                                                  \n",
            " lambda_53 (Lambda)             (None, 5)            0           ['reshape_26[0][0]']             \n",
            "                                                                                                  \n",
            " lambda_54 (Lambda)             (None, 7, 7, 5)      0           ['reshape_27[0][0]']             \n",
            "                                                                                                  \n",
            " multiply_54 (Multiply)         (None, 7, 7, 5)      0           ['lambda_53[0][0]',              \n",
            "                                                                  'lambda_54[0][0]']              \n",
            "                                                                                                  \n",
            " lambda_55 (Lambda)             (None, 7, 7, 1)      0           ['multiply_54[0][0]']            \n",
            "                                                                                                  \n",
            " multiply_55 (Multiply)         (None, 7, 7, 512)    0           ['multiply_53[0][0]',            \n",
            "                                                                  'lambda_55[0][0]']              \n",
            "                                                                                                  \n",
            " global_average_pooling2d_13 (G  (None, 512)         0           ['multiply_55[0][0]']            \n",
            " lobalAveragePooling2D)                                                                           \n",
            "                                                                                                  \n",
            " dense_13 (Dense)               (None, 5)            2565        ['global_average_pooling2d_13[0][\n",
            "                                                                 0]']                             \n",
            "                                                                                                  \n",
            "==================================================================================================\n",
            "Total params: 15,255,490\n",
            "Trainable params: 15,255,440\n",
            "Non-trainable params: 50\n",
            "__________________________________________________________________________________________________\n"
          ]
        }
      ],
      "source": [
        "os.environ[\"CUDA_VISIBLE_DEVICES\"] = \"0\"    \n",
        "loss_fun= 'categorical_crossentropy'  \n",
        "gpu_num=1\n",
        "k=5\n",
        "lr1=0.005\n",
        "lr2=0.0001\n",
        "batch_size= 32\n",
        "image_size=224\n",
        "classes=5\n",
        "\n",
        "base_model=get_base_model('vgg16',image_size)  \n",
        "base_in=base_model.input\n",
        "base_out=base_model.output\n",
        "\n",
        "x=Global_attention_block(base_out)\n",
        "base_out=Category_attention_block(x,classes,k)\n",
        "\n",
        "x=GlobalAveragePooling2D()(base_out)\n",
        "out=Dense(classes,activation='softmax')(x)\n",
        "\n",
        "if gpu_num>1:\n",
        "    model=Model(base_model.input,out)\n",
        "    #model.summary()\n",
        "    parallel_model = multi_gpu_model(model, gpus=gpu_num)\n",
        "    parallel_model.summary()\n",
        "else:\n",
        "    parallel_model=Model(base_model.input,out)\n",
        "    parallel_model.summary()\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 67,
      "id": "m72K6fc2TcyE",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 339
        },
        "id": "m72K6fc2TcyE",
        "outputId": "e1987933-0fdb-43e2-f926-c6055d44f302",
        "scrolled": true
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Found 21073 images belonging to 5 classes.\n",
            "Found 7023 images belonging to 5 classes.\n"
          ]
        },
        {
          "output_type": "error",
          "ename": "NameError",
          "evalue": "ignored",
          "traceback": [
            "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
            "\u001b[0;31mNameError\u001b[0m                                 Traceback (most recent call last)",
            "\u001b[0;32m<ipython-input-67-d4820445f904>\u001b[0m in \u001b[0;36m<module>\u001b[0;34m\u001b[0m\n\u001b[0;32m----> 1\u001b[0;31m \u001b[0mhistory\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mtrain_model\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mparallel_model\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m'DR'\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0mimage_size\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0mbatch_size\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m'vgg16'\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0mlr1\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0mlr2\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;36m1\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;36m2\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      2\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n",
            "\u001b[0;32m<ipython-input-65-8b6e7992b356>\u001b[0m in \u001b[0;36mtrain_model\u001b[0;34m(model, dataset, image_size, batch_size, save_name, lr1, lr2, Epochs1, Epochs2)\u001b[0m\n\u001b[1;32m     37\u001b[0m     \u001b[0;31m#  Apply pruning to the model\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     38\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 39\u001b[0;31m     \u001b[0mpruned_model\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mcompile\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0moptimizer\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mAdam\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mlr\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mlr1\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0mdecay\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;36m0.00001\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0mloss\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mloss_fun\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0mmetrics\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;34m'acc'\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[0m\u001b[1;32m     40\u001b[0m     pruned_model.fit_generator(train_data,\n\u001b[1;32m     41\u001b[0m                         \u001b[0msteps_per_epoch\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mtrain_num\u001b[0m\u001b[0;34m/\u001b[0m\u001b[0mbatch_size\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
            "\u001b[0;31mNameError\u001b[0m: name 'pruned_model' is not defined"
          ]
        }
      ],
      "source": [
        "history=train_model(parallel_model,'DR',image_size,batch_size,'vgg16',lr1,lr2,1,2)\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "HFC8Zktza0DN",
      "metadata": {
        "id": "HFC8Zktza0DN"
      },
      "outputs": [],
      "source": [
        "plotmodel(history,'vgg16')"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "bc70b154",
      "metadata": {
        "id": "bc70b154"
      },
      "outputs": [],
      "source": [
        "# from keras.preprocessing.image import ImageDataGenerator\n",
        "# from keras.models import load_model\n",
        "# from keras import backend as K\n",
        "# from keras.preprocessing import image\n",
        "# import os\n",
        "# import numpy as np"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "6ef87605",
      "metadata": {
        "id": "6ef87605"
      },
      "outputs": [],
      "source": []
    },
    {
      "cell_type": "markdown",
      "id": "819a36ad",
      "metadata": {
        "id": "819a36ad"
      },
      "source": [
        "def label_smooth(y_true, y_pred):\n",
        "\n",
        "    y_true=((1-0.1)*y_true+0.05)\n",
        "    return K.categorical_crossentropy(y_true, y_pred)\n",
        "    \n",
        "\n",
        "def weight_kappa(result,test_num):\n",
        "    weight=np.zeros((5,5),dtype='float')\n",
        "    for i in range(5):\n",
        "        for j in range(5):\n",
        "            weight[i,j]=(i-j)*(i-j)/16\n",
        "    fenzi=0\n",
        "    for i in range(5):\n",
        "        for j in range(5):\n",
        "            fenzi=fenzi+result[i,j]*weight[i,j]\n",
        "    fenmu=0\n",
        "    for i in range(5):\n",
        "        for j in range(5):\n",
        "            fenmu=fenmu+weight[i,j]*result[:,j].sum()*result[i,:].sum()\n",
        "\n",
        "    weght_kappa=1-(fenzi/(fenmu/test_num))\n",
        "    return  weght_kappa"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "874bb33e",
      "metadata": {
        "id": "874bb33e"
      },
      "source": [
        "os.environ[\"CUDA_VISIBLE_DEVICES\"] = \"7\"\n",
        "image_size=512\n",
        "batch_size=64\n",
        "model_name='inceptionv370.h5'\n",
        "test_dir=''  \n",
        "custom_test=True\n",
        "\n",
        "\n",
        "if custom_test== False:   \n",
        "    model=load_model('new/'+model_name)\n",
        "else:\n",
        "    model=load_model('new/'+model_name,custom_objects={'label_smooth': label_smooth})\n",
        "\n",
        "test_num=0\n",
        "result=np.zeros((5,5),dtype=int)\n",
        "recall=np.zeros((1,5),dtype=float)\n",
        "for i in range(5):\n",
        "        datadirs=test_dir+str(i)+'/'\n",
        "        filenames=os.listdir(datadirs)\n",
        "        num=len(filenames)\n",
        "        test_num=test_num+num\n",
        "        valid = ImageDataGenerator()\n",
        "        valid_data=valid.flow_from_directory(directory=test_dir,target_size=(image_size,image_size),\n",
        "                                             batch_size=batch_size,class_mode=None,classes=str(i))\n",
        "        predict=model.predict_generator(valid_data,steps=num/batch_size,verbose=1,workers=1)\n",
        "        predict=np.argmax(predict,axis=-1)\n",
        "        for j in range(5):\n",
        "            result[i,j]=np.sum(predict==j)\n",
        "\n",
        "right=result[0,0]+result[1,1]+result[2,2]+result[3,3]+result[4,4]\n",
        "print('Acc:',right/test_num)\n",
        "\n",
        "w_kappa=weight_kappa(result,test_num)\n",
        "print('w_kappa:',w_kappa)\n"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "pip install -q tensorflow-model-optimization"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "h-jCyL8sdGj0",
        "outputId": "a8d06053-ffb1-4fef-d42a-5503b86948ab"
      },
      "id": "h-jCyL8sdGj0",
      "execution_count": 17,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "\u001b[?25l     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m0.0/238.9 KB\u001b[0m \u001b[31m?\u001b[0m eta \u001b[36m-:--:--\u001b[0m\r\u001b[2K     \u001b[91m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m\u001b[90m╺\u001b[0m \u001b[32m235.5/238.9 KB\u001b[0m \u001b[31m8.5 MB/s\u001b[0m eta \u001b[36m0:00:01\u001b[0m\r\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m238.9/238.9 KB\u001b[0m \u001b[31m6.3 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25h"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "id": "2bcc4762",
      "metadata": {
        "id": "2bcc4762"
      },
      "source": [
        "## Post Training Quantization"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "c65f71e8",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "c65f71e8",
        "outputId": "18712134-8dee-4ee7-a2bc-ae10fde548b6"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "TensorFlow version: 2.11.0\n",
            "Keras version: 2.11.0\n",
            "TensorFlow Model Optimization version: 0.7.3\n"
          ]
        }
      ],
      "source": [
        "import tensorflow as tf\n",
        "import tensorflow_model_optimization as tfmot\n",
        "\n",
        "# Print TensorFlow info\n",
        "print('TensorFlow version: '+tf.__version__)\n",
        "print('Keras version: '+tf.keras.__version__)\n",
        "print('TensorFlow Model Optimization version: '+tfmot.__version__)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "c6ffb14d",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "c6ffb14d",
        "outputId": "63106178-e9ba-441d-bdc5-bbe50ff8c759"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "WARNING:absl:Found untraced functions such as _jit_compiled_convolution_op, _jit_compiled_convolution_op, _jit_compiled_convolution_op, _jit_compiled_convolution_op, _jit_compiled_convolution_op while saving (showing 5 of 16). These functions will not be directly callable after loading.\n"
          ]
        }
      ],
      "source": [
        "#Without Quantization\n",
        "tflite_model = tf.keras.models.load_model(\"/content/drive/MyDrive/CABnet-master/vgg1621.h5\")\n",
        "converter = tf.lite.TFLiteConverter.from_keras_model(tflite_model)\n",
        "tflite_save = converter.convert()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "9a656301",
      "metadata": {
        "id": "9a656301",
        "outputId": "98a4ae81-1ee6-411b-80ce-32af6e894ef5",
        "colab": {
          "base_uri": "https://localhost:8080/"
        }
      },
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "61036660"
            ]
          },
          "metadata": {},
          "execution_count": 60
        }
      ],
      "source": [
        "len(tflite_save)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "d359c10e",
      "metadata": {
        "id": "d359c10e",
        "outputId": "1e45cd48-9e93-42af-cf26-3f54755a9e43",
        "colab": {
          "base_uri": "https://localhost:8080/"
        }
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "WARNING:absl:Found untraced functions such as _jit_compiled_convolution_op, _jit_compiled_convolution_op, _jit_compiled_convolution_op, _jit_compiled_convolution_op, _jit_compiled_convolution_op while saving (showing 5 of 16). These functions will not be directly callable after loading.\n"
          ]
        }
      ],
      "source": [
        "#With Quantization\n",
        "tflite_model = tf.keras.models.load_model(\"/content/drive/MyDrive/CABnet-master/vgg1621.h5\")\n",
        "converter = tf.lite.TFLiteConverter.from_keras_model(tflite_model)\n",
        "converter.optimizations = [tf.lite.Optimize.DEFAULT]\n",
        "tflite_post_quant_model = converter.convert()\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "63edb957",
      "metadata": {
        "id": "63edb957",
        "outputId": "105c8553-6634-4eb5-f660-b612c963e030",
        "colab": {
          "base_uri": "https://localhost:8080/"
        }
      },
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "15350064"
            ]
          },
          "metadata": {},
          "execution_count": 90
        }
      ],
      "source": [
        "len(tflite_post_quant_model)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "042831bb",
      "metadata": {
        "id": "042831bb"
      },
      "outputs": [],
      "source": [
        "with open(\"tflite_save.tflite\", \"wb\") as f:\n",
        "    f.write(tflite_save)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "847fbea9",
      "metadata": {
        "id": "847fbea9"
      },
      "outputs": [],
      "source": [
        "with open(\"tflite_quant_model.tflite\", \"wb\") as f:\n",
        "    f.write(tflite_post_quant_model)"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "8df1bb8c",
      "metadata": {
        "id": "8df1bb8c"
      },
      "source": [
        "## Full Integer Quantization:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "66c2958f",
      "metadata": {
        "scrolled": true,
        "id": "66c2958f",
        "outputId": "a2eb41e4-2c3a-44e0-ccca-2e5cd0962abe",
        "colab": {
          "base_uri": "https://localhost:8080/"
        }
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "WARNING:absl:Found untraced functions such as _jit_compiled_convolution_op, _jit_compiled_convolution_op, _jit_compiled_convolution_op, _jit_compiled_convolution_op, _jit_compiled_convolution_op while saving (showing 5 of 16). These functions will not be directly callable after loading.\n"
          ]
        }
      ],
      "source": [
        "\n",
        "# Load the pre-trained model\n",
        "tflite_model = tf.keras.models.load_model(\"/content/drive/MyDrive/CABnet-master/vgg1621.h5\")\n",
        "\n",
        "\n",
        "# Create a full integer quantization object\n",
        "#quantize_provider = tfmot.quantization.keras.quantize_annotate_layer(to_annotate, quantize_config=None)\n",
        "#quantize_model = quantize_provider(tflite_model)\n",
        "quantize_model = tflite_model\n",
        "\n",
        "\n",
        "# Define the quantization parameters\n",
        "converter = tf.lite.TFLiteConverter.from_keras_model(quantize_model)\n",
        "\n",
        "converter.optimizations = [tf.lite.Optimize.DEFAULT]\n",
        "\n",
        "# Define the representative dataset\n",
        "def representative_dataset_gen():\n",
        "    for i in range(100):\n",
        "        yield [x_train[i:i+1]]\n",
        "        \n",
        "\n",
        "converter.representative_dataset_gen = representative_dataset_gen()\n",
        "\n",
        "\n",
        "# Convert the model to TFLite format\n",
        "tflite_full_integer_model = converter.convert()\n",
        "\n",
        "# Save the TFLite model to disk\n",
        "with open(\"model.tflite\", \"wb\") as f:\n",
        "    f.write(tflite_full_integer_model)\n",
        "\n"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "adba60e0",
      "metadata": {
        "id": "adba60e0"
      },
      "source": [
        "## Quantization Aware Training"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "a9d0aae9",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "a9d0aae9",
        "outputId": "7dfb44db-3e5e-45a9-c4f9-5efcbd8e1195"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "WARNING:tensorflow:From /usr/local/lib/python3.8/dist-packages/tensorflow/python/autograph/pyct/static_analysis/liveness.py:83: Analyzer.lamba_check (from tensorflow.python.autograph.pyct.static_analysis.liveness) is deprecated and will be removed after 2023-09-23.\n",
            "Instructions for updating:\n",
            "Lambda fuctions will be no more assumed to be used in the statement where they are used, or at least in the same block. /~https://github.com/tensorflow/tensorflow/issues/56089\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Model: \"vgg16\"\n",
            "_________________________________________________________________\n",
            " Layer (type)                Output Shape              Param #   \n",
            "=================================================================\n",
            " input_1 (InputLayer)        [(None, 224, 224, 3)]     0         \n",
            "                                                                 \n",
            " quantize_layer (QuantizeLay  (None, 224, 224, 3)      3         \n",
            " er)                                                             \n",
            "                                                                 \n",
            " quant_block1_conv1 (Quantiz  (None, 224, 224, 64)     1923      \n",
            " eWrapperV2)                                                     \n",
            "                                                                 \n",
            " quant_block1_conv2 (Quantiz  (None, 224, 224, 64)     37059     \n",
            " eWrapperV2)                                                     \n",
            "                                                                 \n",
            " quant_block1_pool (Quantize  (None, 112, 112, 64)     1         \n",
            " WrapperV2)                                                      \n",
            "                                                                 \n",
            " quant_block2_conv1 (Quantiz  (None, 112, 112, 128)    74115     \n",
            " eWrapperV2)                                                     \n",
            "                                                                 \n",
            " quant_block2_conv2 (Quantiz  (None, 112, 112, 128)    147843    \n",
            " eWrapperV2)                                                     \n",
            "                                                                 \n",
            " quant_block2_pool (Quantize  (None, 56, 56, 128)      1         \n",
            " WrapperV2)                                                      \n",
            "                                                                 \n",
            " quant_block3_conv1 (Quantiz  (None, 56, 56, 256)      295683    \n",
            " eWrapperV2)                                                     \n",
            "                                                                 \n",
            " quant_block3_conv2 (Quantiz  (None, 56, 56, 256)      590595    \n",
            " eWrapperV2)                                                     \n",
            "                                                                 \n",
            " quant_block3_conv3 (Quantiz  (None, 56, 56, 256)      590595    \n",
            " eWrapperV2)                                                     \n",
            "                                                                 \n",
            " quant_block3_pool (Quantize  (None, 28, 28, 256)      1         \n",
            " WrapperV2)                                                      \n",
            "                                                                 \n",
            " quant_block4_conv1 (Quantiz  (None, 28, 28, 512)      1181187   \n",
            " eWrapperV2)                                                     \n",
            "                                                                 \n",
            " quant_block4_conv2 (Quantiz  (None, 28, 28, 512)      2360835   \n",
            " eWrapperV2)                                                     \n",
            "                                                                 \n",
            " quant_block4_conv3 (Quantiz  (None, 28, 28, 512)      2360835   \n",
            " eWrapperV2)                                                     \n",
            "                                                                 \n",
            " quant_block4_pool (Quantize  (None, 14, 14, 512)      1         \n",
            " WrapperV2)                                                      \n",
            "                                                                 \n",
            " quant_block5_conv1 (Quantiz  (None, 14, 14, 512)      2360835   \n",
            " eWrapperV2)                                                     \n",
            "                                                                 \n",
            " quant_block5_conv2 (Quantiz  (None, 14, 14, 512)      2360835   \n",
            " eWrapperV2)                                                     \n",
            "                                                                 \n",
            " quant_block5_conv3 (Quantiz  (None, 14, 14, 512)      2360835   \n",
            " eWrapperV2)                                                     \n",
            "                                                                 \n",
            " quant_block5_pool (Quantize  (None, 7, 7, 512)        1         \n",
            " WrapperV2)                                                      \n",
            "                                                                 \n",
            "=================================================================\n",
            "Total params: 14,723,183\n",
            "Trainable params: 14,714,688\n",
            "Non-trainable params: 8,495\n",
            "_________________________________________________________________\n"
          ]
        }
      ],
      "source": [
        "import tensorflow_model_optimization as tfmot\n",
        "\n",
        "quantize_model = tfmot.quantization.keras.quantize_model\n",
        "\n",
        "# q_aware stands for for quantization aware.\n",
        "q_aware_model = quantize_model(base_model)\n",
        "\n",
        "# `quantize_model` requires a recompile.\n",
        "\n",
        "q_aware_model.compile(optimizer='adam',\n",
        "              loss='sparse_categorical_crossentropy',\n",
        "              metrics=['accuracy'])\n",
        "q_aware_model.summary()\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "5bca5f7b",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "5bca5f7b",
        "outputId": "09f5ad97-cf1d-4624-a24f-27bd3c888136"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "WARNING:absl:Found untraced functions such as _jit_compiled_convolution_op, _jit_compiled_convolution_op, _jit_compiled_convolution_op, _jit_compiled_convolution_op, _jit_compiled_convolution_op while saving (showing 5 of 13). These functions will not be directly callable after loading.\n"
          ]
        }
      ],
      "source": [
        "import tensorflow as tf\n",
        "from tensorflow import keras\n",
        "from tensorflow.keras import layers\n",
        "from tensorflow.keras.preprocessing.image import ImageDataGenerator\n",
        "\n",
        "# Convert the model to a quantization aware model\n",
        "quantize_model = tf.keras.models.clone_model(base_model)\n",
        "quantize_model = tfmot.quantization.keras.quantize_model(quantize_model)\n",
        "\n",
        "# Convert the quantization aware model to TensorFlow Lite\n",
        "converter = tf.lite.TFLiteConverter.from_keras_model(base_model)\n",
        "converter.optimizations = [tf.lite.Optimize.DEFAULT]\n",
        "tflite_model = converter.convert()\n",
        "\n",
        "# Save the TensorFlow Lite model to disk\n",
        "with open('model.tflite', 'wb') as f:\n",
        "    f.write(tflite_model)\n",
        "\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "import tensorflow as tf\n",
        "from tensorflow.keras.preprocessing.image import ImageDataGenerator\n",
        "from tensorflow_model_optimization.sparsity import keras as sparsity\n",
        "\n",
        "\n",
        "# Define the paths to the training and validation data\n",
        "train_dir = '/content/drive/MyDrive/CABnet-master/data/train'\n",
        "val_dir = '/content/drive/MyDrive/CABnet-master/data/val'\n",
        "\n",
        "# Define the input shape of the model\n",
        "input_shape = (224, 224, 3)\n",
        "\n",
        "# Create a data generator for the training data\n",
        "train_datagen = ImageDataGenerator(\n",
        "    rescale=1./255,\n",
        "    shear_range=0.2,\n",
        "    zoom_range=0.2,\n",
        "    horizontal_flip=True\n",
        ")\n",
        "train_data = train_datagen.flow_from_directory(\n",
        "    train_dir,\n",
        "    target_size=input_shape[:2],\n",
        "    batch_size=32,\n",
        "    class_mode='categorical'\n",
        ")\n",
        "\n",
        "# Create a data generator for the validation data\n",
        "val_datagen = ImageDataGenerator(rescale=1./255)\n",
        "val_data = val_datagen.flow_from_directory(\n",
        "    val_dir,\n",
        "    target_size=input_shape[:2],\n",
        "    batch_size=32,\n",
        "    class_mode='categorical'\n",
        ")\n",
        "\n",
        "# Load the Keras model from the .h5 file\n",
        "#keras_model_path = '/path/to/my_keras_model.h5'\n",
        "#keras_model = tf.keras.models.load_model(base_model)\n",
        "\n",
        "# Define the pruning parameters\n",
        "pruning_params = {\n",
        "    'pruning_schedule': sparsity.PolynomialDecay(initial_sparsity=0.50,\n",
        "                                                  final_sparsity=0.90,\n",
        "                                                  begin_step=0,\n",
        "                                                  end_step=5,\n",
        "                                                  frequency=100)\n",
        "}\n",
        "\n",
        "# Apply pruning to the model\n",
        "pruned_model = sparsity.prune_low_magnitude(base_model, **pruning_params)\n",
        "\n",
        "# Compile the pruned model\n",
        "pruned_model.compile(optimizer='adam', loss='categorical_crossentropy', metrics=['accuracy'])\n",
        "\n",
        "# Train the pruned model\n",
        "num_epochs = 10\n",
        "pruned_model.fit(train_data, epochs=num_epochs, validation_data=val_data)\n",
        "\n",
        "# Strip the pruning parameters from the pruned model\n",
        "stripped_pruned_model = sparsity.strip_pruning(pruned_model)\n",
        "\n",
        "# Save the pruned model to a .h5 file\n",
        "tf.keras.models.save_model(stripped_pruned_model, 'my_pruned_model.h5', include_optimizer=False)\n"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 831
        },
        "id": "FJ-LRj242dGT",
        "outputId": "db07671b-abc1-4b34-e02d-95991127cfa4"
      },
      "id": "FJ-LRj242dGT",
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Found 21073 images belonging to 5 classes.\n",
            "Found 7023 images belonging to 5 classes.\n",
            "Epoch 1/10\n"
          ]
        },
        {
          "output_type": "error",
          "ename": "ValueError",
          "evalue": "ignored",
          "traceback": [
            "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
            "\u001b[0;31mValueError\u001b[0m                                Traceback (most recent call last)",
            "\u001b[0;32m<ipython-input-20-aa6c0c98c31e>\u001b[0m in \u001b[0;36m<module>\u001b[0;34m\u001b[0m\n\u001b[1;32m     54\u001b[0m \u001b[0;31m# Train the pruned model\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     55\u001b[0m \u001b[0mnum_epochs\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;36m10\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 56\u001b[0;31m \u001b[0mpruned_model\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mfit\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mtrain_data\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mepochs\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mnum_epochs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mvalidation_data\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mval_data\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     57\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     58\u001b[0m \u001b[0;31m# Strip the pruning parameters from the pruned model\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
            "\u001b[0;32m/usr/local/lib/python3.8/dist-packages/keras/utils/traceback_utils.py\u001b[0m in \u001b[0;36merror_handler\u001b[0;34m(*args, **kwargs)\u001b[0m\n\u001b[1;32m     68\u001b[0m             \u001b[0;31m# To get the full stack trace, call:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     69\u001b[0m             \u001b[0;31m# `tf.debugging.disable_traceback_filtering()`\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 70\u001b[0;31m             \u001b[0;32mraise\u001b[0m \u001b[0me\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mwith_traceback\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mfiltered_tb\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;32mfrom\u001b[0m \u001b[0;32mNone\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     71\u001b[0m         \u001b[0;32mfinally\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     72\u001b[0m             \u001b[0;32mdel\u001b[0m \u001b[0mfiltered_tb\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
            "\u001b[0;32m/usr/local/lib/python3.8/dist-packages/keras/engine/training.py\u001b[0m in \u001b[0;36mtf__train_function\u001b[0;34m(iterator)\u001b[0m\n\u001b[1;32m     13\u001b[0m                 \u001b[0;32mtry\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[0mdo_return\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;32mTrue\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 15\u001b[0;31m                     \u001b[0mretval_\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mag__\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mconverted_call\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mag__\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mld\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mstep_function\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m(\u001b[0m\u001b[0mag__\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mld\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mag__\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mld\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0miterator\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;32mNone\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mfscope\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;32mexcept\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     17\u001b[0m                     \u001b[0mdo_return\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;32mFalse\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
            "\u001b[0;31mValueError\u001b[0m: in user code:\n\n    File \"/usr/local/lib/python3.8/dist-packages/keras/engine/training.py\", line 1249, in train_function  *\n        return step_function(self, iterator)\n    File \"/usr/local/lib/python3.8/dist-packages/keras/engine/training.py\", line 1233, in step_function  **\n        outputs = model.distribute_strategy.run(run_step, args=(data,))\n    File \"/usr/local/lib/python3.8/dist-packages/keras/engine/training.py\", line 1222, in run_step  **\n        outputs = model.train_step(data)\n    File \"/usr/local/lib/python3.8/dist-packages/keras/engine/training.py\", line 1024, in train_step\n        loss = self.compute_loss(x, y, y_pred, sample_weight)\n    File \"/usr/local/lib/python3.8/dist-packages/keras/engine/training.py\", line 1082, in compute_loss\n        return self.compiled_loss(\n    File \"/usr/local/lib/python3.8/dist-packages/keras/engine/compile_utils.py\", line 265, in __call__\n        loss_value = loss_obj(y_t, y_p, sample_weight=sw)\n    File \"/usr/local/lib/python3.8/dist-packages/keras/losses.py\", line 152, in __call__\n        losses = call_fn(y_true, y_pred)\n    File \"/usr/local/lib/python3.8/dist-packages/keras/losses.py\", line 284, in call  **\n        return ag_fn(y_true, y_pred, **self._fn_kwargs)\n    File \"/usr/local/lib/python3.8/dist-packages/keras/losses.py\", line 2004, in categorical_crossentropy\n        return backend.categorical_crossentropy(\n    File \"/usr/local/lib/python3.8/dist-packages/keras/backend.py\", line 5532, in categorical_crossentropy\n        target.shape.assert_is_compatible_with(output.shape)\n\n    ValueError: Shapes (None, None) and (None, 7, 7, 512) are incompatible\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [],
      "metadata": {
        "id": "5VRedpvX2coa"
      },
      "id": "5VRedpvX2coa"
    },
    {
      "cell_type": "code",
      "source": [
        "from google.colab import drive\n",
        "drive.mount('/content/drive')"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "pvLsTaa9dbiY",
        "outputId": "d3db1c68-bc8b-42a6-ba1e-7efd796bf305"
      },
      "id": "pvLsTaa9dbiY",
      "execution_count": 6,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Mounted at /content/drive\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "81b02926",
      "metadata": {
        "id": "81b02926",
        "outputId": "3e2c4d89-7997-4ca9-8cd3-faa33c038d6f"
      },
      "outputs": [
        {
          "ename": "NameError",
          "evalue": "name 'X_train' is not defined",
          "output_type": "error",
          "traceback": [
            "\u001b[1;31m---------------------------------------------------------------------------\u001b[0m",
            "\u001b[1;31mNameError\u001b[0m                                 Traceback (most recent call last)",
            "Cell \u001b[1;32mIn [118], line 1\u001b[0m\n\u001b[1;32m----> 1\u001b[0m q_aware_model\u001b[38;5;241m.\u001b[39mfit(\u001b[43mX_train\u001b[49m, y_train, epochs\u001b[38;5;241m=\u001b[39m\u001b[38;5;241m1\u001b[39m)\n",
            "\u001b[1;31mNameError\u001b[0m: name 'X_train' is not defined"
          ]
        }
      ],
      "source": [
        "q_aware_model.fit(X_train, y_train, epochs=1)\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "63a710c0",
      "metadata": {
        "id": "63a710c0"
      },
      "outputs": [],
      "source": [
        "q_aware_model.evaluate(X_test, y_test)\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "5c4808cb",
      "metadata": {
        "id": "5c4808cb"
      },
      "outputs": [],
      "source": [
        "converter = tf.lite.TFLiteConverter.from_keras_model(q_aware_model)\n",
        "converter.optimizations = [tf.lite.Optimize.DEFAULT]\n",
        "\n",
        "tflite_qaware_model = converter.convert()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "c5386f02",
      "metadata": {
        "id": "c5386f02"
      },
      "outputs": [],
      "source": [
        "len(tflite_qaware_model)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "a519ff04",
      "metadata": {
        "id": "a519ff04"
      },
      "outputs": [],
      "source": [
        "with open(\"tflite_qaware_model.tflite\", 'wb') as f:\n",
        "    f.write(tflite_qaware_model)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "13770734",
      "metadata": {
        "id": "13770734"
      },
      "outputs": [],
      "source": []
    }
  ],
  "metadata": {
    "accelerator": "TPU",
    "colab": {
      "provenance": [],
      "include_colab_link": true
    },
    "gpuClass": "standard",
    "kernelspec": {
      "display_name": "Python 3 (ipykernel)",
      "language": "python",
      "name": "python3"
    },
    "language_info": {
      "codemirror_mode": {
        "name": "ipython",
        "version": 3
      },
      "file_extension": ".py",
      "mimetype": "text/x-python",
      "name": "python",
      "nbconvert_exporter": "python",
      "pygments_lexer": "ipython3",
      "version": "3.10.7"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 5
}