Distillation_Toy_Example.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "Distillation Toy Example.ipynb",
      "provenance": [],
      "collapsed_sections": [],
      "authorship_tag": "ABX9TyPsr2km9BDyh8GSWlvpA85Y",
      "include_colab_link": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "accelerator": "GPU"
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/sayakpaul/Knowledge-Distillation-in-Keras/blob/master/Distillation_Toy_Example.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "cBctyswACU4c",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "# Imports\n",
        "import tensorflow as tf\n",
        "\n",
        "from tensorflow.keras import models\n",
        "from tensorflow.keras import layers\n",
        "\n",
        "tf.random.set_seed(666)"
      ],
      "execution_count": 1,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "YrhXamQACk6S",
        "colab_type": "code",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 34
        },
        "outputId": "b2cea7cc-962f-44d4-bd25-7b8105504bd8"
      },
      "source": [
        "# Load the FashionMNIST dataset, scale the pixel values\n",
        "(X_train, y_train), (X_test, y_test) = tf.keras.datasets.fashion_mnist.load_data()\n",
        "X_train = X_train/255.\n",
        "X_test = X_test/255.\n",
        "\n",
        "X_train.shape, X_test.shape, y_train.shape, y_test.shape"
      ],
      "execution_count": 2,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "((60000, 28, 28), (10000, 28, 28), (60000,), (10000,))"
            ]
          },
          "metadata": {
            "tags": []
          },
          "execution_count": 2
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "ZYeuzIyPCor2",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "# Change the pixel values to float32 and reshape input data\n",
        "X_train = X_train.astype(\"float32\").reshape(-1, 28, 28, 1)\n",
        "X_test = X_test.astype(\"float32\").reshape(-1, 28, 28, 1)"
      ],
      "execution_count": 3,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "7R-PxhlfCqtu",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "# Define utility function for building a basic shallow Convnet \n",
        "def get_teacher_model():\n",
        "    model = models.Sequential()\n",
        "    model.add(layers.Conv2D(16, (5, 5), activation=\"relu\",\n",
        "        input_shape=(28, 28, 1)))\n",
        "    model.add(layers.MaxPooling2D(pool_size=(2, 2)))\n",
        "    model.add(layers.Conv2D(32, (5, 5), activation=\"relu\"))\n",
        "    model.add(layers.MaxPooling2D(pool_size=(2, 2)))\n",
        "    model.add(layers.Dropout(0.2))\n",
        "    model.add(layers.Flatten())\n",
        "    model.add(layers.Dense(128, activation=\"relu\"))\n",
        "    model.add(layers.Dense(10))\n",
        "    \n",
        "    return model"
      ],
      "execution_count": 4,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "l07x1M5ADDWt",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "# Define loss function and optimizer\n",
        "loss_func = tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True)\n",
        "optimizer = tf.keras.optimizers.Adam()"
      ],
      "execution_count": 5,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "lcBBDW2JDI6y",
        "colab_type": "code",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 374
        },
        "outputId": "edac5d3c-b1fb-4d19-dc0d-1684bed5715a"
      },
      "source": [
        "# Prepare TF dataset\n",
        "train_ds = tf.data.Dataset.from_tensor_slices((X_train, y_train)).shuffle(100).batch(64)\n",
        "test_ds = tf.data.Dataset.from_tensor_slices((X_test, y_test)).batch(64)\n",
        "\n",
        "# Train the teacher model\n",
        "teacher_model = get_teacher_model()\n",
        "teacher_model.compile(loss=loss_func, optimizer=optimizer, metrics=[\"accuracy\"])\n",
        "teacher_model.fit(train_ds,\n",
        "                  validation_data=test_ds,\n",
        "                  epochs=10)"
      ],
      "execution_count": 6,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Epoch 1/10\n",
            "938/938 [==============================] - 3s 3ms/step - loss: 0.5794 - accuracy: 0.7885 - val_loss: 0.4403 - val_accuracy: 0.8405\n",
            "Epoch 2/10\n",
            "938/938 [==============================] - 3s 3ms/step - loss: 0.3885 - accuracy: 0.8584 - val_loss: 0.3942 - val_accuracy: 0.8509\n",
            "Epoch 3/10\n",
            "938/938 [==============================] - 3s 3ms/step - loss: 0.3375 - accuracy: 0.8763 - val_loss: 0.3468 - val_accuracy: 0.8737\n",
            "Epoch 4/10\n",
            "938/938 [==============================] - 3s 3ms/step - loss: 0.3070 - accuracy: 0.8873 - val_loss: 0.3303 - val_accuracy: 0.8798\n",
            "Epoch 5/10\n",
            "938/938 [==============================] - 3s 3ms/step - loss: 0.2877 - accuracy: 0.8945 - val_loss: 0.3120 - val_accuracy: 0.8846\n",
            "Epoch 6/10\n",
            "938/938 [==============================] - 3s 3ms/step - loss: 0.2703 - accuracy: 0.8995 - val_loss: 0.2943 - val_accuracy: 0.8920\n",
            "Epoch 7/10\n",
            "938/938 [==============================] - 3s 3ms/step - loss: 0.2544 - accuracy: 0.9056 - val_loss: 0.2818 - val_accuracy: 0.8960\n",
            "Epoch 8/10\n",
            "938/938 [==============================] - 3s 3ms/step - loss: 0.2427 - accuracy: 0.9098 - val_loss: 0.2795 - val_accuracy: 0.8969\n",
            "Epoch 9/10\n",
            "938/938 [==============================] - 3s 3ms/step - loss: 0.2327 - accuracy: 0.9141 - val_loss: 0.2767 - val_accuracy: 0.8998\n",
            "Epoch 10/10\n",
            "938/938 [==============================] - 3s 3ms/step - loss: 0.2222 - accuracy: 0.9158 - val_loss: 0.2726 - val_accuracy: 0.9020\n"
          ],
          "name": "stdout"
        },
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "<tensorflow.python.keras.callbacks.History at 0x7f93604dc828>"
            ]
          },
          "metadata": {
            "tags": []
          },
          "execution_count": 6
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "OXADzI35Dw3g",
        "colab_type": "code",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 51
        },
        "outputId": "b33dd53b-70ff-496b-b1ea-af4543f2b815"
      },
      "source": [
        "# Evaluate and serialize\n",
        "print(\"Test accuracy: {:.2f}\".format(teacher_model.evaluate(test_ds)[1]*100))\n",
        "teacher_model.save_weights(\"teacher_model.h5\")"
      ],
      "execution_count": 7,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "157/157 [==============================] - 0s 2ms/step - loss: 0.2726 - accuracy: 0.9020\n",
            "Test accuracy: 90.20\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "shnrhMFQFKwZ",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "# Student model utility\n",
        "def get_student_model():\n",
        "    model = models.Sequential()\n",
        "    model.add(layers.Input(shape=(28, 28, 1)))\n",
        "    model.add(layers.Flatten())\n",
        "    model.add(layers.Dense(48, activation=\"relu\"))\n",
        "    model.add(layers.Dense(10))\n",
        "    \n",
        "    return model"
      ],
      "execution_count": 8,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "dPFOtO4mGLIr",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "# Credits: https://github.com/google-research/simclr/blob/master/colabs/distillation_self_training.ipynb\n",
        "def get_kd_loss(student_logits, teacher_logits, temperature=0.5):\n",
        "    teacher_probs = tf.nn.softmax(teacher_logits / temperature)\n",
        "    kd_loss = tf.compat.v1.losses.softmax_cross_entropy(\n",
        "        teacher_probs, student_logits / temperature, temperature**2)\n",
        "    return kd_loss"
      ],
      "execution_count": 9,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "KDZ5DWUeGkK2",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "# Model, optimizer\n",
        "student_model = get_student_model()\n",
        "optimizer = tf.keras.optimizers.Adam(learning_rate=0.01)\n",
        "\n",
        "# Average the loss across the batch size within an epoch\n",
        "train_loss = tf.keras.metrics.Mean(name=\"train_loss\")\n",
        "valid_loss = tf.keras.metrics.Mean(name=\"test_loss\")\n",
        "\n",
        "# Specify the performance metric\n",
        "train_acc = tf.keras.metrics.SparseCategoricalAccuracy(name=\"train_acc\")\n",
        "valid_acc = tf.keras.metrics.SparseCategoricalAccuracy(name=\"valid_acc\")"
      ],
      "execution_count": 16,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "5w1sCCqQGeTe",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "# Train utils\n",
        "@tf.function\n",
        "def model_train(images, labels, teacher_model, \n",
        "                student_model, optimizer, temperature):\n",
        "    teacher_logits = teacher_model(images)\n",
        "\n",
        "    with tf.GradientTape() as tape:\n",
        "        student_logits = student_model(images)\n",
        "        loss = get_kd_loss(student_logits, teacher_logits, temperature)\n",
        "    \n",
        "    gradients = tape.gradient(loss, student_model.trainable_variables)\n",
        "    optimizer.apply_gradients(zip(gradients, student_model.trainable_variables))\n",
        "\n",
        "    train_loss(loss)\n",
        "    train_acc(labels, tf.nn.softmax(student_logits))"
      ],
      "execution_count": 17,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "qXjapT-hHeP1",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "# Validation utils\n",
        "@tf.function\n",
        "def model_validate(images, labels, teacher_model, \n",
        "                   student_model, temperature):\n",
        "    teacher_logits = teacher_model(images)\n",
        "\n",
        "    student_logits = student_model(images)\n",
        "    loss = get_kd_loss(student_logits, teacher_logits, temperature)\n",
        "\n",
        "    valid_loss(loss)\n",
        "    valid_acc(labels, tf.nn.softmax(student_logits))"
      ],
      "execution_count": 18,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "ph4r4J_zHqFE",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "# Tie everything together\n",
        "def train_model(epochs, teacher_model, student_model, optimizer, temperature=0.5):\n",
        "    for epoch in range(epochs):\n",
        "        for (images, labels) in train_ds:\n",
        "            model_train(images, labels, teacher_model, student_model, optimizer, temperature)\n",
        "\n",
        "        for (images, labels) in test_ds:\n",
        "            model_validate(images, labels, teacher_model, student_model, temperature)\n",
        "            \n",
        "        (loss, acc) = train_loss.result(), train_acc.result()\n",
        "        (val_loss, val_acc) = valid_loss.result(), valid_acc.result()\n",
        "        \n",
        "        train_loss.reset_states(), train_acc.reset_states()\n",
        "        valid_loss.reset_states(), valid_acc.reset_states()\n",
        "        \n",
        "        template = \"Epoch {}, loss: {:.3f}, acc: {:.3f}, val_loss: {:.3f}, val_acc: {:.3f}\"\n",
        "        print (template.format(epoch+1,\n",
        "                            loss,\n",
        "                            acc,\n",
        "                            val_loss,\n",
        "                            val_acc))\n",
        "        \n",
        "    \n",
        "    return teacher_model, student_model"
      ],
      "execution_count": 19,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "O-breyl1dwNR",
        "colab_type": "code",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 187
        },
        "outputId": "78143bf4-f890-4d72-bbef-3003b0bcf627"
      },
      "source": [
        "_, student_model = train_model(10, teacher_model, student_model, optimizer)"
      ],
      "execution_count": 20,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Epoch 1, loss: 0.116, acc: 0.816, val_loss: 0.097, val_acc: 0.825\n",
            "Epoch 2, loss: 0.091, acc: 0.848, val_loss: 0.091, val_acc: 0.838\n",
            "Epoch 3, loss: 0.086, acc: 0.853, val_loss: 0.088, val_acc: 0.841\n",
            "Epoch 4, loss: 0.084, acc: 0.857, val_loss: 0.086, val_acc: 0.846\n",
            "Epoch 5, loss: 0.082, acc: 0.858, val_loss: 0.089, val_acc: 0.838\n",
            "Epoch 6, loss: 0.081, acc: 0.861, val_loss: 0.085, val_acc: 0.848\n",
            "Epoch 7, loss: 0.080, acc: 0.862, val_loss: 0.088, val_acc: 0.840\n",
            "Epoch 8, loss: 0.079, acc: 0.863, val_loss: 0.092, val_acc: 0.838\n",
            "Epoch 9, loss: 0.078, acc: 0.864, val_loss: 0.085, val_acc: 0.850\n",
            "Epoch 10, loss: 0.078, acc: 0.864, val_loss: 0.086, val_acc: 0.845\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "H0DHWweqcqIJ",
        "colab_type": "text"
      },
      "source": [
        "This can be further improved with longer training time and more careful hyperparameter tuning. "
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "fmLLgmpybLYi",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "# Serialize\n",
        "student_model.save_weights(\"student_model.h5\")"
      ],
      "execution_count": 21,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "rJeZK9enJct9",
        "colab_type": "code",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 51
        },
        "outputId": "765a834b-ac7d-4818-ed86-b2953c678d17"
      },
      "source": [
        "# Investigate the sizes\n",
        "!ls -lh *.h5"
      ],
      "execution_count": 22,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "-rw-r--r-- 1 root root 163K Aug 31 07:47 student_model.h5\n",
            "-rw-r--r-- 1 root root 335K Aug 31 07:44 teacher_model.h5\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "SNfgaGNncnSt",
        "colab_type": "text"
      },
      "source": [
        "Let's check the total number of trainable params."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "cSHNvya8cYLP",
        "colab_type": "code",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 425
        },
        "outputId": "43bb53d7-5e8f-4d9a-bfdb-77e278d9ff20"
      },
      "source": [
        "teacher_model.summary()"
      ],
      "execution_count": 23,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Model: \"sequential\"\n",
            "_________________________________________________________________\n",
            "Layer (type)                 Output Shape              Param #   \n",
            "=================================================================\n",
            "conv2d (Conv2D)              (None, 24, 24, 16)        416       \n",
            "_________________________________________________________________\n",
            "max_pooling2d (MaxPooling2D) (None, 12, 12, 16)        0         \n",
            "_________________________________________________________________\n",
            "conv2d_1 (Conv2D)            (None, 8, 8, 32)          12832     \n",
            "_________________________________________________________________\n",
            "max_pooling2d_1 (MaxPooling2 (None, 4, 4, 32)          0         \n",
            "_________________________________________________________________\n",
            "dropout (Dropout)            (None, 4, 4, 32)          0         \n",
            "_________________________________________________________________\n",
            "flatten (Flatten)            (None, 512)               0         \n",
            "_________________________________________________________________\n",
            "dense (Dense)                (None, 128)               65664     \n",
            "_________________________________________________________________\n",
            "dense_1 (Dense)              (None, 10)                1290      \n",
            "=================================================================\n",
            "Total params: 80,202\n",
            "Trainable params: 80,202\n",
            "Non-trainable params: 0\n",
            "_________________________________________________________________\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "T0-Y1gpDccZ_",
        "colab_type": "code",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 255
        },
        "outputId": "d7a27cdb-13b4-4736-a1a0-a8a336357b30"
      },
      "source": [
        "student_model.summary()"
      ],
      "execution_count": 24,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Model: \"sequential_2\"\n",
            "_________________________________________________________________\n",
            "Layer (type)                 Output Shape              Param #   \n",
            "=================================================================\n",
            "flatten_2 (Flatten)          (None, 784)               0         \n",
            "_________________________________________________________________\n",
            "dense_4 (Dense)              (None, 48)                37680     \n",
            "_________________________________________________________________\n",
            "dense_5 (Dense)              (None, 10)                490       \n",
            "=================================================================\n",
            "Total params: 38,170\n",
            "Trainable params: 38,170\n",
            "Non-trainable params: 0\n",
            "_________________________________________________________________\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "AzC3KhO_J42N",
        "colab_type": "text"
      },
      "source": [
        "Further size decrease is possible with TFLite. "
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Z8d0R_ypVp8y",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "# Credits: https://www.tensorflow.org/lite/performance/post_training_quant\n",
        "\n",
        "def representative_data_gen():\n",
        "    for input_value in tf.data.Dataset.from_tensor_slices(X_train).batch(1).take(100):\n",
        "        yield [input_value]\n",
        "\n",
        "def convert_to_tflite(model, tflite_file):\n",
        "    converter = tf.lite.TFLiteConverter.from_keras_model(model)\n",
        "    converter.optimizations = [tf.lite.Optimize.DEFAULT]\n",
        "    converter.representative_dataset = representative_data_gen\n",
        "    converter.target_spec.supported_ops = [tf.lite.OpsSet.TFLITE_BUILTINS_INT8]\n",
        "    converter.inference_input_type = tf.int8\n",
        "    converter.inference_output_type = tf.int8\n",
        "    tflite_quant_model = converter.convert()\n",
        "\n",
        "    open(tflite_file, 'wb').write(tflite_quant_model)"
      ],
      "execution_count": 25,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "bZgxSge7Y3hU",
        "colab_type": "code",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 190
        },
        "outputId": "71470926-d629-4bdb-f207-a6ed22a6f599"
      },
      "source": [
        "convert_to_tflite(teacher_model, \"teacher.tflite\")\n",
        "convert_to_tflite(student_model, \"student.tflite\")"
      ],
      "execution_count": 26,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "WARNING:tensorflow:From /usr/local/lib/python3.6/dist-packages/tensorflow/python/training/tracking/tracking.py:111: Model.state_updates (from tensorflow.python.keras.engine.training) is deprecated and will be removed in a future version.\n",
            "Instructions for updating:\n",
            "This property should not be used in TensorFlow 2.0, as updates are applied automatically.\n",
            "WARNING:tensorflow:From /usr/local/lib/python3.6/dist-packages/tensorflow/python/training/tracking/tracking.py:111: Layer.updates (from tensorflow.python.keras.engine.base_layer) is deprecated and will be removed in a future version.\n",
            "Instructions for updating:\n",
            "This property should not be used in TensorFlow 2.0, as updates are applied automatically.\n",
            "INFO:tensorflow:Assets written to: /tmp/tmp5020kbxi/assets\n",
            "INFO:tensorflow:Assets written to: /tmp/tmp2t19bpk6/assets\n"
          ],
          "name": "stdout"
        },
        {
          "output_type": "stream",
          "text": [
            "INFO:tensorflow:Assets written to: /tmp/tmp2t19bpk6/assets\n"
          ],
          "name": "stderr"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "-f8eGqRtZA-w",
        "colab_type": "code",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 51
        },
        "outputId": "2c34b9f9-b22a-4e2f-9385-893c778ed2ea"
      },
      "source": [
        "!ls -lh *.tflite"
      ],
      "execution_count": 27,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "-rw-r--r-- 1 root root 40K Aug 31 07:48 student.tflite\n",
            "-rw-r--r-- 1 root root 85K Aug 31 07:48 teacher.tflite\n"
          ],
          "name": "stdout"
        }
      ]
    }
  ]
}