CIFAR10.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# BinaryNet on CIFAR-10 Dataset Implemtation\n",
    "\n",
    "Implementation of [BinaryNet](https://arxiv.org/abs/1602.02830) on the CIFAR-10 dataset, following the original paper."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import tensorflow as tf\n",
    "import larq as lq\n",
    "from typing import Tuple\n",
    "\n",
    "import numpy as np\n",
    "\n",
    "from sklearn.model_selection import train_test_split\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "import os"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Data\n",
    "\n",
    "We fetch the data using inbuilt Keras datasets to make loading on Colab easier.\n",
    "\n",
    "The CIFAR-10 dataset consists of 60000 32x32 colour images in 10 classes, with 6000 images per class. There are 50000 training images and 10000 \n",
    "test images."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "(x_train, y_train), (x_test, y_test) = tf.keras.datasets.cifar10.load_data()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "x_train, x_val, y_train, y_val = train_test_split(x_train, y_train, test_size=0.1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<matplotlib.image.AxesImage at 0x107511c88>"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "iVBORw0KGgoAAAANSUhEUgAAAP8AAAD8CAYAAAC4nHJkAAAABHNCSVQICAgIfAhkiAAAAAlwSFlzAAALEgAACxIB0t1+/AAAADl0RVh0U29mdHdhcmUAbWF0cGxvdGxpYiB2ZXJzaW9uIDMuMC4zLCBodHRwOi8vbWF0cGxvdGxpYi5vcmcvnQurowAAGqdJREFUeJztnXuMXPV1x7/n3nns04/F6we28QNMgkGAyYpEIk1okkYkigKRmihRhfgDxVEV1EZNKyEiNVTqH0mVh/JHlcopVkibBGheoBQlITQtiiiPhYCxMQEbbGJ7vV6D3+xjZu7pHzNuFvt3zs7Ozt6x/ft+pNXO3nN/93fub+bMnb3fOeeIqoIQEh9Jpx0ghHQGBj8hkcLgJyRSGPyERAqDn5BIYfATEikMfkIihcFPSKQw+AmJlMJcBovITQC+BSAF8K+q+hVv//KCAe1dujJo875nKC341vrx2u3JuYLn+4X8LU/r3M6d51LE8tF+XkTC1+2To/sxcezNpk6u5eAXkRTAPwP4MwD7ADwtIg+p6ovWmN6lK/Ghr/0saMvcE7XOxT5HRWYfz/tKs9rjrA9KgtTxw8Hzwzxnbz2cFXHGeF/xbvXr35aP7qtSa/bxWnyDMv331sOzteiHJPYxi8lUcHuqVXtMsT+4/cG/+njTPs3lY//1AHap6quqOgXgPgA3z+F4hJAcmUvwrwTwh2l/72tsI4ScB8wl+EOfY876TCQim0VkWESGJ4+/OYfpCCHtZC7Bvw/A6ml/rwJw4MydVHWLqg6p6lB5wcAcpiOEtJO5BP/TADaIyDoRKQH4NICH2uMWIWS+afluv6pWReQOAL9EXerbqqo7vDFZLcX48cVBWyWz74bWJOymdyda3Lv2rd7dNt4rdR6+LuHcFk8Mmccb541RZ6184bO9cpm4Co3zfDpuaGacgXP33VUC5uFufyE9Edyein23v2Tc7c9qtvJ01rxN7xlAVR8G8PBcjkEI6Qz8hh8hkcLgJyRSGPyERAqDn5BIYfATEilzuts/68mkioHiwaBt0lY1UCn0BrdrMbwdACp2joibuJGplwATlps8qcxN0LFH+ck7ZhaYfVBPRvPkzSSZvawIAJoZa2UPQcE5L+8q5R3TTuyxx2TirZUzmUOS2KFWSsLyt2ThhB8AKBR7jImal/p45SckUhj8hEQKg5+QSGHwExIpDH5CIiXXu/3jxw9jx6+3hm01+w5rYeHy4Par33+LOWai0Gfaqs4d/cxIIgJg3i9vtcSUR+IkgvilwYzNbS7HNZMn1nRpal9vUi/XpsVkrKwWln3cYm3O2re6HpLYtoJ1TEMxAYBi0SiTNotcK175CYkUBj8hkcLgJyRSGPyERAqDn5BIYfATEim5Sn1ZpYK39h8K2kolOyEhe8uQSY5NmmN6L7JbCIxXbdml6iyJlRCUeOKbW4vPNrZd6nOGtFqJz+30Y8yYOrUEC57E5s3lJekYSVc1Z0VcGdBNuLLHiZMgVUzHw2OcDkalYthLrzThWfs2vysh5EKCwU9IpDD4CYkUBj8hkcLgJyRSGPyERMqcpD4R2QPgBIAagKqqDnn7J0kBfb3hemVlmTDHvVUN1zI7eXjMHHPx0sttPxxxa8qpgabGOD+rz8n0cka1mllmzWa2rZqJVlphOaSFVrP6HBw3MqMuYM2pu1h1zit1JDsvY9HrsJYmxeD2xPGjXDDGzCKtrx06/5+q6uE2HIcQkiP82E9IpMw1+BXAr0TkGRHZ3A6HCCH5MNeP/Teo6gERWQrgERF5SVUfm75D401hMwCUusNthQkh+TOnK7+qHmj8PgTgpwCuD+yzRVWHVHWoUO6ey3SEkDbScvCLSK+I9J9+DODDALa3yzFCyPwyl4/9ywD8tCE7FQD8QFV/4Q1ICgm6l4RbbE0dDrfxAgCk5eDmQwft95p1V19h2jKr1REARVhC8Ujh9AZzWlB5kl3iaEOenGNl02VOMcj5wMr4S1NbSi06UpmXreZlF9aqVgFPZ4zbvsz235Nn4bTrKqIrPCSzX4vWWs0mQ7Pl4FfVVwFc0+p4QkhnodRHSKQw+AmJFAY/IZHC4CckUhj8hERKrgU8C6USBlatCdr2H3zJHNfdvzC4/cQbB+zJpo7ZfiweMG3FzCkwmVWD2xPY8o9b1NGV8xxbC1mEWebIg14fP0+9cmQvyw9P3kw9m+OI26vPOKYnfGbeNdF7Qr0io242YPiYxaodnqWw+g1nmrP3bX5XQsiFBIOfkEhh8BMSKQx+QiKFwU9IpOR7t79YxuCqdUHbvqftJIbESOyZOnnCHHNyzFYCBpdfZtomwzf0AQBpzbiD7d8SN03qvPdmzjG9llf23f4Wawl67alaaPTlJr84KoZXd9FLnvKEDIvM6//l+OGdWi2ZvSNFZ0jJCJdZlPDjlZ+QWGHwExIpDH5CIoXBT0ikMPgJiRQGPyGRkntiz9JVlwRtaXefOa5mSC8Fp3be4dftRKGVGzaatp5SuJ4aAJSM1XJFnKKRgQGg6shGE47k6EtsYZuX/OIdrejIik7HK/OYXi6QullQtskzWmctngbonJeXOOO5byUYAXbdRaezGYql8HZXST1z3+Z3JYRcSDD4CYkUBj8hkcLgJyRSGPyERAqDn5BImVHqE5GtAD4G4JCqXtXYNgDgfgBrAewB8ClVPTLjZIUiLhpcFrQNLF1hjhsZCbfy6u8Nt/4CgNG9r5i2HY//2rQtWLjItPX3hWXA7l5bpiz2OLbeBaZt8aKlpm2y5r1nG1pPK+ltAIpepp2b/WaMcVyveeXxZp/ICACwkhmdREAUCrbR6TbmSrA1Z8LMyBZNnfUtWzX82iz1fRfATWdsuxPAo6q6AcCjjb8JIecRMwa/qj4G4M0zNt8M4N7G43sB3NJmvwgh80yr//MvU9URAGj8tj+jEkLOSeb9hp+IbBaRYREZPnX08HxPRwhpklaDf1REVgBA4/cha0dV3aKqQ6o61LtoSYvTEULaTavB/xCA2xqPbwPwYHvcIYTkRTNS3w8B3AhgiYjsA/BlAF8B8ICI3A7gdQCfbGayJE3Qa8hzl6xda457fffu4Pb+BbZUVqlMmbYXn3nctNUyO6WrVAzrPP0LbT96jFZjALB4yaBp+8jH7XuolZJ9TMt/dQp4eipgwZGvfFUpfNDE0fqcTmnwVMVWVExP6ksd7dBvv2bbak4KpJXNKI6TVouv2RTwnDH4VfUzhumDzU9DCDnX4Df8CIkUBj8hkcLgJyRSGPyERAqDn5BIybWAZyqCnnJYLlu/do057mkJV7M0kqHqc5W6Tdsiuy0gMkfqM9PHpmxZ8cj+103ba7970rQNXR4udAoAy9+xybSJkXaWFu33+SSxXwaJk8ZmZaPVHQlv9rLiPMXOUSpdrS8zbF7PwFQcJx0vvX6IlZr9urKeM68loxivYRbwJITMCIOfkEhh8BMSKQx+QiKFwU9IpDD4CYmUXKW+RIA+Q+q7bP1ac9zi/nDhzKNTdq++tNvuudflKDkFr9+aodZ4Y6rV46btyJ5dpu0X9/27aVsz9KppW7g4XIC0u9uWPpctDxdVBYB3XnmdactcSSyM2+vOUw69Ip2OOmv2KHTS3zyJTZ1cRs+PqpOymFjHtF/eQE94czqLyzmv/IRECoOfkEhh8BMSKQx+QiKFwU9IpOSb2JMAC0thW9fFy81xK5aHa92N7DlqjimX+k1br9OOKauFk4gAYMpokyVONkWvU8Nv9frLTdvBsTP7pPyR0SeeMm09feHz7jbqDwLA+jV2q7SrLl1n2noX29WYK8Zd9sypZVdykmbcVBvn9nzN6PPlp295Vhs3J8zxMTGUh6xqH7BYDo8xSvuF521+V0LIhQSDn5BIYfATEikMfkIihcFPSKQw+AmJlGbadW0F8DEAh1T1qsa2uwF8FsBYY7e7VPXhmY6VCrDAmLHoSHPv2HBZcPszO//L9nvRxaYtqU6atoKT8FErhVuNdXfZRQELztvrYNnIzgDg1YqbcnpXqWHrc+oWHhsZMW1Hdu80bZe///2m7Xg1nJVSsZVUlMWWtlJHTs0ceatiLKNbE7BFvBp+3nQFKzPMWaseQzJ3SjWeRTO7fhfATYHt31TVaxs/MwY+IeTcYsbgV9XHANjfOCGEnJfM5X/+O0Rkm4hsFZHFbfOIEJILrQb/twFcCuBaACMAvm7tKCKbRWRYRIaPHB6zdiOE5ExLwa+qo6paU9UMwHcAXO/su0VVh1R1yOtHTwjJl5aCX0SmZ4J8AsD29rhDCMmLZqS+HwK4EcASEdkH4MsAbhSRa1FXMPYA+Fwzk6UCLDIkisRJ27rqneHst8LPfmGO0ZrdQmtsbL9pW7rUrme3aCCcoafVCXPM5IRtE6dNVlqwbTJla0BpMazpFZxMxn27d5i2x3/1S9P27k1Xm7YFvWFZVJ3sQi9zT5xCiV6pu6oji5pzefX9nHE1R+rzZMDUeG7ESdHrMRZrNlUVZwx+Vf1MYPM9s5iDEHIOwm/4ERIpDH5CIoXBT0ikMPgJiRQGPyGRkm8BTwH6newyi3euvyS4faC3bI45Ojlu2ipTFdM2NWVn/JVr4XEnT9gtucol+4THx20fS2X73LqcbMCqIUZNVmzJcXT0sGn75b69pq1/of2t7uVG4c/L1oczNAHgXddcadoSJ6uv6mXTGZc3X85rLeXPyy50unWZncMksf3oMp7n2VzNeeUnJFIY/IRECoOfkEhh8BMSKQx+QiKFwU9IpOQq9SUAeo23G83s3Kz+FeFMu8tW2/39hvfbMtryi+3edNVJWxI7+ma4GElX2ZbzalVbVjxx/JhpGxy0++B5a6USzutSJ4Nw1aXvMG2nDtrFPZ94YbdpK+4dDW5f98oBc8yGtWtM27JBu+dh6vTWS60eeeLIeS0W91TvmG5y4ewntITg2eQw8spPSKQw+AmJFAY/IZHC4CckUhj8hERK7nf7uwxb5rRqKnaFE1muvmKDOeap3Y+btnRhn2l7y7nbv/Ci8D3WbicJ58ABu17gooV2izIvuaRasesTJl3Wudnv88suWW/adHk4QQcA4Jz3qUpYbXnxD3YS0c5ddhLRiiV2vUAvV6w4+xJ+s7tl3p6BQTLnNWAFLu/2E0JmhMFPSKQw+AmJFAY/IZHC4CckUhj8hERKM+26VgP4HoDlADIAW1T1WyIyAOB+AGtRb9n1KVU9MtPxrHZCvkQRHrXp6qvMEaX//K1pO3bMTqiBkzTTbdTj8xJ0Con9/qpVu+2WOuPKvXYNv1oaXqvE6YdWzWzpECWjv5ozFwAkRqup8XG7RuLz23eatvdu2mjaykX7ZZyY7bq8ZJr2SnYto60mCjVHM1f+KoAvquoVAN4D4PMishHAnQAeVdUNAB5t/E0IOU+YMfhVdURVn208PgFgJ4CVAG4GcG9jt3sB3DJfThJC2s+s/ucXkbUANgF4EsAyVR0B6m8QAJa22zlCyPzRdPCLSB+AHwP4gqraherPHrdZRIZFZHhsLFwMgxCSP00Fv4gUUQ/876vqTxqbR0VkRcO+AsCh0FhV3aKqQ6o6NDg42A6fCSFtYMbgl3prk3sA7FTVb0wzPQTgtsbj2wA82H73CCHzRTNZfTcAuBXACyLyXGPbXQC+AuABEbkdwOsAPjkXR1pRLi6/7FLTtvQiu+bby3sOmrZSassrY4fCdek8RabXkeWmJm3Z6w1jLgBYuqbXtKVp+ClVZ4W9dlep2HJkltk2ScK2rqKdvfnyrl2mbfSwrSKvXRWu8QjAUfScV1yrMlqLtf/M6az6g21ixuBX1d/C9u+D7XWHEJIX/IYfIZHC4CckUhj8hEQKg5+QSGHwExIpuRbw9Jn9+9CSQfsbxetW2O2uXttrF9UslO2iml29C4Lbu/vsMerogGm3nUGYpXY23aSTDdhrtIyaqtpzwcx8A9LEluYSp+gqNGwrd9lFP8eO2tmRTzy/w7SV+xzp01h/r0BqktqvxdTJZHRxZEBLXXYSO9HbG17H2aiNvPITEikMfkIihcFPSKQw+AmJFAY/IZHC4CckUs4Lqc9Sy9K02xyzcf1a0/bc9pdNW9eAXXNg4YKw1NezcJE5ZmqqYtqyzJbKyl1WV0PAqcWJUiks2xWcVLVM7JdBltjikYjtiBpSX5bYc0m3vR4//58nTNv/Pv+SaStbWY5qy6UFJyo8GbDs9C705NSCho9ZgP3a+ZN3XxPcfurUW7YPZ8ArPyGRwuAnJFIY/IRECoOfkEhh8BMSKefQ3f72cuXGK0yb3v8j03bg94dN28je14LbFzgJRosW2rUEL3KqGfeVw63BAADOHWezc1XBPl5m3G0GgCnn1vekkyxUq4Vt1nYAEGeu8Sm7pdi+MbuSfNGog3f8pD2mq8e+a1902pelXms2R20xVZOJE+aYtyYngtuPnTxpjjkTXvkJiRQGPyGRwuAnJFIY/IRECoOfkEhh8BMSKTNKfSKyGsD3ACwHkAHYoqrfEpG7AXwWwOnWu3ep6sPz4WS9XeDsuOrKK03bl/7ub0zbgZER07Z3f7jN15snTpljDjudiQ++9DvTVqnaSR3HnOQNNRJnFvQvNsf09IQTlgCg5EiVacmWxMSQvSTx2mTZiUJeucDUk9EMabG7ZEuf5W67xVqXk7zj1Wv0srHUqAtY7rZ9HBkJv64qFTth6Uya0fmrAL6oqs+KSD+AZ0TkkYbtm6r6taZnI4ScMzTTq28EwEjj8QkR2Qlg5Xw7RgiZX2b1P7+IrAWwCcCTjU13iMg2EdkqIvbnSkLIOUfTwS8ifQB+DOALqnocwLcBXArgWtQ/GXzdGLdZRIZFZHjM+f+XEJIvTQW/iBRRD/zvq+pPAEBVR1W1pvWSLd8BcH1orKpuUdUhVR0adL7LTgjJlxmDX+q32u8BsFNVvzFt+4ppu30CwPb2u0cImS+audt/A4BbAbwgIs81tt0F4DMici3qHYL2APjcvHjo4Ekri526eh/6wI2mrVKZNG0TE+FMqklHXjnpZFkdddpTnXTkvNf223LkgZHR4PbKpO1j6shQ1cyWHMcnx02bVbvw1Lg9ZtzJEnzjjSOmrTJuP2epkdWXOarccWc9vHZdniDtKJWoGZfgolEHEQAKU+F1rEzY63vWMWbaQVV/i/B5zYumTwjJB37Dj5BIYfATEikMfkIihcFPSKQw+AmJlPO6gKeX7Zc4xRRLRbsIY8msgAn09RjZXo7Go7CLe3pJYDWnlde7Jmxpq1IJy2WZraLVxVqDRG3JMcts+dA6ZKViF+KcMHwHgPEJW3LUqu2HGE9OltlPWtVZD3U0QmcY1LFWDEkvUXs9uo3syL+4dYfjxRnHb3pPQsgFBYOfkEhh8BMSKQx+QiKFwU9IpDD4CYmUc0bqa6VIp3s8T39zCkX674dhucaTcTy8Uy44/fj6ezwfjXNTb31tWyp9zjhPP7SmsueqOdqnt8KJYzRfB856qJEJOF9kxtmJk9WXFsI+9lhydABe+QmJFAY/IZHC4CckUhj8hEQKg5+QSGHwExIp54zU5/Y5M+Qh8aSh1tQ3XxEzZCNXVnTxxtknIC1Ii956eEeriv0SUU8yNQ7qZjK6Up9t83r1JYbNaxnY8tPZ4iET40WXOQ0KM0vSnYXzvPITEikMfkIihcFPSKQw+AmJFAY/IZEy491+EekC8BiAcmP/H6nql0VkHYD7AAwAeBbArapqF2ibAf8uu+lcK6a209o9+xnGueKH0zLKSkpp8W0+a/XWdwvDPBe9dlfeuBxfBi1j+ugEResK0x9p5iUxCeADqnoN6u24bxKR9wD4KoBvquoGAEcA3D5nbwghuTFj8Gud090mi40fBfABAD9qbL8XwC3z4iEhZF5o6sOgiKSNDr2HADwCYDeAo6p6umbyPgAr58dFQsh80FTwq2pNVa8FsArA9QCuCO0WGisim0VkWESGx8bGWveUENJWZnUbSFWPAvhvAO8BsEjk/7/7uQrAAWPMFlUdUtWhwcHBufhKCGkjMwa/iAyKyKLG424AHwKwE8BvAPx5Y7fbADw4X04SQtpPM4k9KwDcK3V9KQHwgKr+XEReBHCfiPwjgN8BuKeZCa2qb5nTBsmS7bzEHl9GazHr57yghbp6LeKuorXGbkKNfS3ypFu/TdZ5jJeM1QYte8bgV9VtADYFtr+K+v//hJDzEH7Dj5BIYfATEikMfkIihcFPSKQw+AmJFMlT9hKRMQB7G38uAXA4t8lt6MfboR9v53zzY42qNvVtulyD/20Tiwyr6lBHJqcf9IN+8GM/IbHC4CckUjoZ/Fs6OPd06MfboR9v54L1o2P/8xNCOgs/9hMSKR0JfhG5SUR+LyK7ROTOTvjQ8GOPiLwgIs+JyHCO824VkUMisn3atgEReUREXmn8XtwhP+4Wkf2NNXlORD6agx+rReQ3IrJTRHaIyF83tue6Jo4fua6JiHSJyFMi8nzDj39obF8nIk821uN+ESnNaSJVzfUHQIp6GbD1AEoAngewMW8/Gr7sAbCkA/O+D8B1ALZP2/ZPAO5sPL4TwFc75MfdAP425/VYAeC6xuN+AC8D2Jj3mjh+5LomqCc+9zUeFwE8iXoBnQcAfLqx/V8A/OVc5unElf96ALtU9VWtl/q+D8DNHfCjY6jqYwDePGPzzagXQgVyKohq+JE7qjqiqs82Hp9AvVjMSuS8Jo4fuaJ15r1obieCfyWAP0z7u5PFPxXAr0TkGRHZ3CEfTrNMVUeA+osQwNIO+nKHiGxr/Fsw7/9+TEdE1qJeP+JJdHBNzvADyHlN8iia24ngD5Ug6ZTkcIOqXgfgIwA+LyLv65Af5xLfBnAp6j0aRgB8Pa+JRaQPwI8BfEFVj+c1bxN+5L4mOoeiuc3SieDfB2D1tL/N4p/zjaoeaPw+BOCn6GxlolERWQEAjd+HOuGEqo42XngZgO8gpzURkSLqAfd9Vf1JY3PuaxLyo1Nr0ph71kVzm6UTwf80gA2NO5clAJ8G8FDeTohIr4j0n34M4MMAtvuj5pWHUC+ECnSwIOrpYGvwCeSwJiIiqNeA3Kmq35hmynVNLD/yXpPciubmdQfzjLuZH0X9TupuAF/qkA/rUVcangewI08/APwQ9Y+PFdQ/Cd0O4CIAjwJ4pfF7oEN+/BuAFwBsQz34VuTgx3tR/wi7DcBzjZ+P5r0mjh+5rgmAq1EvirsN9Teav5/2mn0KwC4A/wGgPJd5+A0/QiKF3/AjJFIY/IRECoOfkEhh8BMSKQx+QiKFwU9IpDD4CYkUBj8hkfJ/jmkuLAWqC6AAAAAASUVORK5CYII=\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "plt.imshow(x_train[1])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Using the tf.data.Dataset API\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "def preprocess(image: tf.Tensor, labels: tf.Tensor, training: bool) -> Tuple[tf.Tensor, tf.Tensor]:\n",
    "    \"\"\"\n",
    "    Applies preprocessing to images in the ETL pipeline.\n",
    "    :param image: The input image\n",
    "    :param labels: The input images label\n",
    "    :param training: Boolean, True if training, False otherwise\n",
    "    :returns: The processed image and corresponding label\n",
    "    \"\"\"\n",
    "    image = tf.image.per_image_standardization(image)\n",
    "    \n",
    "    # Some transforms we only apply to train data, such as random augmentations to improve learning generalisation\n",
    "    if training:\n",
    "        image = tf.image.resize_with_crop_or_pad(image, 40, 40)\n",
    "        image = tf.image.random_crop(image, (32, 32, 3))\n",
    "        image = tf.image.random_flip_left_right(image)\n",
    "        image = tf.image.random_flip_up_down(image)\n",
    "    return image, labels"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "def dataset_generator(data_x: np.ndarray, data_y: np.ndarray, batch_size: int, training: bool) -> tf.data.Dataset:\n",
    "    \"\"\"\n",
    "    Produces a tf.data.Dataset object for a highly optimised data ETL pipeline.\n",
    "    :param data_x: The input image array\n",
    "    :param data_y: The corresponding labels\n",
    "    :param batch_size: The batch size\n",
    "    :param training: Boolean, True if training, false otherwise\n",
    "    :return: A tf.data.Dataset object\n",
    "    \"\"\"\n",
    "    labels = tf.one_hot(np.squeeze(data_y), 10)  # Convert the labels to a sparse matrix\n",
    "    images = tf.cast(data_x, dtype=tf.float16)\n",
    "    \n",
    "    dataset = tf.data.Dataset.from_tensor_slices((images, labels))  # Convert to tf.data.Dataset\n",
    "    \n",
    "    dataset = dataset.shuffle(data_x.shape[0], reshuffle_each_iteration=True)  # Shuffle the data each epoch\n",
    "    \n",
    "    dataset = dataset.map(lambda x, y: preprocess(image=x, labels=y, training=training))  # Apply preprocessing\n",
    "    \n",
    "    dataset = dataset.batch(batch_size)\n",
    "    dataset = dataset.prefetch(2)\n",
    "    \n",
    "    return dataset    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_batch_size = 64\n",
    "val_batch_size = 128\n",
    "test_batch_size = 128"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "WARNING: Logging before flag parsing goes to stderr.\n",
      "W0710 23:11:38.317013 4513134016 deprecation.py:323] From /anaconda3/envs/TensorFlow-2.0/lib/python3.7/site-packages/tensorflow/python/ops/image_ops_impl.py:1511: div (from tensorflow.python.ops.math_ops) is deprecated and will be removed in a future version.\n",
      "Instructions for updating:\n",
      "Deprecated in favor of operator or tf.math.divide.\n"
     ]
    }
   ],
   "source": [
    "train_dataset = dataset_generator(x_train, y_train, batch_size=train_batch_size, training=True)\n",
    "val_dataset = dataset_generator(x_val, y_val, batch_size=val_batch_size, training=False)\n",
    "test_dataset = dataset_generator(x_train, y_train, batch_size=test_batch_size, training=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<PrefetchDataset shapes: ((None, 32, 32, 3), (None, 10)), types: (tf.float16, tf.float32)>"
      ]
     },
     "execution_count": 15,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "train_dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<PrefetchDataset shapes: ((None, 32, 32, 3), (None, 10)), types: (tf.float16, tf.float32)>"
      ]
     },
     "execution_count": 16,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "val_dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<PrefetchDataset shapes: ((None, 32, 32, 3), (None, 10)), types: (tf.float16, tf.float32)>"
      ]
     },
     "execution_count": 17,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "test_dataset"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Building BinaryNet"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [],
   "source": [
    "class BinaryNet(tf.keras.Model):\n",
    "\n",
    "    def __init__(self, num_classes: int) -> None:\n",
    "        \"\"\"\n",
    "        Builds the BinaryNet model in TensorFlow 2.0 and Larq.\n",
    "        :param num_classes: The number of output classes\n",
    "        \"\"\"\n",
    "        super(BinaryNet, self).__init__(name='BinaryNet')\n",
    "\n",
    "        kwargs = {\n",
    "            'input_quantizer': 'ste_sign',\n",
    "            'kernel_quantizer': 'ste_sign',\n",
    "            'kernel_constraint': 'weight_clip',\n",
    "            'use_bias': False\n",
    "        }\n",
    "        # Conv Block 1\n",
    "        self.conv_1_1 = lq.layers.QuantConv2D(128, 3,\n",
    "                                              padding='same',\n",
    "                                              kernel_quantizer='ste_sign',\n",
    "                                              kernel_constraint='weight_clip',\n",
    "                                              use_bias=False)\n",
    "        self.bn_1_1 = tf.keras.layers.BatchNormalization(momentum=0.999, scale=False)\n",
    "        self.conv_1_2 = lq.layers.QuantConv2D(128, 3, padding='same', **kwargs)\n",
    "        self.max_pool_1 = tf.keras.layers.MaxPool2D(pool_size=(2, 2), strides=(2, 2))\n",
    "        self.bn_1_2 = tf.keras.layers.BatchNormalization(momentum=0.999, scale=False)\n",
    "\n",
    "        # Conv Block 2\n",
    "        self.conv_2_1 = lq.layers.QuantConv2D(256, 3, padding='same', **kwargs)\n",
    "        self.bn_2_1 = tf.keras.layers.BatchNormalization(momentum=0.999, scale=False)\n",
    "        self.conv_2_2 = lq.layers.QuantConv2D(256, 3, padding='same', **kwargs)\n",
    "        self.max_pool_2 = tf.keras.layers.MaxPool2D(pool_size=(2, 2), strides=(2, 2))\n",
    "        self.bn_2_2 = tf.keras.layers.BatchNormalization(momentum=0.999, scale=False)\n",
    "\n",
    "        # Conv Block 3\n",
    "        self.conv_3_1 = lq.layers.QuantConv2D(512, 3, padding='same', **kwargs)\n",
    "        self.bn_3_1 = tf.keras.layers.BatchNormalization(momentum=0.999, scale=False)\n",
    "        self.conv_3_2 = lq.layers.QuantConv2D(512, 3, padding='same', **kwargs)\n",
    "        self.max_pool_3 = tf.keras.layers.MaxPool2D(pool_size=(2, 2), strides=(2, 2))\n",
    "        self.bn_3_2 = tf.keras.layers.BatchNormalization(momentum=0.999, scale=False)\n",
    "\n",
    "        # Flatten\n",
    "        self.flatten_4 = tf.keras.layers.Flatten()\n",
    "\n",
    "        # FC Block\n",
    "        self.fc_5 = lq.layers.QuantDense(1024, **kwargs)\n",
    "        self.bn_5 = tf.keras.layers.BatchNormalization(momentum=0.999, scale=False)\n",
    "\n",
    "        # FC Block\n",
    "        self.fc_6 = lq.layers.QuantDense(1024, **kwargs)\n",
    "        self.bn_6 = tf.keras.layers.BatchNormalization(momentum=0.999, scale=False)\n",
    "\n",
    "        # FC Block\n",
    "        self.fc_7 = lq.layers.QuantDense(num_classes, **kwargs)\n",
    "        self.bn_7 = tf.keras.layers.BatchNormalization(momentum=0.999, scale=False)\n",
    "\n",
    "        # Output\n",
    "        self.network_output = tf.keras.layers.Softmax()\n",
    "\n",
    "    def call(self, inputs: tf.Tensor) -> tf.keras.Model:\n",
    "        \"\"\"\n",
    "        Builds the BinaryNet model when called.\n",
    "        :param inputs: The input tensor\n",
    "        :return: The built BinaryNet model\n",
    "        \"\"\"\n",
    "\n",
    "        x = self.conv_1_1(inputs)\n",
    "        x = self.bn_1_1(x)\n",
    "        x = self.conv_1_2(x)\n",
    "        x = self.max_pool_1(x)\n",
    "        x = self.bn_1_2(x)\n",
    "\n",
    "        x = self.conv_2_1(x)\n",
    "        x = self.bn_2_1(x)\n",
    "        x = self.conv_2_2(x)\n",
    "        x = self.max_pool_2(x)\n",
    "        x = self.bn_2_2(x)\n",
    "\n",
    "        x = self.conv_3_1(x)\n",
    "        x = self.bn_3_1(x)\n",
    "        x = self.conv_3_2(x)\n",
    "        x = self.max_pool_3(x)\n",
    "        x = self.bn_3_2(x)\n",
    "\n",
    "        x = self.flatten_4(x)\n",
    "\n",
    "        x = self.fc_5(x)\n",
    "        x = self.bn_5(x)\n",
    "\n",
    "        x = self.fc_6(x)\n",
    "        x = self.bn_6(x)\n",
    "\n",
    "        x = self.fc_7(x)\n",
    "        x = self.bn_7(x)\n",
    "\n",
    "        x = self.network_output(x)\n",
    "\n",
    "        return x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [],
   "source": [
    "model = BinaryNet(10)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [],
   "source": [
    "model.build((None,32,32,3))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Model: \"BinaryNet\"\n",
      "_________________________________________________________________\n",
      "Layer (type)                 Output Shape              Param #   \n",
      "=================================================================\n",
      "quant_conv2d (QuantConv2D)   multiple                  6914      \n",
      "_________________________________________________________________\n",
      "batch_normalization (BatchNo multiple                  384       \n",
      "_________________________________________________________________\n",
      "quant_conv2d_1 (QuantConv2D) multiple                  294914    \n",
      "_________________________________________________________________\n",
      "max_pooling2d (MaxPooling2D) multiple                  0         \n",
      "_________________________________________________________________\n",
      "batch_normalization_1 (Batch multiple                  384       \n",
      "_________________________________________________________________\n",
      "quant_conv2d_2 (QuantConv2D) multiple                  589826    \n",
      "_________________________________________________________________\n",
      "batch_normalization_2 (Batch multiple                  768       \n",
      "_________________________________________________________________\n",
      "quant_conv2d_3 (QuantConv2D) multiple                  1179650   \n",
      "_________________________________________________________________\n",
      "max_pooling2d_1 (MaxPooling2 multiple                  0         \n",
      "_________________________________________________________________\n",
      "batch_normalization_3 (Batch multiple                  768       \n",
      "_________________________________________________________________\n",
      "quant_conv2d_4 (QuantConv2D) multiple                  2359298   \n",
      "_________________________________________________________________\n",
      "batch_normalization_4 (Batch multiple                  1536      \n",
      "_________________________________________________________________\n",
      "quant_conv2d_5 (QuantConv2D) multiple                  4718594   \n",
      "_________________________________________________________________\n",
      "max_pooling2d_2 (MaxPooling2 multiple                  0         \n",
      "_________________________________________________________________\n",
      "batch_normalization_5 (Batch multiple                  1536      \n",
      "_________________________________________________________________\n",
      "flatten (Flatten)            multiple                  0         \n",
      "_________________________________________________________________\n",
      "quant_dense (QuantDense)     multiple                  16777218  \n",
      "_________________________________________________________________\n",
      "batch_normalization_6 (Batch multiple                  3072      \n",
      "_________________________________________________________________\n",
      "quant_dense_1 (QuantDense)   multiple                  2097154   \n",
      "_________________________________________________________________\n",
      "batch_normalization_7 (Batch multiple                  3072      \n",
      "_________________________________________________________________\n",
      "quant_dense_2 (QuantDense)   multiple                  20482     \n",
      "_________________________________________________________________\n",
      "batch_normalization_8 (Batch multiple                  30        \n",
      "_________________________________________________________________\n",
      "softmax (Softmax)            multiple                  0         \n",
      "=================================================================\n",
      "Total params: 28,055,600\n",
      "Trainable params: 14,025,866\n",
      "Non-trainable params: 14,029,734\n",
      "_________________________________________________________________\n"
     ]
    }
   ],
   "source": [
    "model.summary()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "+BinaryNet stats------------------------------------------------------------+\n",
      "| Layer                  Input prec.   Outputs   # 1-bit  # 32-bit   Memory |\n",
      "|                              (bit)                                   (kB) |\n",
      "+---------------------------------------------------------------------------+\n",
      "| quant_conv2d                     -  multiple      3456         0     0.42 |\n",
      "| batch_normalization              -  multiple         0       384     1.50 |\n",
      "| quant_conv2d_1                   1  multiple    147456         0    18.00 |\n",
      "| max_pooling2d                    -  multiple         0         0     0.00 |\n",
      "| batch_normalization_1            -  multiple         0       384     1.50 |\n",
      "| quant_conv2d_2                   1  multiple    294912         0    36.00 |\n",
      "| batch_normalization_2            -  multiple         0       768     3.00 |\n",
      "| quant_conv2d_3                   1  multiple    589824         0    72.00 |\n",
      "| max_pooling2d_1                  -  multiple         0         0     0.00 |\n",
      "| batch_normalization_3            -  multiple         0       768     3.00 |\n",
      "| quant_conv2d_4                   1  multiple   1179648         0   144.00 |\n",
      "| batch_normalization_4            -  multiple         0      1536     6.00 |\n",
      "| quant_conv2d_5                   1  multiple   2359296         0   288.00 |\n",
      "| max_pooling2d_2                  -  multiple         0         0     0.00 |\n",
      "| batch_normalization_5            -  multiple         0      1536     6.00 |\n",
      "| flatten                          -  multiple         0         0     0.00 |\n",
      "| quant_dense                      1  multiple   8388608         0  1024.00 |\n",
      "| batch_normalization_6            -  multiple         0      3072    12.00 |\n",
      "| quant_dense_1                    1  multiple   1048576         0   128.00 |\n",
      "| batch_normalization_7            -  multiple         0      3072    12.00 |\n",
      "| quant_dense_2                    1  multiple     10240         0     1.25 |\n",
      "| batch_normalization_8            -  multiple         0        30     0.12 |\n",
      "| softmax                          -  multiple         0         0     0.00 |\n",
      "+---------------------------------------------------------------------------+\n",
      "| Total                                         14022016     11550  1756.79 |\n",
      "+---------------------------------------------------------------------------+\n",
      "+BinaryNet summary----------------+\n",
      "| Total params           14033566 |\n",
      "| Trainable params       14025866 |\n",
      "| Non-trainable params   7700     |\n",
      "| Float-32 Equivalent    53.53 MB |\n",
      "| Compression of Memory  31.20    |\n",
      "+---------------------------------+\n"
     ]
    }
   ],
   "source": [
    "lq.models.summary(model)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "iVBORw0KGgoAAAANSUhEUgAAAG4AAAA8CAYAAACHHY8HAAAAAXNSR0IArs4c6QAACKFJREFUeAHtXGeIFEsQrjvPnDPmnAMqxh8i5gCKCkYUFAMGUBATZhFEMUeMKAYUA+YEBu7MCX+omBUDRjDn1K++eq+HCXt7c3f7vBnshr3prq7uqamvq7pqtvfiFBcyJWwaSIoPm8RG3n81YIAL6UowwBngQqqBkIqd4Jb73LlzNH/+fDfZtDNQA6tXr6a8efM6JPC4ysePH9P27dsdTKaRMRp49eqVYPH161ePAB6L0xzbtm3TVXPNIA0kJSVRs2bNIt7dY3ERuQwxcBowwAUOEn8CGeD86SlwXAa4wEHiTyADnD89BY7LABc4SPwJZIDzp6fAcRngAgeJP4EMcP70FDguA1zgIPEnkAHOn54Cx2WACxwk/gQywPnTU+C4DHCBg8SfQAY4f3oKHFey38f5lRTflj98+NDBHhcXR7ly5aLGjRvL90m5c+d29O/du5dGjBhBq1atojZt2jj6Mqpx9uxZ2rp1q3V7yFehQgWrjcq8efPo0aNHFq1kyZI0ZswYq/0nK+m2uH79+lHp0qVp8eLF8smTJ4+037x5QwMGDKACBQrQ6NGjHc/04MEDARvXoJRGjRpR9+7dBTw8S+fOnenTp08O8fr06UNZs2aV5/zx4wf17NnT0Z9SI6ZHWHEg1l74m28ckLWTUqy/f/9exmAcW5/FjzoDJ32HDx+26KjY+RwdGdwYNmyY9SzdunXzSPP8+XPpf/r0qacvGuHo0aNqypQp0Vg8fYmJiXIv3NNVEtNtcVhlcIVwj+4CS2zatKmQL1y44OhGX1BLvXr1RDScvZk1a5ZDzEKFCkk7f/78Dnq0BraS3r17069fv6Kxpaov3XtcSnfj1SIsRYoUsVhv3rxJmzdvFpc6aNAgoV++fJl27twpe2OHDh1o5cqVdOvWLerfvz/17dvXGnvmzBnatWsXYSFgHx0yZAh17NhR+n/+/Els2bR27VpauHAhLVmyhK5du0Z16tShd+/eCU9CQgKNHz+eSpQoQcuWLZP+bNmy0YIFC6x7QMktW7akOXPm0MSJE2V8u3btpF8vUH3Vg3bv3k34vHjxggD82LFj5WTWvXv3qHnz5vTy5UvpZ0uloUOHUoMGDfTQtF1dJqjS4ioxBz+ImPWdO3fUhw8f1KVLlxQLKLS6desq3hPkVgyYqlWrltBZKUK7fv26atKkidCKFy+uatasqRgslS9fPhUfH68YJOE7deqUtNevX6++ffumhg8frjJlyqQYYOlnhcscrAnFylNlypSRNviLFSsm9QMHDggv/mAOBlKdPHnSosFVzp07V/EiUG3btpUxbF3q7t27wsNWIzQ+eWWNmTRpkuJDPQry8VE6lTlzZlWlShX1+fNnBTc3c+ZMGcPxgDpx4oTQrMFRKtFcpWczSy9wbAUiJJSHT6VKlRQf+XOId/DgQenTwKETigF/jRo11MePH4V/+fLlQuPIVdqTJ0+WNluntPfv3y/tTZs2SRt/oEDMs2jRIgUlX7x4Ua56rsGDB1u8R44cUbzyrTYqGjjUX79+rSpWrCjzYTFBLjdwV65ckcXD0SaGSNGAs9eQtpZz+vTp/3H4u0QDLiZ7HCvKKmw9hIjy2LFjNHXqVELkyKuPNmzYYPEgMnOX7NmzC6lUqVKUM2dOqVevXl2ucDMovLKJgRBXdP78eVqxYoXQcRZUF7ZYqSIqZGul+vXryxXRb+HChUUOnFdEYcDFFUsjwh/sY3B/cMlwuZjDXTAH9q4uXbrIvXA/7GlIJbR71mPc7lXT03KNOXAQgl0ctWjRgqZNm0a8yohdhuxFCKFTU9gNOtizZMlCb9++lf1hzZo11L59e0c/GgALxT0WCwO5GQ6XYm+DTIcOHaJevXoJf3J/2AMI2FD6jh07PMEK9mGkQLw1WJ8bN24QexBPjhd44OxKaNiwoTS/fPlCz549s3elug7FIUgYOHAg4Vh21apVUzUHu0GxZnabtGXLFmrVqpUsspQmgTWxmxY2WL295MiRgzgdkkDKTkfdnacGDjj22G6Zrfbt27elXrlyZYkirY40VGbPni1uqUePHjJaJ8jR7m+/DV4GAHS4SkR2kVwfLFLPax8L79GpUyfEBHayRJwgIPq09yH6HTdunPBqwDgYcoxNTyMmrhL+XwutrQph8Lp162jChAkinw770cAKRbG/PtJpg31fwF6BwkGCXBHKo2DVw/oQ1qMgNdizZ4/U9f51//59abv/jBo1ijBP0aJFqXXr1u5uQqpy9epVDx3K37hxI1WrVs3Rx8GO5LFIZfD6DvsuLBsphU4x9L7LQZm405j8qIYV7iipjSrxNoCTaYm8+IkkLeD9RMJsfkjVtWtXxQ9l3WPp0qWqfPnyws97g+J8RzHIEt1hPMJ/fv+nTp8+rcqVKyd8/EsVxbmZ2rdvn2KFS7jNSleILsuWLatwvxkzZih+BSV9mAfRLO+D1n3tFYzlBWUnSSivUxIGSdKTJ0+eOHjQQOoBeezpAFuXQhqD++KDNOT48ePWWKQdtWvXlj7oiheG1RetEi2qjFk6EE2AWPZ9//5dsSuzpkR+iHzJb4ESCxYsqJBvprWwNavfv387hqONOZHWuPvAiDQC43D1W6IB97+/OeEVGNPCyS3howvcnnahmhbpysqUiBNvVfCtBednkdh80dgTePjgSqPNiWg30jjPRD4JoQPO53M52LDvjBw5UvJJhOn4+VLYS0yCk6ArAe8lkdchYML7RyTJYS9/hcXhJTTe5gA8P241DKD+FcABiEiv2cIAUHIy/hWuMrmHDzPdABdS9AxwBriQaiCkYhuLM8CFVAMhFdtYnAEupBoIqdjG4gxwIdVASMU2FhdS4JJ9V4kfQJiSsRrQxzAiSeGxOJxr5B87ROI1tD+sAZwDBRY4Iu8ucfga3U007cBrwPz7+sBDlIyAHleZDJ8hB0wDBriAAeJXHAOcX00FjO8fDiHegN5OCnMAAAAASUVORK5CYII=\n",
      "text/plain": [
       "<IPython.core.display.Image object>"
      ]
     },
     "execution_count": 23,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "tf.keras.utils.plot_model(model, expand_nested=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Train"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [],
   "source": [
    "os.makedirs('./checkpoints', exist_ok=True)\n",
    "os.makedirs('./logs', exist_ok=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Optimisers"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "initial_lr = 1e-3\n",
    "var_decay = 1e-5\n",
    "\n",
    "model.compile(\n",
    "    optimizer=tf.keras.optimizers.Adam(lr=initial_lr, decay=var_decay),\n",
    "    loss=tf.keras.losses.SquaredHinge(),\n",
    "    metrics=[tf.keras.metrics.CategoricalAccuracy()]\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "trained = model.fit(train_dataset, epochs=10)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [],
   "source": [
    "initial = 0.01\n",
    "\n",
    "def scheduler(epoch: int) -> float:\n",
    "    \"\"\"\n",
    "    Computers the updated learning rate for each epoch\n",
    "    > we decay the learning rate by using a 1-bit right shift every 50 epochs.\n",
    "    :param epoch: The current epoch\n",
    "    :return: The updated learning rate\n",
    "    \"\"\"\n",
    "    batch = epoch // 50\n",
    "    multiplier = 0.1**batch\n",
    "    return initial * multiplier"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [],
   "source": [
    "optimiser = tf.keras.optimizers.Adam(learning_rate=initial, decay=1e-5)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {},
   "outputs": [],
   "source": [
    "model.compile(\n",
    "    optimizer=optimiser,\n",
    "    loss=tf.keras.losses.CategoricalCrossentropy(),\n",
    "    metrics=[tf.keras.metrics.CategoricalAccuracy(), \n",
    "             tf.keras.metrics.TopKCategoricalAccuracy()]\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Initialise Callbacks"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {},
   "outputs": [],
   "source": [
    "callbacks = [\n",
    "    lq.callbacks.\n",
    "    tf.keras.callbacks.EarlyStopping(patience=50),\n",
    "    tf.keras.callbacks.ModelCheckpoint('./checkpoints', save_best_only=True),\n",
    "    tf.keras.callbacks.TensorBoard(),\n",
    "    tf.keras.callbacks.LearningRateScheduler(scheduler, verbose=1)\n",
    "]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Fit"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "Epoch 00001: LearningRateScheduler reducing learning rate to 0.01.\n",
      "Epoch 1/500\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "W0710 23:12:02.873087 4513134016 deprecation.py:323] From /anaconda3/envs/TensorFlow-2.0/lib/python3.7/site-packages/tensorflow/python/ops/math_grad.py:1250: add_dispatch_support.<locals>.wrapper (from tensorflow.python.ops.array_ops) is deprecated and will be removed in a future version.\n",
      "Instructions for updating:\n",
      "Use tf.where in 2.0, which has the same broadcast rule as np.where\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      " 39/704 [>.............................] - ETA: 24:15 - loss: 2.3034 - categorical_accuracy: 0.1103 - top_k_categorical_accuracy: 0.5610 - quant_conv2d/flip_ratio: 0.0555 - quant_conv2d_1/flip_ratio: 0.0177 - quant_conv2d_2/flip_ratio: 0.0161 - quant_conv2d_3/flip_ratio: 0.0164 - quant_conv2d_4/flip_ratio: 0.0191 - quant_conv2d_5/flip_ratio: 0.0223 - quant_dense/flip_ratio: 0.0274 - quant_dense_1/flip_ratio: 0.0181 - quant_dense_2/flip_ratio: 0.0073"
     ]
    },
    {
     "ename": "KeyboardInterrupt",
     "evalue": "",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mKeyboardInterrupt\u001b[0m                         Traceback (most recent call last)",
      "\u001b[0;32m<ipython-input-29-542b814bf8ff>\u001b[0m in \u001b[0;36m<module>\u001b[0;34m\u001b[0m\n\u001b[1;32m      3\u001b[0m     \u001b[0mepochs\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;36m500\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m      4\u001b[0m     \u001b[0mcallbacks\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mcallbacks\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m----> 5\u001b[0;31m     \u001b[0mvalidation_data\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mval_dataset\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m      6\u001b[0m )\n",
      "\u001b[0;32m/anaconda3/envs/TensorFlow-2.0/lib/python3.7/site-packages/tensorflow/python/keras/engine/training.py\u001b[0m in \u001b[0;36mfit\u001b[0;34m(self, x, y, batch_size, epochs, verbose, callbacks, validation_split, validation_data, shuffle, class_weight, sample_weight, initial_epoch, steps_per_epoch, validation_steps, validation_freq, max_queue_size, workers, use_multiprocessing, **kwargs)\u001b[0m\n\u001b[1;32m    641\u001b[0m         \u001b[0mmax_queue_size\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mmax_queue_size\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    642\u001b[0m         \u001b[0mworkers\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mworkers\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 643\u001b[0;31m         use_multiprocessing=use_multiprocessing)\n\u001b[0m\u001b[1;32m    644\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    645\u001b[0m   def evaluate(self,\n",
      "\u001b[0;32m/anaconda3/envs/TensorFlow-2.0/lib/python3.7/site-packages/tensorflow/python/keras/engine/training_generator.py\u001b[0m in \u001b[0;36mfit\u001b[0;34m(self, model, x, y, batch_size, epochs, verbose, callbacks, validation_split, validation_data, shuffle, class_weight, sample_weight, initial_epoch, steps_per_epoch, validation_steps, validation_freq, **kwargs)\u001b[0m\n\u001b[1;32m    692\u001b[0m         \u001b[0mshuffle\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mshuffle\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    693\u001b[0m         \u001b[0minitial_epoch\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0minitial_epoch\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 694\u001b[0;31m         steps_name='steps_per_epoch')\n\u001b[0m\u001b[1;32m    695\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    696\u001b[0m   def evaluate(self,\n",
      "\u001b[0;32m/anaconda3/envs/TensorFlow-2.0/lib/python3.7/site-packages/tensorflow/python/keras/engine/training_generator.py\u001b[0m in \u001b[0;36mmodel_iteration\u001b[0;34m(model, data, steps_per_epoch, epochs, verbose, callbacks, validation_data, validation_steps, validation_freq, class_weight, max_queue_size, workers, use_multiprocessing, shuffle, initial_epoch, mode, batch_size, steps_name, **kwargs)\u001b[0m\n\u001b[1;32m    262\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    263\u001b[0m       \u001b[0mis_deferred\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;32mnot\u001b[0m \u001b[0mmodel\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_is_compiled\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 264\u001b[0;31m       \u001b[0mbatch_outs\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mbatch_function\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0mbatch_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    265\u001b[0m       \u001b[0;32mif\u001b[0m \u001b[0;32mnot\u001b[0m \u001b[0misinstance\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mbatch_outs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mlist\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    266\u001b[0m         \u001b[0mbatch_outs\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;34m[\u001b[0m\u001b[0mbatch_outs\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m/anaconda3/envs/TensorFlow-2.0/lib/python3.7/site-packages/tensorflow/python/keras/engine/training.py\u001b[0m in \u001b[0;36mtrain_on_batch\u001b[0;34m(self, x, y, sample_weight, class_weight, reset_metrics)\u001b[0m\n\u001b[1;32m    916\u001b[0m       \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_update_sample_weight_modes\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0msample_weights\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0msample_weights\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    917\u001b[0m       \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_make_train_function\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 918\u001b[0;31m       \u001b[0moutputs\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mtrain_function\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mins\u001b[0m\u001b[0;34m)\u001b[0m  \u001b[0;31m# pylint: disable=not-callable\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    919\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    920\u001b[0m     \u001b[0;32mif\u001b[0m \u001b[0mreset_metrics\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m/anaconda3/envs/TensorFlow-2.0/lib/python3.7/site-packages/tensorflow/python/keras/backend.py\u001b[0m in \u001b[0;36m__call__\u001b[0;34m(self, inputs)\u001b[0m\n\u001b[1;32m   3508\u001b[0m         \u001b[0mvalue\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mmath_ops\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mcast\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mvalue\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mtensor\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mdtype\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   3509\u001b[0m       \u001b[0mconverted_inputs\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mappend\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mvalue\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m-> 3510\u001b[0;31m     \u001b[0moutputs\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_graph_fn\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0mconverted_inputs\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   3511\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   3512\u001b[0m     \u001b[0;31m# EagerTensor.numpy() will often make a copy to ensure memory safety.\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m/anaconda3/envs/TensorFlow-2.0/lib/python3.7/site-packages/tensorflow/python/eager/function.py\u001b[0m in \u001b[0;36m__call__\u001b[0;34m(self, *args, **kwargs)\u001b[0m\n\u001b[1;32m    570\u001b[0m       raise TypeError(\"Keyword arguments {} unknown. Expected {}.\".format(\n\u001b[1;32m    571\u001b[0m           list(kwargs.keys()), list(self._arg_keywords)))\n\u001b[0;32m--> 572\u001b[0;31m     \u001b[0;32mreturn\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_call_flat\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0margs\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    573\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    574\u001b[0m   \u001b[0;32mdef\u001b[0m \u001b[0m_filtered_call\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m/anaconda3/envs/TensorFlow-2.0/lib/python3.7/site-packages/tensorflow/python/eager/function.py\u001b[0m in \u001b[0;36m_call_flat\u001b[0;34m(self, args)\u001b[0m\n\u001b[1;32m    669\u001b[0m     \u001b[0;31m# Only need to override the gradient in graph mode and when we have outputs.\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    670\u001b[0m     \u001b[0;32mif\u001b[0m \u001b[0mcontext\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mexecuting_eagerly\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;32mor\u001b[0m \u001b[0;32mnot\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0moutputs\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 671\u001b[0;31m       \u001b[0moutputs\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_inference_function\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mcall\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mctx\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0margs\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    672\u001b[0m     \u001b[0;32melse\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    673\u001b[0m       \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_register_gradient\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m/anaconda3/envs/TensorFlow-2.0/lib/python3.7/site-packages/tensorflow/python/eager/function.py\u001b[0m in \u001b[0;36mcall\u001b[0;34m(self, ctx, args)\u001b[0m\n\u001b[1;32m    443\u001b[0m             attrs=(\"executor_type\", executor_type,\n\u001b[1;32m    444\u001b[0m                    \"config_proto\", config),\n\u001b[0;32m--> 445\u001b[0;31m             ctx=ctx)\n\u001b[0m\u001b[1;32m    446\u001b[0m       \u001b[0;31m# Replace empty list with None\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    447\u001b[0m       \u001b[0moutputs\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0moutputs\u001b[0m \u001b[0;32mor\u001b[0m \u001b[0;32mNone\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m/anaconda3/envs/TensorFlow-2.0/lib/python3.7/site-packages/tensorflow/python/eager/execute.py\u001b[0m in \u001b[0;36mquick_execute\u001b[0;34m(op_name, num_outputs, inputs, attrs, ctx, name)\u001b[0m\n\u001b[1;32m     59\u001b[0m     tensors = pywrap_tensorflow.TFE_Py_Execute(ctx._handle, device_name,\n\u001b[1;32m     60\u001b[0m                                                \u001b[0mop_name\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0minputs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mattrs\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 61\u001b[0;31m                                                num_outputs)\n\u001b[0m\u001b[1;32m     62\u001b[0m   \u001b[0;32mexcept\u001b[0m \u001b[0mcore\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_NotOkStatusException\u001b[0m \u001b[0;32mas\u001b[0m \u001b[0me\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     63\u001b[0m     \u001b[0;32mif\u001b[0m \u001b[0mname\u001b[0m \u001b[0;32mis\u001b[0m \u001b[0;32mnot\u001b[0m \u001b[0;32mNone\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;31mKeyboardInterrupt\u001b[0m: "
     ]
    }
   ],
   "source": [
    "trained_model = model.fit(\n",
    "    train_dataset,\n",
    "    epochs=500, \n",
    "    callbacks=callbacks,\n",
    "    validation_data=val_dataset\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}