rnn-unrolling-tflite/RNN_models_for_Arm_Ethos_U55_U65.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#  Copyright (c) 2021 Arm Limited. All rights reserved.\n",
    "#  SPDX-License-Identifier: Apache-2.0\n",
    "#\n",
    "#  Licensed under the Apache License, Version 2.0 (the \"License\");\n",
    "#  you may not use this file except in compliance with the License.\n",
    "#  You may obtain a copy of the License at\n",
    "#\n",
    "#      http://www.apache.org/licenses/LICENSE-2.0\n",
    "#\n",
    "#  Unless required by applicable law or agreed to in writing, software\n",
    "#  distributed under the License is distributed on an \"AS IS\" BASIS,\n",
    "#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
    "#  See the License for the specific language governing permissions and\n",
    "#  limitations under the License."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Overview"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Recurrent neural networks (RNN) are a type of neural network that are very good for working with sequential data or data that can be expressed in a sequence. They work by processing multiple inputs sequentially in time, and have an internal state that allows them to encode information about the inputs they have previously seen. Importantly, all the weights of the RNN are shared between each step. Some popular usecases for recurrent neural networks include: machine translation, speech synthesis and time series prediction. Many of the target use cases are ones that are perfect for running on small embedded platforms, and as such can benefit from acceleration by running on specialist hardware such as a Neural Processing Unit (NPU)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<img src=\"resources/rnn.png\">\n",
    "<h6 align=\"center\">Figure 1. Layout of a simple Recurrent Neural Network for t timesteps.</h6>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In this guide we will show you how to use TensorFlow to train and quantize a simple recurrent neural network, and also to export it into a format that is compatible with Arm's embedded NPU the [Arm Ethos-U55](https://www.arm.com/products/silicon-ip-cpu/ethos/ethos-u55) and [Arm Ethos-U65](https://www.arm.com/products/silicon-ip-cpu/ethos/ethos-u65)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Before you begin"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "You should ensure you have Python3 installed on your machine.   \n",
    "\n",
    "Running the following commands will install the libraries necessary to run the code sample."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip3 install --upgrade pip3\n",
    "!pip3 install tensorflow==2.5.0\n",
    "!pip3 install numpy==1.19.5\n",
    "!pip3 install ethos-u-vela==3.0.0"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Training"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To start with we are going to train a type of RNN called a Gated Recurrent Unit or GRU for short. This is a slightly more complex type of RNN that has features that allow it to hold on to important information from the past longer than a standard RNN.\n",
    "\n",
    "The task that we will use our GRU for is classifying handwritten digits using the MNIST dataset. Image classification is not something we would normally use an RNN for, but we choose to do here as it is a simple and common enough task that many people will be familiar with. To classify an image using our GRU we will treat each row of the 28x28 input images as one time step's worth of input to the GRU. This way we will have a sequence of 28 image rows to input to our GRU. After performing 28 time step's worth of calculations the GRU will have 'seen' the whole image and we can pass the final output to a fully connected layer that will do the classification for us."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<img src=\"resources/gru_network.png\">\n",
    "<h6 align=\"center\">Figure 2. Layout of the GRU model we will train.</h6>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import tensorflow as tf\n",
    "import numpy as np"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "1) Define a model consisting of a single GRU layer followed by a fully connected layer in TensorFlow. The following code shows how to do this:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "def rnn_model(time_steps):\n",
    "    model_input = tf.keras.Input(shape=(time_steps, 28), name='input')\n",
    "\n",
    "    gru_out = tf.keras.layers.GRU(units=256)(model_input)\n",
    "\n",
    "    model_output = tf.keras.layers.Dense(10, activation='softmax')(gru_out)\n",
    "\n",
    "    model = tf.keras.Model([model_input], [model_output])\n",
    "\n",
    "    return model"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "2) Next, load the data for training and testing the RNN model and normalize it with the following code:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "(x_train, y_train), (x_test, y_test) = tf.keras.datasets.mnist.load_data()\n",
    "\n",
    "# Normalize between [0, 1] to help training.\n",
    "x_train = x_train / 255.0  \n",
    "x_test = x_test / 255.0"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "3) Create an instance of our previously defined model using the following code:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "model = rnn_model(time_steps=28)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The time steps is set to be twenty eight, which is the number of rows in each image."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "4) Compile the model and then train it using the standard 'fit' method call. To save time we train only for one epoch, training for longer will of course give a more accurate model. This is shown in the code below:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1875/1875 [==============================] - 21s 5ms/step - loss: 0.6053 - accuracy: 0.7912 - val_loss: 0.0945 - val_accuracy: 0.9700\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "<tensorflow.python.keras.callbacks.History at 0x7ff2d455b780>"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model.compile(optimizer=tf.keras.optimizers.Adam(learning_rate=0.001),\n",
    "              loss=tf.keras.losses.SparseCategoricalCrossentropy(),\n",
    "              metrics=['accuracy'])\n",
    "model.fit(x=x_train, y=y_train, epochs=1, validation_data=(x_test, y_test))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "5) After training for only one epoch we reach an accuracy of around 96-98% on the test dataset, not a bad result considering this is not an ideal task for an RNN. Once training is finished we save the weights of our model out to file so we can use them again later with the following code:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "model.save_weights('gru_mnist')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Deployment to TFLite"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "6) Now we have trained model weights we will define a slightly different graph to use in deployment. This graph will be more optimal for inference and conversion to TFLite format - and ultimately allows the model to be quantized and run on Arm Ethos-U55 or Arm Ethos-U65. The following code example shows the new graph definition in TensorFlow:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "def rnn_model_tflite(time_steps):\n",
    "    # We need to specify batch_size=1 to get an optimal graph for inference.\n",
    "    input_node = tf.keras.Input(shape=(time_steps, 28), batch_size=1, name='input')\n",
    "\n",
    "    gru_out = tf.keras.layers.GRU(units=256, unroll=True)(input_node)\n",
    "\n",
    "    prediction = tf.keras.layers.Dense(10, activation='softmax')(gru_out)\n",
    "\n",
    "    model = tf.keras.Model([input_node], [prediction])\n",
    "\n",
    "    return model"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "It looks nearly identical to the previous definition but has two important changes. The first is setting the batch_size parameter on the input node to be one. As we are only planning to do single inferences at a time with our trained model, we can safely do this. Doing so will also have the benefit of removing any extra operations that may have to be added automatically to read the batch size or deal with batch sizes greater than one.\n",
    "\n",
    "The second change is in the parameters used when we define the GRU layer. This is the most important change and it involves setting unroll to be True. "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Unrolling it all"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Normally an RNN has no defined end to the number of timesteps that it can loop for. As such, it can work with any length of input given to it. However, for a lot of real life use cases we know ahead of time the number of time steps we will need our RNN to loop for. In our example we know that we need to have 28 time steps to account for the 28 rows in each input image. Setting unroll to be True will remove the loop from the RNN, like in Figure 3, and unroll it to the number of time steps 't' set by the input to the layer. The result will be a strictly feedforward neural network that can then be quantized and converted easily to TensorFlow Lite."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<img src=\"resources/unroll.png\">\n",
    "<h6 align=\"center\">Figure 3. Unrolling of an RNN.</h6>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "7) We create an instance of our new 'deployment ready' model and load the previously saved weights into this model with the following code:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "WARNING:tensorflow:Layer gru_1 will not use cuDNN kernel since it doesn't meet the cuDNN kernel criteria. It will use generic GPU kernel as fallback when running on GPU\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "<tensorflow.python.training.tracking.util.CheckpointLoadStatus at 0x7ff2c85cb668>"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model_for_tflite = rnn_model_tflite(time_steps=28)\n",
    "model_for_tflite.load_weights('gru_mnist')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Quantize and TFLite"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now that we have our optimal model graph created, and the weights from the previous training have been loaded we want to go ahead and quantize it. Quantization is an essential step when deploying our model on small edge devices as it reduces model size and can speed up inference times. Moreover, NPUs like the Arm Ethos-U55 or Arm Ethos-U65 only support inference on models that have been quantized to 8-bits.\n",
    "\n",
    "The following steps on quantizing the model follow from the public TensorFlow tutorials on [post-training quantization](https://www.tensorflow.org/lite/performance/post_training_quantization#full_integer_quantization)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "8) As required for post-training quantization we define a generator function that will yield images from the training dataset. The following code shows how this is done:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "def rep_dataset():\n",
    "    for i in range(50):  # Only need a few examples.\n",
    "        img = x_train[i].astype(np.float32)\n",
    "        img = np.expand_dims(img, 0)\n",
    "        yield [img]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Then we load the new keras model into the TFLiteConverter, set all the attributes for post-training quantization and convert the model to TensorFlow Lite format. Finally we save the model to file ready for optimization with Vela and deployment on Arm Ethos-U55 or Arm Ethos-U65. These steps are shown in the code below:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "351008"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "converter_quant = tf.lite.TFLiteConverter.from_keras_model(model_for_tflite)\n",
    "converter_quant.optimizations = [tf.lite.Optimize.DEFAULT]\n",
    "converter_quant.representative_dataset = rep_dataset\n",
    "\n",
    "converter_quant.target_spec.supported_ops = [tf.lite.OpsSet.TFLITE_BUILTINS_INT8]\n",
    "converter_quant.inference_input_type = tf.int8\n",
    "converter_quant.inference_output_type = tf.int8\n",
    "\n",
    "tflite_model = converter_quant.convert()\n",
    "open('gru_mnist.tflite', 'wb').write(tflite_model)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Testing quantized accuracy"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As always, when quantizing your neural network model there can be some accuracy loss experienced. You should therefore check the accuracy of your model after quantization. Doing this for our GRU model we see that the accuracy has stayed virtually the same which is great news!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Test quantized TFLite model on the test set.\n",
    "interpreter = tf.lite.Interpreter(model_path='gru_mnist.tflite')\n",
    "interpreter.allocate_tensors()\n",
    "\n",
    "input_details = interpreter.get_input_details()\n",
    "output_details = interpreter.get_output_details()\n",
    "\n",
    "print(\"Testing quantized TFLite model...\")\n",
    "num_correct = 0\n",
    "for data, label in zip(x_test, y_test):\n",
    "    image = ((data * 255) - 128).astype(np.int8)  # Change data to int8 range to match quantized model.\n",
    "\n",
    "    interpreter.set_tensor(input_details[0]['index'], np.expand_dims(image, 0))\n",
    "    interpreter.invoke()\n",
    "\n",
    "    outputs = interpreter.get_tensor(output_details[0]['index'])\n",
    "\n",
    "    predicted = np.argmax(outputs)\n",
    "    num_correct += np.sum(predicted == label)\n",
    "\n",
    "print(f\"Accuracy of quantized TFLite model on test set: {100 * num_correct / len(y_test)}%\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Feeding and getting hidden states"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "It might be the case for your specific model you want to manually set the initial hidden states or retrieve the final hidden states from your RNN layers. Luckily this is simple enough to do and just needs an alteration to how you define the model when exporting to TensorFlow Lite. For instance, to get a TensorFlow Lite model that has access to these input and output states we will simply define our model like in the code snippet below:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def rnn_model_tflite_with_state_transfer(time_steps):\n",
    "    # We need to specify batch_size=1 to get an optimal graph for inference.\n",
    "    input_node = tf.keras.Input(shape=(time_steps, 28), batch_size=1, name='input')\n",
    "\n",
    "    # Input tensor for the initial GRU hidden state.\n",
    "    initial_gru_state = tf.keras.Input(shape=256, batch_size=1, name='initial_gru_state')\n",
    "\n",
    "    # Return the final GRU hidden state.\n",
    "    gru_out, final_gru_state = tf.keras.layers.GRU(units=256, unroll=True, return_state=True)(input_node, initial_state=initial_gru_state)\n",
    "\n",
    "    prediction = tf.keras.layers.Dense(10, activation='softmax')(gru_out)\n",
    "\n",
    "    model = tf.keras.Model([input_node, initial_gru_state], [prediction, final_gru_state])\n",
    "\n",
    "    return model"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In this model defintion we have now defined a new Input tensor that is the size of the GRU hidden layer, this is how we will feed an initial GRU state. In the GRU layer we set return_state to be True so that we can capture the final hidden state. When we call the layer we provide the initial state that we defined in the previous line. Finally, when making the keras Model we supply this additional input and output to the corresponding list of input and output nodes.\n",
    "\n",
    "If you were using an LSTM layer instead of a basic RNN or GRU layer then the code would need only a small modification from above. You would need an additional Input tensor for the cell state and you would also need to capture the additional final cell state. These would also need to be added to the list of inputs and outputs when making the TensorFlow keras Model."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**NOTE: When doing post-training quantization on a model with this state transfer enabled you will need to supply example data for these input hidden states in your representative data generator function. If your usecase involves feeding final states as new initial states to your RNN, then it is advised to collect examples of these final states and iteratively add them to your represenative dataset of input states.**"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Vela"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The final step in preparing the model for deployment on Arm Ethos-U55 or Arm Ethos-U65 is to run the quantized TensorFlow Lite file through the [Vela compiler](https://pypi.org/project/ethos-u-vela/).\n",
    "\n",
    "Please see instruction on the Vela page on how to run your model through the Vela optimizer. After following the steps in this guide and running your model through Vela you should find yourself with a model that can be completely run on Arm Ethos-U55 or Arm Ethos-U65."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Conclusion"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "By following this guide, you successfully trained a simple RNN based model in TensorFlow. You have then seen how to unroll your RNN model so that is can be quantized and converted easily to TensorFlow Lite format ready for deploying on Arm Ethos-U55 or Arm Ethos-U65. You have also seen how it is possible to create a TFLite model where an initial RNN state can be provided as input to the model and how to capture final RNN states from the model."
   ]
  }
 ],
 "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.6.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}