part6_cnns.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "deletable": false,
    "editable": false
   },
   "source": [
    "# Part 6: Convolutional Neural Networks in hls4ml\n",
    "\n",
    "In this notebook you will learn how to train a pruned and quantized convolutional neural network (CNN) and deploy it using hls4ml. For this exercise, we will use the Street View House Numbers (SVHN) Dataset (http://ufldl.stanford.edu/housenumbers/).\n",
    "\n",
    "The SVHN dataset consists of real-world images of house numbers extracted from Google Street View images. The format is similar to that of the MNIST dataset, but is a much more challenging real-world problem, as illustrated by the examples shown below.\n",
    "\n",
    "All the images are in RGB format and have been cropped to 32x32 pixels. \n",
    "Unlike MNIST, more than one digit can be present in the same image and in these cases, the center digit is used to assign a label to the image.\n",
    "Each image can belong to one of 10 classes, corresponding to digits 0 through 9.\n",
    "\n",
    "![alt text](images/test.png \"SVHN examples from the test dataset\")\n",
    "\n",
    "The SVHN dataset consists of 73,257 images for training (and 531,131 extra samples that are easier to classify and can be used as additional training data) and 26,032 images for testing."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "deletable": false,
    "editable": false
   },
   "source": [
    "## Start with the neccessary imports"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "import time\n",
    "import tensorflow.compat.v2 as tf\n",
    "import tensorflow_datasets as tfds\n",
    "\n",
    "os.environ['PATH'] = os.environ['XILINX_VITIS'] + '/bin:' + os.environ['PATH']"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "deletable": false,
    "editable": false
   },
   "source": [
    "## Fetch the SVHN dataset using Tensorflow Dataset\n",
    "\n",
    "In this part we will fetch the trainining, validation and test dataset using Tensorflow Datasets (https://www.tensorflow.org/datasets). We will not use the 'extra' training in order to save time, but you could fetch it by adding `split='train[:90%]+extra'`. We will use the first 90% of the training data for training and the last 10% for validation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ds_train, info = tfds.load('svhn_cropped', split='train[:90%]', with_info=True, as_supervised=True)\n",
    "ds_test = tfds.load('svhn_cropped', split='test', shuffle_files=True, as_supervised=True)\n",
    "ds_val = tfds.load('svhn_cropped', split='train[-10%:]', shuffle_files=True, as_supervised=True)\n",
    "\n",
    "assert isinstance(ds_train, tf.data.Dataset)\n",
    "train_size = int(info.splits['train'].num_examples)\n",
    "input_shape = info.features['image'].shape\n",
    "n_classes = info.features['label'].num_classes\n",
    "\n",
    "print('Training on {} samples of input shape {}, belonging to {} classes'.format(train_size, input_shape, n_classes))\n",
    "fig = tfds.show_examples(ds_train, info)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "deletable": false,
    "editable": false
   },
   "source": [
    "We'll use TensorFlow Dataset to prepare our datasets. We'll fetch the training dataset as tuples, and the test dataset as numpy arrays"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def preprocess(image, label, nclasses=10):\n",
    "    image = tf.cast(image, tf.float32) / 255.0\n",
    "    label = tf.one_hot(tf.squeeze(label), nclasses)\n",
    "    return image, label"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "batch_size = 1024\n",
    "\n",
    "train_data = ds_train.map(preprocess, n_classes)  # Get dataset as image and one-hot encoded labels, divided by max RGB\n",
    "train_data = train_data.shuffle(4096).batch(batch_size).prefetch(tf.data.experimental.AUTOTUNE)\n",
    "\n",
    "for example in train_data.take(1):\n",
    "    break\n",
    "print(\"X train batch shape = {}, Y train batch shape = {} \".format(example[0].shape, example[1].shape))\n",
    "\n",
    "val_data = ds_val.map(preprocess, n_classes)\n",
    "val_data = val_data.batch(batch_size)\n",
    "val_data = val_data.prefetch(tf.data.experimental.AUTOTUNE)\n",
    "\n",
    "# For  testing, we get the full dataset in memory as it's rather small.\n",
    "# We fetch it as numpy arrays to have access to labels and images separately\n",
    "X_test, Y_test = tfds.as_numpy(tfds.load('svhn_cropped', split='test', batch_size=-1, as_supervised=True))\n",
    "X_test, Y_test = preprocess(X_test, Y_test, nclasses=n_classes)\n",
    "print(\"X test batch shape = {}, Y test batch shape = {} \".format(X_test.shape, Y_test.shape))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "deletable": false,
    "editable": false
   },
   "source": [
    "## Defining the model\n",
    "\n",
    "We then need to define a model. For the lowest possible latency, each layer should have a maximum number of trainable parameters of 4096. This is due to fixed limits in the Vivado compiler, beyond which maximally unrolled (=parallel) compilation will fail. This will allow us to use `strategy = 'latency'` in the hls4ml part, rather than `strategy = 'resource'`, in turn resulting in lower latency"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from tensorflow.keras.layers import Input\n",
    "from tensorflow.keras.layers import BatchNormalization\n",
    "from tensorflow.keras.layers import Conv2D\n",
    "from tensorflow.keras.regularizers import l1\n",
    "from tensorflow.keras.layers import MaxPooling2D\n",
    "from tensorflow.keras.layers import Activation\n",
    "from tensorflow.keras.layers import Flatten\n",
    "from tensorflow.keras.layers import Dense\n",
    "\n",
    "from tensorflow.keras.models import Model\n",
    "\n",
    "filters_per_conv_layer = [16, 16, 24]\n",
    "neurons_per_dense_layer = [42, 64]\n",
    "\n",
    "x = x_in = Input(input_shape)\n",
    "\n",
    "for i, f in enumerate(filters_per_conv_layer):\n",
    "    print(('Adding convolutional block {} with N={} filters').format(i, f))\n",
    "    x = Conv2D(\n",
    "        int(f),\n",
    "        kernel_size=(3, 3),\n",
    "        strides=(1, 1),\n",
    "        kernel_initializer='lecun_uniform',\n",
    "        kernel_regularizer=l1(0.0001),\n",
    "        use_bias=False,\n",
    "        name='conv_{}'.format(i),\n",
    "    )(x)\n",
    "    x = BatchNormalization(name='bn_conv_{}'.format(i))(x)\n",
    "    x = Activation('relu', name='conv_act_%i' % i)(x)\n",
    "    x = MaxPooling2D(pool_size=(2, 2), name='pool_{}'.format(i))(x)\n",
    "x = Flatten()(x)\n",
    "\n",
    "for i, n in enumerate(neurons_per_dense_layer):\n",
    "    print(('Adding dense block {} with N={} neurons').format(i, n))\n",
    "    x = Dense(n, kernel_initializer='lecun_uniform', kernel_regularizer=l1(0.0001), name='dense_%i' % i, use_bias=False)(x)\n",
    "    x = BatchNormalization(name='bn_dense_{}'.format(i))(x)\n",
    "    x = Activation('relu', name='dense_act_%i' % i)(x)\n",
    "x = Dense(int(n_classes), name='output_dense')(x)\n",
    "x_out = Activation('softmax', name='output_softmax')(x)\n",
    "\n",
    "model = Model(inputs=[x_in], outputs=[x_out], name='keras_baseline')\n",
    "\n",
    "model.summary()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "deletable": false,
    "editable": false
   },
   "source": [
    "Lets check if this model can be implemented completely unrolled (=parallel)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for layer in model.layers:\n",
    "    if layer.__class__.__name__ in ['Conv2D', 'Dense']:\n",
    "        w = layer.get_weights()[0]\n",
    "        layersize = np.prod(w.shape)\n",
    "        print(\"{}: {}\".format(layer.name, layersize))  # 0 = weights, 1 = biases\n",
    "        if layersize > 4096:  # assuming that shape[0] is batch, i.e., 'None'\n",
    "            print(\"Layer {} is too large ({}), are you sure you want to train?\".format(layer.name, layersize))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "deletable": false,
    "editable": false
   },
   "source": [
    "Looks good! It's below the Vivado-enforced unroll limit of 4096.\n",
    "\n",
    "## Prune dense and convolutional layers\n",
    "Since we've seen in the previous notebooks that pruning can be done at no accuracy cost, let's prune the convolutional and dense layers to 50% sparsity, skipping the output layer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import tensorflow_model_optimization as tfmot\n",
    "from tensorflow_model_optimization.sparsity import keras as sparsity\n",
    "from tensorflow_model_optimization.python.core.sparsity.keras import pruning_callbacks\n",
    "\n",
    "NSTEPS = int(train_size * 0.9) // batch_size  # 90% train, 10% validation in 10-fold cross validation\n",
    "print('Number of training steps per epoch is {}'.format(NSTEPS))\n",
    "\n",
    "\n",
    "# Prune all convolutional and dense layers gradually from 0 to 50% sparsity every 2 epochs,\n",
    "# ending by the 10th epoch\n",
    "def pruneFunction(layer):\n",
    "    pruning_params = {\n",
    "        'pruning_schedule': sparsity.PolynomialDecay(\n",
    "            initial_sparsity=0.0, final_sparsity=0.50, begin_step=NSTEPS * 2, end_step=NSTEPS * 10, frequency=NSTEPS\n",
    "        )\n",
    "    }\n",
    "    if isinstance(layer, tf.keras.layers.Conv2D):\n",
    "        return tfmot.sparsity.keras.prune_low_magnitude(layer, **pruning_params)\n",
    "    if isinstance(layer, tf.keras.layers.Dense) and layer.name != 'output_dense':\n",
    "        return tfmot.sparsity.keras.prune_low_magnitude(layer, **pruning_params)\n",
    "    return layer\n",
    "\n",
    "\n",
    "model_pruned = tf.keras.models.clone_model(model, clone_function=pruneFunction)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "deletable": false,
    "editable": false
   },
   "source": [
    "## Train baseline\n",
    "\n",
    "We're now ready to train the model! We defined the batch size and n epochs above. We won't use callbacks that store the best weights only, since this might select a weight configuration that has not yet reached 50% sparsity."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "train = True  # True if you want to retrain, false if you want to load a previsously trained model\n",
    "\n",
    "n_epochs = 30\n",
    "\n",
    "if train:\n",
    "    LOSS = tf.keras.losses.CategoricalCrossentropy()\n",
    "    OPTIMIZER = tf.keras.optimizers.Adam(learning_rate=3e-3, beta_1=0.9, beta_2=0.999, epsilon=1e-07, amsgrad=True)\n",
    "\n",
    "    model_pruned.compile(loss=LOSS, optimizer=OPTIMIZER, metrics=[\"accuracy\"])\n",
    "\n",
    "    callbacks = [\n",
    "        tf.keras.callbacks.EarlyStopping(patience=10, verbose=1),\n",
    "        tf.keras.callbacks.ReduceLROnPlateau(monitor='val_loss', factor=0.5, patience=3, verbose=1),\n",
    "        pruning_callbacks.UpdatePruningStep(),\n",
    "    ]\n",
    "\n",
    "    start = time.time()\n",
    "    model_pruned.fit(train_data, epochs=n_epochs, validation_data=val_data, callbacks=callbacks)\n",
    "    end = time.time()\n",
    "\n",
    "    print('It took {} minutes to train Keras model'.format((end - start) / 60.0))\n",
    "\n",
    "    model_pruned.save('pruned_cnn_model.h5')\n",
    "\n",
    "else:\n",
    "    from qkeras.utils import _add_supported_quantized_objects\n",
    "    from tensorflow_model_optimization.python.core.sparsity.keras import pruning_wrapper\n",
    "\n",
    "    co = {}\n",
    "    _add_supported_quantized_objects(co)\n",
    "    co['PruneLowMagnitude'] = pruning_wrapper.PruneLowMagnitude\n",
    "    model_pruned = tf.keras.models.load_model('pruned_cnn_model.h5', custom_objects=co)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "deletable": false,
    "editable": false
   },
   "source": [
    "You'll notice the accuracy is lower than that in the hls4ml CNN paper (https://arxiv.org/abs/2101.05108) despite the model being the same. The reson for this is that we didn't use the ``extra`` training data in order to save time. If you want to futher optimize the network, increasing the training data is a good place to start. Enlarging the model architecture comes at a high latency/resource cost."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "deletable": false,
    "editable": false
   },
   "source": [
    "## Quantization and the fused Conv2D+BatchNormalization layer in QKeras\n",
    "Let's now create a pruned an quantized model using QKeras. For this, we will use a fused Convolutional and BatchNormalization (BN) layer from QKeras, which will further speed up the implementation when we implement the model using hls4ml. \n",
    "There is currently no fused Dense+BatchNoralization layer available in QKeras, so we'll use Keras BatchNormalization when BN follows a Dense layer for now. We'll use the same precision everywhere, namely a bit width of 6 and 0 integer bits (this will be implemented as``<6,1>`` in hls4ml, due to the missing sign-bit). For now, make sure to set ```use_bias=True``` in ```QConv2DBatchnorm``` to avoid problems during synthesis."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from qkeras import QActivation\n",
    "from qkeras import QDense, QConv2DBatchnorm\n",
    "\n",
    "x = x_in = Input(shape=input_shape)\n",
    "\n",
    "for i, f in enumerate(filters_per_conv_layer):\n",
    "    print(('Adding fused QConv+BN block {} with N={} filters').format(i, f))\n",
    "    x = QConv2DBatchnorm(\n",
    "        int(f),\n",
    "        kernel_size=(3, 3),\n",
    "        strides=(1, 1),\n",
    "        kernel_quantizer=\"quantized_bits(6,0,alpha=1)\",\n",
    "        bias_quantizer=\"quantized_bits(6,0,alpha=1)\",\n",
    "        kernel_initializer='lecun_uniform',\n",
    "        kernel_regularizer=l1(0.0001),\n",
    "        use_bias=True,\n",
    "        name='fused_convbn_{}'.format(i),\n",
    "    )(x)\n",
    "    x = QActivation('quantized_relu(6)', name='conv_act_%i' % i)(x)\n",
    "    x = MaxPooling2D(pool_size=(2, 2), name='pool_{}'.format(i))(x)\n",
    "x = Flatten()(x)\n",
    "\n",
    "for i, n in enumerate(neurons_per_dense_layer):\n",
    "    print(('Adding QDense block {} with N={} neurons').format(i, n))\n",
    "    x = QDense(\n",
    "        n,\n",
    "        kernel_quantizer=\"quantized_bits(6,0,alpha=1)\",\n",
    "        kernel_initializer='lecun_uniform',\n",
    "        kernel_regularizer=l1(0.0001),\n",
    "        name='dense_%i' % i,\n",
    "        use_bias=False,\n",
    "    )(x)\n",
    "    x = BatchNormalization(name='bn_dense_{}'.format(i))(x)\n",
    "    x = QActivation('quantized_relu(6)', name='dense_act_%i' % i)(x)\n",
    "x = Dense(int(n_classes), name='output_dense')(x)\n",
    "x_out = Activation('softmax', name='output_softmax')(x)\n",
    "qmodel = Model(inputs=[x_in], outputs=[x_out], name='qkeras')\n",
    "\n",
    "qmodel.summary()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Print the quantized layers\n",
    "from qkeras.autoqkeras.utils import print_qmodel_summary\n",
    "\n",
    "print_qmodel_summary(qmodel)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "deletable": false,
    "editable": false
   },
   "source": [
    "You see that a bias quantizer is defined, although we are not using a bias term for the layers. This is set automatically by QKeras. In addition, you'll note that ``alpha='1'``. This sets the weight scale per channel to 1 (no scaling). The default is ``alpha='auto_po2'``, which sets the weight scale per channel to be a power-of-2, such that an actual hardware implementation can be performed by just shifting the result of the convolutional/dense layer to the right or left by checking the sign of the scale and then taking the log2 of the scale.\n",
    "\n",
    "Let's now prune and train this model! If you want, you can also train the unpruned version, ``qmodel`` and see how the performance compares. We will stick to the pruned one here. Again, we do not use a model checkpoint which stores the best weights, in order to ensure the model is trained to the desired sparsity."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "qmodel_pruned = tf.keras.models.clone_model(qmodel, clone_function=pruneFunction)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "train = True\n",
    "\n",
    "n_epochs = 30\n",
    "if train:\n",
    "    LOSS = tf.keras.losses.CategoricalCrossentropy()\n",
    "    OPTIMIZER = tf.keras.optimizers.Adam(learning_rate=3e-3, beta_1=0.9, beta_2=0.999, epsilon=1e-07, amsgrad=True)\n",
    "    qmodel_pruned.compile(loss=LOSS, optimizer=OPTIMIZER, metrics=[\"accuracy\"])\n",
    "\n",
    "    callbacks = [\n",
    "        tf.keras.callbacks.EarlyStopping(patience=10, verbose=1),\n",
    "        tf.keras.callbacks.ReduceLROnPlateau(monitor='val_loss', factor=0.5, patience=3, verbose=1),\n",
    "        pruning_callbacks.UpdatePruningStep(),\n",
    "    ]\n",
    "\n",
    "    start = time.time()\n",
    "    history = qmodel_pruned.fit(train_data, epochs=n_epochs, validation_data=val_data, callbacks=callbacks, verbose=1)\n",
    "    end = time.time()\n",
    "    print('\\n It took {} minutes to train!\\n'.format((end - start) / 60.0))\n",
    "\n",
    "    qmodel_pruned.save('quantized_pruned_cnn_model.h5')\n",
    "\n",
    "else:\n",
    "    from qkeras.utils import _add_supported_quantized_objects\n",
    "    from tensorflow_model_optimization.python.core.sparsity.keras import pruning_wrapper\n",
    "\n",
    "    co = {}\n",
    "    _add_supported_quantized_objects(co)\n",
    "    co['PruneLowMagnitude'] = pruning_wrapper.PruneLowMagnitude\n",
    "    qmodel_pruned = tf.keras.models.load_model('quantized_pruned_cnn_model.h5', custom_objects=co)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "deletable": false,
    "editable": false
   },
   "source": [
    "We note that training a model quantization aware, takes around twice as long as when not quantizing during training!\n",
    "The validation accuracy is very similar to that of the floating point model equivalent, despite containing significantly less information \n",
    "\n",
    "## Performance\n",
    "Let's look at some ROC curves to compare the performance. Lets choose a few numbers so it doesn't get confusing. Feel free to change the numbers in ``labels``."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "predict_baseline = model_pruned.predict(X_test)\n",
    "test_score_baseline = model_pruned.evaluate(X_test, Y_test)\n",
    "\n",
    "predict_qkeras = qmodel_pruned.predict(X_test)\n",
    "test_score_qkeras = qmodel_pruned.evaluate(X_test, Y_test)\n",
    "\n",
    "print('Keras accuracy = {} , QKeras 6-bit accuracy = {}'.format(test_score_baseline[1], test_score_qkeras[1]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "import pandas as pd\n",
    "from sklearn import metrics\n",
    "\n",
    "\n",
    "labels = ['%i' % nr for nr in range(0, n_classes)]  # If you want to look at all the labels\n",
    "# labels = ['0','1','9'] # Look at only a few labels, here for digits 0, 1 and 9\n",
    "print('Plotting ROC for labels {}'.format(labels))\n",
    "\n",
    "df = pd.DataFrame()\n",
    "df_q = pd.DataFrame()\n",
    "fpr = {}\n",
    "tpr = {}\n",
    "auc1 = {}\n",
    "fpr_q = {}\n",
    "tpr_q = {}\n",
    "auc1_q = {}\n",
    "%matplotlib inline\n",
    "colors = ['#67001f', '#b2182b', '#d6604d', '#f4a582', '#fddbc7', '#d1e5f0', '#92c5de', '#4393c3', '#2166ac', '#053061']\n",
    "fig, ax = plt.subplots(figsize=(10, 10))\n",
    "for i, label in enumerate(labels):\n",
    "    df[label] = Y_test[:, int(label)]\n",
    "    df[label + '_pred'] = predict_baseline[:, int(label)]\n",
    "    fpr[label], tpr[label], threshold = metrics.roc_curve(df[label], df[label + '_pred'])\n",
    "    auc1[label] = metrics.auc(fpr[label], tpr[label])\n",
    "\n",
    "    df_q[label] = Y_test[:, int(label)]\n",
    "    df_q[label + '_pred'] = predict_qkeras[:, int(label)]\n",
    "    fpr_q[label], tpr_q[label], threshold_q = metrics.roc_curve(df_q[label], df_q[label + '_pred'])\n",
    "    auc1_q[label] = metrics.auc(fpr_q[label], tpr_q[label])\n",
    "\n",
    "    plt.plot(\n",
    "        fpr[label],\n",
    "        tpr[label],\n",
    "        label=r'{}, AUC Keras = {:.1f}% AUC QKeras = {:.1f}%)'.format(label, auc1[label] * 100, auc1_q[label] * 100),\n",
    "        linewidth=1.5,\n",
    "        c=colors[i],\n",
    "        linestyle='solid',\n",
    "    )\n",
    "    plt.plot(fpr_q[label], tpr_q[label], linewidth=1.5, c=colors[i], linestyle='dotted')\n",
    "\n",
    "plt.semilogx()\n",
    "plt.ylabel(\"True Positive Rate\")\n",
    "plt.xlabel(\"False Positive Rate\")\n",
    "plt.xlim(0.01, 1.0)\n",
    "plt.ylim(0.5, 1.1)\n",
    "plt.legend(loc='lower right')\n",
    "plt.figtext(\n",
    "    0.2,\n",
    "    0.83,\n",
    "    r'Accuracy Keras = {:.1f}% QKeras 8-bit = {:.1f}%'.format(test_score_baseline[1] * 100, test_score_qkeras[1] * 100),\n",
    "    wrap=True,\n",
    "    horizontalalignment='left',\n",
    "    verticalalignment='center',\n",
    ")\n",
    "from matplotlib.lines import Line2D\n",
    "\n",
    "lines = [Line2D([0], [0], ls='-'), Line2D([0], [0], ls='--')]\n",
    "from matplotlib.legend import Legend\n",
    "\n",
    "leg = Legend(ax, lines, labels=['Keras', 'QKeras'], loc='lower right', frameon=False)\n",
    "ax.add_artist(leg)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "deletable": false,
    "editable": false
   },
   "source": [
    "The difference in AUC between the fp32 Keras model and the 8-bit QKeras model, is small, as we have seen for the previous examples. You can find a bonus exercise below, **Bonus: Automatic quantization**, where we'll use AutoQKeras to find the best heterogeneously quantized model, given a set of resource and accuracy constriants.\n",
    "### Check sparsity\n",
    "Let's also check the per-layer sparsity:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def doWeights(model):\n",
    "    allWeightsByLayer = {}\n",
    "    for layer in model.layers:\n",
    "        if (layer._name).find(\"batch\") != -1 or len(layer.get_weights()) < 1:\n",
    "            continue\n",
    "        weights = layer.weights[0].numpy().flatten()\n",
    "        allWeightsByLayer[layer._name] = weights\n",
    "        print('Layer {}: % of zeros = {}'.format(layer._name, np.sum(weights == 0) / np.size(weights)))\n",
    "\n",
    "    labelsW = []\n",
    "    histosW = []\n",
    "\n",
    "    for key in reversed(sorted(allWeightsByLayer.keys())):\n",
    "        labelsW.append(key)\n",
    "        histosW.append(allWeightsByLayer[key])\n",
    "\n",
    "    fig = plt.figure(figsize=(10, 10))\n",
    "    bins = np.linspace(-1.5, 1.5, 50)\n",
    "    plt.hist(histosW, bins, histtype='stepfilled', stacked=True, label=labelsW, edgecolor='black')\n",
    "    plt.legend(frameon=False, loc='upper left')\n",
    "    plt.ylabel('Number of Weights')\n",
    "    plt.xlabel('Weights')\n",
    "    plt.figtext(0.2, 0.38, model._name, wrap=True, horizontalalignment='left', verticalalignment='center')\n",
    "\n",
    "\n",
    "doWeights(model_pruned)\n",
    "doWeights(qmodel_pruned)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "deletable": false,
    "editable": false
   },
   "source": [
    "We see that 50% of the weights per layer are set to zero, as expected.\n",
    "Now, let's synthesize the floating point Keras model and the QKeras quantized model!"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "deletable": false,
    "editable": false
   },
   "source": [
    "## CNNs in hls4ml\n",
    "\n",
    "In this part, we will take the two models we trained above (the floating-point 32 Keras model and the 6-bit QKeras model), and synthesize them with hls4ml. Although your models are probably already in memory, let's load them from scratch. We need to pass the appropriate custom QKeras/pruning layers when loading, and remove the pruning parameters that were saved together with the model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from tensorflow_model_optimization.sparsity.keras import strip_pruning\n",
    "from tensorflow_model_optimization.python.core.sparsity.keras import pruning_wrapper\n",
    "\n",
    "from qkeras.utils import _add_supported_quantized_objects\n",
    "\n",
    "co = {}\n",
    "_add_supported_quantized_objects(co)\n",
    "co['PruneLowMagnitude'] = pruning_wrapper.PruneLowMagnitude\n",
    "\n",
    "model = tf.keras.models.load_model('pruned_cnn_model.h5', custom_objects=co)\n",
    "model = strip_pruning(model)\n",
    "\n",
    "qmodel = tf.keras.models.load_model('quantized_pruned_cnn_model.h5', custom_objects=co)\n",
    "qmodel = strip_pruning(qmodel)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now, we need to define the hls4ml and Vivado configurations. Two things will change with respect to what was done in the previous exercises. First, we will use ``io_type='io_stream'`` in the Vitis_HLS configuration.\n",
    "\n",
    "---\n",
    "****You must use ``io_type='io_stream'`` if attempting to synthesize a large convolutional neural network.****\n",
    "\n",
    "---\n",
    "The CNN implementation in hls4ml is based on streams, which are synthesized in hardware as first in, first out (FIFO) buffers. Shift registers are used to keep track of the last  ``<kernel height - 1>`` rows of input pixels, and maintains a shifting snapshot of the convolution kernel.\n",
    "\n",
    "This is illustrated  in the gif below. Here, the input image is at the top-left and the output image at the bottom left. The top right image shows the internal state of the shift registers and convolutional kernel. The red square indicates the current pixels contained within the convolutional kernel.\n",
    "\n",
    "![alt text](images/conv2d_animation.gif \"The implementation of convolutional layers in hls4ml.\")\n",
    "\n",
    "Lastly, we will use ``['Strategy'] = 'Latency'`` for all the layers in the hls4ml configuration. If one layer would have >4096 elements, we sould set ``['Strategy'] = 'Resource'`` for that layer, or increase the reuse factor by hand. You can find examples of how to do this below."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import hls4ml\n",
    "import plotting\n",
    "\n",
    "# First, the baseline model\n",
    "hls_config = hls4ml.utils.config_from_keras_model(\n",
    "    model, granularity='name', backend='Vitis', default_precision='ap_fixed<16,6>'\n",
    ")\n",
    "\n",
    "plotting.print_dict(hls_config)\n",
    "\n",
    "\n",
    "hls_model = hls4ml.converters.convert_from_keras_model(\n",
    "    model,\n",
    "    hls_config=hls_config,\n",
    "    backend='Vitis',\n",
    "    output_dir='model_1/hls4ml_prj',\n",
    "    part='xcu250-figd2104-2L-e',\n",
    "    io_type='io_stream',\n",
    ")\n",
    "hls_model.compile()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "deletable": false,
    "editable": false
   },
   "source": [
    "Let's get a nice overview over the various shapes and precisions used for each layer through ``hls4ml.utils.plot_model``, as well as look at the weight profile using ``hls4ml.model.profiling.numerical``. The weight profiling returns two plots: Before (top) and after (bottom) various optimizations applied to the HLS model before the final translation to HLS, for instance the fusing of Dense and BatchNormalization layers."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "hls4ml.utils.plot_model(hls_model, show_shapes=True, show_precision=True, to_file=None)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "hls4ml.model.profiling.numerical(model=model, hls_model=hls_model)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "deletable": false,
    "editable": false
   },
   "source": [
    "The colored boxes are the distribution of the weights of the model, and the gray band illustrates the numerical range covered by the chosen fixed point precision. As we configured, this model uses a precision of ``ap_fixed<16,6>`` for all layers of the model. Let's now build our QKeras model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Then the QKeras model\n",
    "hls_config_q = hls4ml.utils.config_from_keras_model(qmodel, granularity='name', backend='Vitis')\n",
    "hls_config_q['Model']['ReuseFactor'] = 1\n",
    "hls_config['Model']['Precision'] = 'ap_fixed<16,6>'\n",
    "hls_config_q['LayerName']['output_softmax']['Strategy'] = 'Stable'\n",
    "plotting.print_dict(hls_config_q)\n",
    "\n",
    "cfg_q = hls4ml.converters.create_config(backend='Vitis')\n",
    "cfg_q['IOType'] = 'io_stream'  # Must set this if using CNNs!\n",
    "cfg_q['HLSConfig'] = hls_config_q\n",
    "cfg_q['KerasModel'] = qmodel\n",
    "cfg_q['OutputDir'] = 'quantized_pruned_cnn/'\n",
    "cfg_q['XilinxPart'] = 'xcu250-figd2104-2L-e'\n",
    "\n",
    "hls_model_q = hls4ml.converters.keras_to_hls(cfg_q)\n",
    "hls_model_q.compile()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "deletable": false,
    "editable": false
   },
   "source": [
    "Let's plot the model and profile the weights her too"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "hls4ml.model.profiling.numerical(model=qmodel, hls_model=hls_model_q)\n",
    "hls4ml.utils.plot_model(hls_model_q, show_shapes=True, show_precision=True, to_file=None)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "deletable": false,
    "editable": false
   },
   "source": [
    "For the 6-bit QKeras model, we see that different precisions are used for different layers."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "deletable": false,
    "editable": false
   },
   "source": [
    "### Accuracy with bit-accurate emulation \n",
    "Let's check that the hls4ml accuracy matches the original. This usually takes some time, so let's do it over a reduced dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "X_test_reduced = X_test[:3000]\n",
    "Y_test_reduced = Y_test[:3000]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "y_predict = model.predict(X_test_reduced)\n",
    "y_predict_hls4ml = hls_model.predict(np.ascontiguousarray(X_test_reduced))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "y_predict_q = qmodel.predict(X_test_reduced)\n",
    "y_predict_hls4ml_q = hls_model_q.predict(np.ascontiguousarray(X_test_reduced))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import plotting\n",
    "from sklearn.metrics import accuracy_score\n",
    "\n",
    "\n",
    "def plotROC(Y, y_pred, y_pred_hls4ml, label=\"Model\"):\n",
    "    accuracy_keras = float(accuracy_score(np.argmax(Y, axis=1), np.argmax(y_pred, axis=1)))\n",
    "    accuracy_hls4ml = float(accuracy_score(np.argmax(Y, axis=1), np.argmax(y_pred_hls4ml, axis=1)))\n",
    "\n",
    "    print(\"Accuracy Keras:  {}\".format(accuracy_keras))\n",
    "    print(\"Accuracy hls4ml: {}\".format(accuracy_hls4ml))\n",
    "\n",
    "    fig, ax = plt.subplots(figsize=(9, 9))\n",
    "    _ = plotting.makeRoc(Y, y_pred, labels=['%i' % nr for nr in range(n_classes)])\n",
    "    plt.gca().set_prop_cycle(None)  # reset the colors\n",
    "    _ = plotting.makeRoc(Y, y_pred_hls4ml, labels=['%i' % nr for nr in range(n_classes)], linestyle='--')\n",
    "\n",
    "    from matplotlib.lines import Line2D\n",
    "\n",
    "    lines = [Line2D([0], [0], ls='-'), Line2D([0], [0], ls='--')]\n",
    "    from matplotlib.legend import Legend\n",
    "\n",
    "    leg = Legend(ax, lines, labels=['Keras', 'hls4ml'], loc='lower right', frameon=False)\n",
    "    ax.add_artist(leg)\n",
    "    plt.figtext(0.2, 0.38, label, wrap=True, horizontalalignment='left', verticalalignment='center')\n",
    "    plt.ylim(0.01, 1.0)\n",
    "    plt.xlim(0.7, 1.0)\n",
    "\n",
    "\n",
    "# Plot the pruned floating point model:\n",
    "plotROC(Y_test_reduced, y_predict, y_predict_hls4ml, label=\"Keras\")\n",
    "\n",
    "# Plot the pruned and quantized QKeras model\n",
    "plotROC(Y_test_reduced, y_predict_q, y_predict_hls4ml_q, label=\"QKeras\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "deletable": false,
    "editable": false
   },
   "source": [
    "Looks good! Let's synthesize the models. \n",
    "## Logic synthesis\n",
    "This takes quite a while for CNN models, up to one hour for the models considered here. In the interest of time, we have therefore provided the neccessary reports for the models considered. You can also synthesize them yourself if you have time, and as usual follow the progress using ``tail -f pruned_cnn/vivado_hls.log`` and ``tail -f quantized_pruned_cnn/vivado_hls.log``.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "synth = False  # Only if you want to synthesize the models yourself (>1h per model) rather than look at the provided reports.\n",
    "if synth:\n",
    "    hls_model.build(csim=False, synth=True, vsynth=True)\n",
    "    hls_model_q.build(csim=False, synth=True, vsynth=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "deletable": false,
    "editable": false
   },
   "source": [
    "We extract the latency from the C synthesis, namely the report in ```<project_dir>/myproject_prj/solution1/syn/report/myproject_csynth.rpt```. A more accurate latency estimate can be obtained from running cosim by passing ```hls_model.build(csim=False, synth=True, vsynth=True, cosim=True)``` ( = C/RTL cosimulation, synthesised HLS code is run on a simulator and tested on C test bench) but this takes a lot of time so we will skip it here.\n",
    "The resource estimates are obtained from the Vivado logic synthesis, and can be extracted from the report in ```<project_dir>/vivado_synth.rpt```. Let's fetch the most relevant numbers:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def getReports(indir):\n",
    "    data_ = {}\n",
    "\n",
    "    report_vsynth = Path('{}/vivado_synth.rpt'.format(indir))\n",
    "    report_csynth = Path('{}/myproject_prj/solution1/syn/report/myproject_csynth.rpt'.format(indir))\n",
    "\n",
    "    if report_vsynth.is_file() and report_csynth.is_file():\n",
    "        print('Found valid vsynth and synth in {}! Fetching numbers'.format(indir))\n",
    "\n",
    "        # Get the resources from the logic synthesis report\n",
    "        with report_vsynth.open() as report:\n",
    "            lines = np.array(report.readlines())\n",
    "            data_['lut'] = int(lines[np.array(['CLB LUTs*' in line for line in lines])][0].split('|')[2])\n",
    "            data_['ff'] = int(lines[np.array(['CLB Registers' in line for line in lines])][0].split('|')[2])\n",
    "            data_['bram'] = float(lines[np.array(['Block RAM Tile' in line for line in lines])][0].split('|')[2])\n",
    "            data_['dsp'] = int(lines[np.array(['DSPs' in line for line in lines])][0].split('|')[2])\n",
    "            data_['lut_rel'] = float(lines[np.array(['CLB LUTs*' in line for line in lines])][0].split('|')[5])\n",
    "            data_['ff_rel'] = float(lines[np.array(['CLB Registers' in line for line in lines])][0].split('|')[5])\n",
    "            data_['bram_rel'] = float(lines[np.array(['Block RAM Tile' in line for line in lines])][0].split('|')[5])\n",
    "            data_['dsp_rel'] = float(lines[np.array(['DSPs' in line for line in lines])][0].split('|')[5])\n",
    "\n",
    "        with report_csynth.open() as report:\n",
    "            lines = np.array(report.readlines())\n",
    "            lat_line = lines[np.argwhere(np.array(['Latency (cycles)' in line for line in lines])).flatten()[0] + 3]\n",
    "            data_['latency_clks'] = int(lat_line.split('|')[2])\n",
    "            data_['latency_mus'] = float(lat_line.split('|')[2]) * 5.0 / 1000.0\n",
    "            data_['latency_ii'] = int(lat_line.split('|')[6])\n",
    "\n",
    "    return data_"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from pathlib import Path\n",
    "import pprint\n",
    "\n",
    "data_pruned_ref = getReports('pruned_cnn')\n",
    "data_quantized_pruned = getReports('quantized_pruned_cnn')\n",
    "\n",
    "print(\"\\n Resource usage and latency: Pruned\")\n",
    "pprint.pprint(data_pruned_ref)\n",
    "print(\"\\n Resource usage and latency: Pruned + quantized\")\n",
    "pprint.pprint(data_quantized_pruned)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We see that the latency is of around 5 microseconds for both the quantized and the unquantized model, but that the resources are signifcantly reduced using QKeras.\n",
    "\n",
    "Congratulations! You have now reached the end of this notebook. If you have some spare time, you can have a look at the bonus exercise below, where you will learn how to perform a bayesian optimization over the QKeras quantizers in order to obtain an optimally heterogeneously quantized model.\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "deletable": false,
    "editable": false
   },
   "source": [
    "## Bonus exercise: Automatic quantization with AutoQKeras\n",
    "\n",
    "In this bonus exercise, you will learn how to find the optimal heterogeneously quantized model using AutoQKeras.\n",
    "For more details, you can look at the [AutoQKeras notebook](https://github.com/google/qkeras/blob/master/notebook/AutoQKeras.ipynb). \n",
    "\n",
    "Let's first check the estimated energy consumption of the QKeras 6-bit model using QTools. By setting ```for_reference=True``` you can print out the unquantized model energy consumption and compare the two. Note that this only works for QKeras layers. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "filters_per_conv_layer = [16, 16, 24]\n",
    "neurons_per_dense_layer = [42, 64]\n",
    "\n",
    "x = x_in = Input(input_shape)\n",
    "\n",
    "for i, f in enumerate(filters_per_conv_layer):\n",
    "    print(('Adding convolutional block {} with N={} filters').format(i, f))\n",
    "    x = Conv2D(\n",
    "        int(f),\n",
    "        kernel_size=(3, 3),\n",
    "        strides=(1, 1),\n",
    "        kernel_initializer='lecun_uniform',\n",
    "        kernel_regularizer=l1(0.0001),\n",
    "        use_bias=False,\n",
    "        name='conv_{}'.format(i),\n",
    "    )(x)\n",
    "    x = BatchNormalization(name='bn_conv_{}'.format(i))(x)\n",
    "    x = Activation('relu', name='conv_act_%i' % i)(x)\n",
    "    x = MaxPooling2D(pool_size=(2, 2), name='pool_{}'.format(i))(x)\n",
    "x = Flatten()(x)\n",
    "\n",
    "for i, n in enumerate(neurons_per_dense_layer):\n",
    "    print(('Adding dense block {} with N={} neurons').format(i, n))\n",
    "    x = Dense(n, kernel_initializer='lecun_uniform', kernel_regularizer=l1(0.0001), name='dense_%i' % i, use_bias=False)(x)\n",
    "    x = BatchNormalization(name='bn_dense_{}'.format(i))(x)\n",
    "    x = Activation('relu', name='dense_act_%i' % i)(x)\n",
    "x = Dense(int(n_classes), name='output_dense')(x)\n",
    "x_out = Activation('softmax', name='output_softmax')(x)\n",
    "\n",
    "baseline_model = Model(inputs=[x_in], outputs=[x_out], name='keras_baseline')\n",
    "\n",
    "LOSS = tf.keras.losses.CategoricalCrossentropy()\n",
    "OPTIMIZER = tf.keras.optimizers.Adam(learning_rate=3e-3, beta_1=0.9, beta_2=0.999, epsilon=1e-07, amsgrad=True)\n",
    "\n",
    "baseline_model.compile(loss=LOSS, optimizer=OPTIMIZER, metrics=[\"accuracy\"])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from qkeras import print_qstats\n",
    "\n",
    "# for automatic quantization\n",
    "import pprint\n",
    "from qkeras.autoqkeras import *\n",
    "from qkeras import *\n",
    "from qkeras.utils import model_quantize\n",
    "\n",
    "from qkeras.qtools import run_qtools\n",
    "from qkeras.qtools import settings as qtools_settings\n",
    "from tensorflow_model_optimization.python.core.sparsity.keras import pruning_wrapper\n",
    "from qkeras import quantized_bits\n",
    "from qkeras import QDense, QActivation\n",
    "\n",
    "q = run_qtools.QTools(\n",
    "    baseline_model,\n",
    "    process=\"horowitz\",\n",
    "    source_quantizers=[quantized_bits(16, 5, 1)],\n",
    "    is_inference=True,\n",
    "    weights_path=None,\n",
    "    keras_quantizer=\"fp16\",\n",
    "    keras_accumulator=\"fp16\",\n",
    "    for_reference=False,\n",
    ")\n",
    "q.qtools_stats_print()\n",
    "\n",
    "energy_dict = q.pe(\n",
    "    weights_on_memory=\"fixed\", activations_on_memory=\"fixed\", min_sram_size=8 * 16 * 1024 * 1024, rd_wr_on_io=False\n",
    ")\n",
    "\n",
    "# get stats of energy distribution in each layer\n",
    "energy_profile = q.extract_energy_profile(qtools_settings.cfg.include_energy, energy_dict)\n",
    "# extract sum of energy of each layer according to the rule specified in\n",
    "# qtools_settings.cfg.include_energy\n",
    "total_energy = q.extract_energy_sum(qtools_settings.cfg.include_energy, energy_dict)\n",
    "\n",
    "pprint.pprint(energy_profile)\n",
    "print()\n",
    "\n",
    "print(\"Total energy: {:.6f} uJ\".format(total_energy / 1000000.0))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now, lets use AutoQKeras to find an optimally heterogeneously quantized model for us. For more details, check the AutoQKeras tutorial linked above. As baseline model, we'll use the pruned floating point Keras model from above."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# These are the quantizers we'll test in the bayesian optimization\n",
    "quantization_config = {\n",
    "    \"kernel\": {\n",
    "        \"quantized_bits(2,0,1,alpha=1.0)\": 2,\n",
    "        \"quantized_bits(4,0,1,alpha=1.0)\": 4,\n",
    "        \"quantized_bits(6,0,1,alpha=1.0)\": 6,\n",
    "        \"quantized_bits(8,0,1,alpha=1.0)\": 8,\n",
    "    },\n",
    "    \"bias\": {\n",
    "        \"quantized_bits(2,0,1,alpha=1.0)\": 2,\n",
    "        \"quantized_bits(4,0,1,alpha=1.0)\": 4,\n",
    "        \"quantized_bits(6,0,1,alpha=1.0)\": 6,\n",
    "        \"quantized_bits(8,0,1,alpha=1.0)\": 8,\n",
    "    },\n",
    "    \"activation\": {\n",
    "        \"quantized_relu(3,1)\": 3,\n",
    "        \"quantized_relu(4,2)\": 4,\n",
    "        \"quantized_relu(8,2)\": 8,\n",
    "        \"quantized_relu(8,4)\": 8,\n",
    "        \"quantized_relu(16,6)\": 16,\n",
    "    },\n",
    "    \"linear\": {\n",
    "        \"quantized_bits(2,0,1,alpha=1.0)\": 2,\n",
    "        \"quantized_bits(4,0,1,alpha=1.0)\": 4,\n",
    "        \"quantized_bits(6,0,1,alpha=1.0)\": 6,\n",
    "        \"quantized_bits(8,0,1,alpha=1.0)\": 8,\n",
    "    },\n",
    "}\n",
    "\n",
    "# These are the layer types we will quantize\n",
    "limit = {\n",
    "    \"Dense\": [8, 8, 16],\n",
    "    \"Conv2D\": [8, 8, 16],\n",
    "    \"Activation\": [16],\n",
    "}\n",
    "\n",
    "# Use this if you want to minimize the model bit size\n",
    "goal_bits = {\n",
    "    \"type\": \"bits\",\n",
    "    \"params\": {\n",
    "        \"delta_p\": 8.0,  # We tolerate up to a +8% accuracy change\n",
    "        \"delta_n\": 8.0,  # We tolerate down to a -8% accuracy change\n",
    "        \"rate\": 2.0,  # We want a x2 times smaller model\n",
    "        \"stress\": 1.0,  # Force the reference model size to be smaller by setting stress<1\n",
    "        \"input_bits\": 8,\n",
    "        \"output_bits\": 8,\n",
    "        \"ref_bits\": 8,\n",
    "        \"config\": {\"default\": [\"parameters\", \"activations\"]},\n",
    "    },\n",
    "}\n",
    "\n",
    "# Use this if you want to minimize the model energy consumption\n",
    "goal_energy = {\n",
    "    \"type\": \"energy\",\n",
    "    \"params\": {\n",
    "        \"delta_p\": 8.0,\n",
    "        \"delta_n\": 8.0,\n",
    "        \"rate\": 2.0,\n",
    "        \"stress\": 1.0,\n",
    "        \"process\": \"horowitz\",\n",
    "        \"parameters_on_memory\": [\"sram\", \"sram\"],\n",
    "        \"activations_on_memory\": [\"sram\", \"sram\"],\n",
    "        \"rd_wr_on_io\": [False, False],\n",
    "        \"min_sram_size\": [0, 0],\n",
    "        \"source_quantizers\": [\"fp32\"],\n",
    "        \"reference_internal\": \"int8\",\n",
    "        \"reference_accumulator\": \"int32\",\n",
    "    },\n",
    "}\n",
    "\n",
    "run_config = {\n",
    "    \"goal\": goal_energy,\n",
    "    \"quantization_config\": quantization_config,\n",
    "    \"learning_rate_optimizer\": False,\n",
    "    \"transfer_weights\": False,  # Randomely initialize weights\n",
    "    \"mode\": \"bayesian\",  # This can be bayesian,random,hyperband\n",
    "    \"seed\": 42,\n",
    "    \"limit\": limit,\n",
    "    \"tune_filters\": \"layer\",\n",
    "    \"tune_filters_exceptions\": \"^output\",\n",
    "    \"distribution_strategy\": None,\n",
    "    \"max_trials\": 5,  # Let's just do 5 trials for this demonstrator, ideally you should do as many as possible\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from qkeras.autoqkeras import AutoQKeras\n",
    "\n",
    "autoqk = AutoQKeras(baseline_model, output_dir=\"autoq_cnn\", metrics=[\"acc\"], custom_objects={}, **run_config)\n",
    "autoqk.fit(train_data, validation_data=val_data, epochs=15)\n",
    "\n",
    "aqmodel = autoqk.get_best_model()\n",
    "print_qmodel_summary(aqmodel)\n",
    "\n",
    "# Train for the full epochs\n",
    "callbacks = [\n",
    "    tf.keras.callbacks.EarlyStopping(patience=10, verbose=1),\n",
    "    tf.keras.callbacks.ReduceLROnPlateau(monitor='val_loss', factor=0.5, patience=3, verbose=1),\n",
    "]\n",
    "\n",
    "start = time.time()\n",
    "history = aqmodel.fit(train_data, epochs=n_epochs, validation_data=val_data, callbacks=callbacks, verbose=1)\n",
    "end = time.time()\n",
    "print('\\n It took {} minutes to train!\\n'.format((end - start) / 60.0))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# This model has some remnants from the optimization procedure attached to it, so let's define a new one\n",
    "aqmodel.save_weights(\"autoqkeras_cnn_weights.h5\")\n",
    "\n",
    "layers = [l for l in aqmodel.layers]\n",
    "x = layers[0].output\n",
    "for i in range(1, len(layers)):\n",
    "    x = layers[i](x)\n",
    "\n",
    "new_model = Model(inputs=[layers[0].input], outputs=[x])\n",
    "LOSS = tf.keras.losses.CategoricalCrossentropy()\n",
    "OPTIMIZER = tf.keras.optimizers.Adam(learning_rate=3e-3, beta_1=0.9, beta_2=0.999, epsilon=1e-07, amsgrad=True)\n",
    "\n",
    "new_model.compile(loss=LOSS, optimizer=OPTIMIZER, metrics=[\"accuracy\"])\n",
    "new_model.summary()\n",
    "new_model.load_weights(\"autoqkeras_cnn_weights.h5\")\n",
    "print_qmodel_summary(new_model)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's check what the best heterogeneously quantized model looks like (keep in mind we only did a few trials, the optimization obviosuly didn't have time to converge at the minimum but yo get the idea!)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "hls_config_aq = hls4ml.utils.config_from_keras_model(new_model, granularity='name')\n",
    "hls_config_aq['Model']['ReuseFactor'] = 1\n",
    "hls_config_aq['Model']['Precision'] = 'ap_fixed<16,6>'\n",
    "hls_config_aq['LayerName']['output_softmax']['Strategy'] = 'Stable'\n",
    "plotting.print_dict(hls_config_aq)\n",
    "\n",
    "cfg_aq = hls4ml.converters.create_config(backend='Vivado')\n",
    "cfg_aq['IOType'] = 'io_stream'  # Must set this if using CNNs!\n",
    "cfg_aq['HLSConfig'] = hls_config_aq\n",
    "cfg_aq['KerasModel'] = new_model\n",
    "cfg_aq['OutputDir'] = 'autoqkeras_cnn/'\n",
    "cfg_aq['XilinxPart'] = 'xcu250-figd2104-2L-e'\n",
    "\n",
    "hls_model_aq = hls4ml.converters.keras_to_hls(cfg_aq)\n",
    "hls_model_aq.compile()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "y_predict_aq = aqmodel.predict(X_test_reduced)\n",
    "y_predict_hls4ml_aq = hls_model_aq.predict(np.ascontiguousarray(X_test_reduced))\n",
    "\n",
    "\n",
    "accuracy_keras = float(accuracy_score(np.argmax(Y_test_reduced, axis=1), np.argmax(y_predict_aq, axis=1)))\n",
    "accuracy_hls4ml = float(accuracy_score(np.argmax(Y_test_reduced, axis=1), np.argmax(y_predict_hls4ml_aq, axis=1)))\n",
    "\n",
    "print(\"Accuracy AutoQ Keras:  {}\".format(accuracy_keras))\n",
    "print(\"Accuracy AutoQ hls4ml: {}\".format(accuracy_hls4ml))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The accuracy is slightly lower for this heterogeneously quantized model. Due to some randomness in the optimization procedure, you're going to have to synthesize this one yourself!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "synth = True\n",
    "if synth:\n",
    "    hls_model_aq.build(csim=False, synth=True, vsynth=True)\n",
    "    data_autoq = getReports('autoq_cnn')\n",
    "\n",
    "    print(\"\\n Resource usage and latency: AutoQ\")\n",
    "    pprint.pprint(data_autoq)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "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.11.11"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}