train_lstm_timeseries_ptq_tf2.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "license"
      },
      "source": [
        "##### *Copyright 2021 Google LLC*\n",
        "*Licensed under the Apache License, Version 2.0 (the \"License\")*"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "both",
        "id": "rKwqeqWBXANA"
      },
      "outputs": [],
      "source": [
        "# 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",
        "# https://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": {
        "id": "ooTA7e2BSeQ8"
      },
      "source": [
        "# Train an LSTM weather forecasting model for the Coral Edge TPU\n",
        "\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "AjUp9jtWoukX"
      },
      "source": [
        "This tutorial shows how you can create an LSTM time series model that's compatible with the Edge TPU (available in [Coral devices](https://coral.ai/products/)).\n",
        "\n",
        "This notebook is based on the [Keras timeseries forecasting tutorial](https://keras.io/examples/timeseries/timeseries_weather_forecasting/). We've mostly just added code to quantize the model with TensorFlow Lite and compile it for the Edge TPU.\n",
        "\n",
        "Here's an example of the results:\n",
        "\n",
        "\u003cimg src=\"https://storage.googleapis.com/site_and_emails_static_assets/Images/lstm-forecast.png?\" width=\"400\" hspace=\"0\"\u003e\n",
        "\n",
        "\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "zHZyMPbMsAns"
      },
      "source": [
        "\u003ca href=\"https://colab.research.google.com/github/google-coral/tutorials/blob/master/train_lstm_timeseries_ptq_tf2.ipynb\" target=\"_parent\"\u003e\u003cimg src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open in Colab\"\u003e\u003c/a\u003e\n",
        "\u0026nbsp;\u0026nbsp;\u0026nbsp;\u0026nbsp;\n",
        "\u003ca href=\"https://github.com/google-coral/tutorials/blob/master/train_lstm_timeseries_ptq_tf2.ipynb\" target=\"_parent\"\u003e\u003cimg src=\"https://img.shields.io/static/v1?logo=GitHub\u0026label=\u0026color=333333\u0026style=flat\u0026message=View%20on%20GitHub\" alt=\"View in GitHub\"\u003e\u003c/a\u003e\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "PzoNZRp4sVxK"
      },
      "source": [
        "To start running all the code in this tutorial, select **Runtime \u003e Run all** in the Colab toolbar."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "pPSeVqU0SeQ-"
      },
      "source": [
        "## Import the required libraries"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "-Ecb5t7qSeQ_"
      },
      "outputs": [],
      "source": [
        "import pandas as pd\n",
        "import matplotlib.pyplot as plt\n",
        "import tensorflow as tf\n",
        "from tensorflow import keras"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "a9uGCJGsjJpc"
      },
      "outputs": [],
      "source": [
        "assert float(tf.__version__[:3]) \u003e= 2.3"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "5Vk9iwhfrP_C"
      },
      "source": [
        "## Prepare the climate dataset"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "NuO24ZxkSeQ_"
      },
      "source": [
        "We will be using Jena Climate dataset recorded by the\n",
        "[Max Planck Institute for Biogeochemistry](https://www.bgc-jena.mpg.de/wetter/).\n",
        "The dataset consists of 14 features such as temperature, pressure, humidity etc, recorded once per\n",
        "10 minutes.\n",
        "\n",
        "**Location**: Weather Station, Max Planck Institute for Biogeochemistry\n",
        "in Jena, Germany\n",
        "\n",
        "**Time-frame Considered**: Jan 10, 2009 - December 31, 2016\n",
        "\n",
        "\n",
        "The table below shows the column names, their value formats, and their description.\n",
        "\n",
        "Index| Features      |Format             |Description\n",
        "-----|---------------|-------------------|-----------------------\n",
        "1    |Date Time      |01.01.2009 00:10:00|Date-time reference\n",
        "2    |p (mbar)       |996.52             |The pascal SI derived unit of pressure used to quantify internal pressure. Meteorological reports typically state atmospheric pressure in millibars.\n",
        "3    |T (degC)       |-8.02              |Temperature in Celsius\n",
        "4    |Tpot (K)       |265.4              |Temperature in Kelvin\n",
        "5    |Tdew (degC)    |-8.9               |Temperature in Celsius relative to humidity. Dew Point is a measure of the absolute amount of water in the air, the DP is the temperature at which the air cannot hold all the moisture in it and water condenses.\n",
        "6    |rh (%)         |93.3               |Relative Humidity is a measure of how saturated the air is with water vapor, the %RH determines the amount of water contained within collection objects.\n",
        "7    |VPmax (mbar)   |3.33               |Saturation vapor pressure\n",
        "8    |VPact (mbar)   |3.11               |Vapor pressure\n",
        "9    |VPdef (mbar)   |0.22               |Vapor pressure deficit\n",
        "10   |sh (g/kg)      |1.94               |Specific humidity\n",
        "11   |H2OC (mmol/mol)|3.12               |Water vapor concentration\n",
        "12   |rho (g/m ** 3) |1307.75            |Airtight\n",
        "13   |wv (m/s)       |1.03               |Wind speed\n",
        "14   |max. wv (m/s)  |1.75               |Maximum wind speed\n",
        "15   |wd (deg)       |152.3              |Wind direction in degrees"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "MOrP2N3jSeRA"
      },
      "outputs": [],
      "source": [
        "from zipfile import ZipFile\n",
        "import os\n",
        "\n",
        "uri = \"https://storage.googleapis.com/tensorflow/tf-keras-datasets/jena_climate_2009_2016.csv.zip\"\n",
        "zip_path = keras.utils.get_file(origin=uri, fname=\"jena_climate_2009_2016.csv.zip\")\n",
        "zip_file = ZipFile(zip_path)\n",
        "zip_file.extractall()\n",
        "csv_path = \"jena_climate_2009_2016.csv\"\n",
        "\n",
        "df = pd.read_csv(csv_path)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "RFyJ8ZwrSeRB"
      },
      "source": [
        "### Visualize the data\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "nvsOYHklrMkx"
      },
      "source": [
        "To give us a sense of the data we are working with, each feature has been plotted below.\n",
        "This shows the distinct pattern of each feature over the time period from 2009 to 2016.\n",
        "It also shows where anomalies are present, which will be addressed during normalization."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "sKPSF7V2SeRB"
      },
      "outputs": [],
      "source": [
        "titles = [\n",
        "    \"Pressure\",\n",
        "    \"Temperature\",\n",
        "    \"Temperature in Kelvin\",\n",
        "    \"Temperature (dew point)\",\n",
        "    \"Relative Humidity\",\n",
        "    \"Saturation vapor pressure\",\n",
        "    \"Vapor pressure\",\n",
        "    \"Vapor pressure deficit\",\n",
        "    \"Specific humidity\",\n",
        "    \"Water vapor concentration\",\n",
        "    \"Airtight\",\n",
        "    \"Wind speed\",\n",
        "    \"Maximum wind speed\",\n",
        "    \"Wind direction in degrees\",\n",
        "]\n",
        "\n",
        "feature_keys = [\n",
        "    \"p (mbar)\",\n",
        "    \"T (degC)\",\n",
        "    \"Tpot (K)\",\n",
        "    \"Tdew (degC)\",\n",
        "    \"rh (%)\",\n",
        "    \"VPmax (mbar)\",\n",
        "    \"VPact (mbar)\",\n",
        "    \"VPdef (mbar)\",\n",
        "    \"sh (g/kg)\",\n",
        "    \"H2OC (mmol/mol)\",\n",
        "    \"rho (g/m**3)\",\n",
        "    \"wv (m/s)\",\n",
        "    \"max. wv (m/s)\",\n",
        "    \"wd (deg)\",\n",
        "]\n",
        "\n",
        "colors = [\n",
        "    \"blue\",\n",
        "    \"orange\",\n",
        "    \"green\",\n",
        "    \"red\",\n",
        "    \"purple\",\n",
        "    \"brown\",\n",
        "    \"pink\",\n",
        "    \"gray\",\n",
        "    \"olive\",\n",
        "    \"cyan\",\n",
        "]\n",
        "\n",
        "date_time_key = \"Date Time\"\n",
        "\n",
        "def show_raw_visualization(data):\n",
        "    time_data = data[date_time_key]\n",
        "    fig, axes = plt.subplots(\n",
        "        nrows=7, ncols=2, figsize=(15, 20), dpi=80, facecolor=\"w\", edgecolor=\"k\"\n",
        "    )\n",
        "    for i in range(len(feature_keys)):\n",
        "        key = feature_keys[i]\n",
        "        c = colors[i % (len(colors))]\n",
        "        t_data = data[key]\n",
        "        t_data.index = time_data\n",
        "        t_data.head()\n",
        "        ax = t_data.plot(\n",
        "            ax=axes[i // 2, i % 2],\n",
        "            color=c,\n",
        "            title=\"{} - {}\".format(titles[i], key),\n",
        "            rot=25,\n",
        "        )\n",
        "        ax.legend([titles[i]])\n",
        "    plt.tight_layout()\n",
        "\n",
        "\n",
        "show_raw_visualization(df)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "VLSMHguASeRC"
      },
      "source": [
        "This heat map shows the correlation between different features."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Z9gL0fmZSeRD"
      },
      "outputs": [],
      "source": [
        "def show_heatmap(data):\n",
        "    plt.matshow(data.corr())\n",
        "    plt.xticks(range(data.shape[1]), data.columns, fontsize=14, rotation=90)\n",
        "    plt.gca().xaxis.tick_bottom()\n",
        "    plt.yticks(range(data.shape[1]), data.columns, fontsize=14)\n",
        "\n",
        "    cb = plt.colorbar()\n",
        "    cb.ax.tick_params(labelsize=14)\n",
        "    plt.title(\"Feature Correlation Heatmap\", fontsize=14)\n",
        "    plt.show()\n",
        "\n",
        "show_heatmap(df)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "eKAw-W7qrZCa"
      },
      "source": [
        "### Preprocess the dataset"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "duK_boN3SeRE"
      },
      "source": [
        "Here we are picking ~300,000 data points for training. Observation is recorded every\n",
        "10 mins, that means 6 times per hour. We will resample one point per hour since no\n",
        "drastic change is expected within 60 minutes. We do this via the `sampling_rate`\n",
        "argument in `timeseries_dataset_from_array` utility.\n",
        "\n",
        "We are tracking data from past 720 timestamps (720/6=120 hours). This data will be\n",
        "used to predict the temperature after 72 timestamps (72/6=12 hours).\n",
        "\n",
        "Since every feature has values with\n",
        "varying ranges, we do normalization to confine feature values to a range of `[0, 1]` before\n",
        "training a neural network.\n",
        "We do this by subtracting the mean and dividing by the standard deviation of each feature.\n",
        "\n",
        "71.5 % of the data will be used to train the model, i.e. 300,693 rows. `split_fraction` can\n",
        "be changed to alter this percentage.\n",
        "\n",
        "The model is shown data for first 5 days i.e. 720 observations, that are sampled every\n",
        "hour. The temperature after 72 (12 hours * 6 observation per hour) observation will be\n",
        "used as a label."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "qzTTe8XnSeRE"
      },
      "outputs": [],
      "source": [
        "split_fraction = 0.715\n",
        "train_split = int(split_fraction * int(df.shape[0]))\n",
        "step = 6\n",
        "\n",
        "past = 720\n",
        "future = 72\n",
        "learning_rate = 0.001\n",
        "batch_size = 256\n",
        "\n",
        "def normalize(data, train_split):\n",
        "    data_mean = data[:train_split].mean(axis=0)\n",
        "    data_std = data[:train_split].std(axis=0)\n",
        "    return (data - data_mean) / data_std"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "MnhDuT5wSeRF"
      },
      "source": [
        "We can see from the correlation heatmap, few parameters like Relative Humidity and\n",
        "Specific Humidity are redundant. Hence we will be using select features, not all."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "dJTTmv8cSeRF"
      },
      "outputs": [],
      "source": [
        "print(\n",
        "    \"The selected parameters are:\",\n",
        "    \", \".join([titles[i] for i in [0, 1, 5, 7, 8, 10, 11]]),\n",
        ")\n",
        "selected_features = [feature_keys[i] for i in [0, 1, 5, 7, 8, 10, 11]]\n",
        "features = df[selected_features]\n",
        "features.index = df[date_time_key]\n",
        "features.head()\n",
        "\n",
        "features = normalize(features.values, train_split)\n",
        "features = pd.DataFrame(features)\n",
        "features.head()\n",
        "\n",
        "train_data = features.loc[0 : train_split - 1]\n",
        "val_data = features.loc[train_split:]"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "dWKybz60sNre"
      },
      "source": [
        "### Split the training and validation dataset"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "LwzOGwWrSeRG"
      },
      "source": [
        "The training dataset labels starts from the 792nd observation (720 + 72)."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "6M4EIkFaSeRG"
      },
      "outputs": [],
      "source": [
        "start = past + future\n",
        "end = start + train_split\n",
        "\n",
        "x_train = train_data[[i for i in range(7)]].values\n",
        "y_train = features.iloc[start:end][[1]]\n",
        "print('training size:', len(x_train))\n",
        "\n",
        "sequence_length = int(past / step)\n",
        "print('window size:', sequence_length)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "WO9ioCrCSeRG"
      },
      "source": [
        "The `timeseries_dataset_from_array` function takes in a sequence of data-points gathered at\n",
        "equal intervals, along with time series parameters such as length of the\n",
        "sequences/windows, spacing between two sequence/windows, etc., to produce batches of\n",
        "sub-timeseries inputs and targets sampled from the main time series."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "-UksQgJySeRG"
      },
      "outputs": [],
      "source": [
        "dataset_train = keras.preprocessing.timeseries_dataset_from_array(\n",
        "    x_train,\n",
        "    y_train,\n",
        "    sequence_length=sequence_length,\n",
        "    sampling_rate=step,\n",
        "    batch_size=batch_size,\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "uEV2EybxSeRH"
      },
      "source": [
        "The validation dataset must not contain the last 792 rows as we won't have label data for\n",
        "those records, hence 792 must be subtracted from the end of the data.\n",
        "\n",
        "The validation label dataset must start from 792 after train_split, hence we must add\n",
        "past + future (792) to label_start."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "C44Bq1KUSeRH"
      },
      "outputs": [],
      "source": [
        "x_end = len(val_data) - past - future\n",
        "\n",
        "label_start = train_split + past + future\n",
        "\n",
        "x_val = val_data.iloc[:x_end][[i for i in range(7)]].values\n",
        "y_val = features.iloc[label_start:][[1]]\n",
        "\n",
        "dataset_val = keras.preprocessing.timeseries_dataset_from_array(\n",
        "    x_val,\n",
        "    y_val,\n",
        "    sequence_length=sequence_length,\n",
        "    sampling_rate=step,\n",
        "    batch_size=batch_size,\n",
        ")\n",
        "\n",
        "for batch in dataset_train.take(1):\n",
        "    inputs, targets = batch\n",
        "\n",
        "print(\"Input shape:\", inputs.numpy().shape)\n",
        "print(\"Target shape:\", targets.numpy().shape)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "GJvaQK128exu"
      },
      "source": [
        "The input shape is the above-defined `batch_size`, the sequence length (120 hours), and the 7 selected features.\n",
        "\n",
        "When we convert the model for the Edge TPU, we'll have to change this batch size to 1 for compiler compatibility."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "5JpFyDJWSeRH"
      },
      "source": [
        "## Build and train the LSTM model"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "M58YjCAqSeRH"
      },
      "outputs": [],
      "source": [
        "inputs = keras.layers.Input(shape=(inputs.shape[1], inputs.shape[2]))\n",
        "lstm_out = keras.layers.LSTM(32)(inputs)\n",
        "outputs = keras.layers.Dense(1)(lstm_out)\n",
        "\n",
        "model = keras.Model(inputs=inputs, outputs=outputs)\n",
        "model.compile(optimizer=keras.optimizers.Adam(learning_rate=learning_rate), loss='mse')\n",
        "model.summary()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "P6R-Kq4USeRI"
      },
      "source": [
        "We'll use the `ModelCheckpoint` callback to regularly save checkpoints, and\n",
        "the `EarlyStopping` callback to interrupt training when the validation loss\n",
        "is not longer improving."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Hvg9m237lHmm"
      },
      "outputs": [],
      "source": [
        "epochs = 10"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "zf9f14n8SeRI"
      },
      "outputs": [],
      "source": [
        "path_checkpoint = \"model_checkpoint.h5\"\n",
        "es_callback = keras.callbacks.EarlyStopping(monitor=\"val_loss\", min_delta=0, patience=5)\n",
        "\n",
        "modelckpt_callback = keras.callbacks.ModelCheckpoint(\n",
        "    monitor=\"val_loss\",\n",
        "    filepath=path_checkpoint,\n",
        "    verbose=1,\n",
        "    save_weights_only=True,\n",
        "    save_best_only=True,\n",
        ")\n",
        "\n",
        "history = model.fit(\n",
        "    dataset_train,\n",
        "    epochs=epochs,\n",
        "    validation_data=dataset_val,\n",
        "    callbacks=[es_callback, modelckpt_callback],\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "6qO3pJejSeRI"
      },
      "source": [
        "We can visualize the loss with the function below. After one point, the loss stops\n",
        "decreasing."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "PJbqa8usSeRJ"
      },
      "outputs": [],
      "source": [
        "def visualize_loss(history, title):\n",
        "    loss = history.history[\"loss\"]\n",
        "    val_loss = history.history[\"val_loss\"]\n",
        "    epochs = range(len(loss))\n",
        "    plt.figure()\n",
        "    plt.plot(epochs, loss, \"b\", label=\"Training loss\")\n",
        "    plt.plot(epochs, val_loss, \"r\", label=\"Validation loss\")\n",
        "    plt.title(title)\n",
        "    plt.xlabel(\"Epochs\")\n",
        "    plt.ylabel(\"Loss\")\n",
        "    plt.legend()\n",
        "    plt.show()\n",
        "    print(\"Final val loss: \", val_loss)\n",
        "\n",
        "visualize_loss(history, \"Training and Validation Loss\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "29JK8UlNsoFl"
      },
      "source": [
        "## Plot some predictions"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "7SzkExG3SeRJ"
      },
      "source": [
        "Now let's see what some of the predictions look like from a new test set. We create a new test dataset (even though it's has the same data as the validation set) because we want a batch size of just 1 (that's required when we compile for the Edge TPU). We also set the `sequence_stride` to something bigger than the sequence/window length, so our test samples aren't almost identical."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "w5YebzWefjw2"
      },
      "outputs": [],
      "source": [
        "dataset_test = keras.preprocessing.timeseries_dataset_from_array(\n",
        "    x_val,\n",
        "    y_val,\n",
        "    sequence_length=sequence_length,\n",
        "    sequence_stride=int(sequence_length * 6),\n",
        "    sampling_rate=step,\n",
        "    batch_size=1,\n",
        ")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "PJc7m1j2SeRJ"
      },
      "outputs": [],
      "source": [
        "def show_plot(plot_data, delta, title):\n",
        "    labels = [\"History\", \"True Future\", \"Model Prediction\"]\n",
        "    marker = [\".-\", \"rx\", \"go\"]\n",
        "    time_steps = list(range(-(plot_data[0].shape[0]), 0))\n",
        "    if delta:\n",
        "        future = delta\n",
        "    else:\n",
        "        future = 0\n",
        "\n",
        "    plt.title(title)\n",
        "    for i, val in enumerate(plot_data):\n",
        "        if i:\n",
        "            plt.plot(future, plot_data[i], marker[i], markersize=10, label=labels[i])\n",
        "        else:\n",
        "            plt.plot(time_steps, plot_data[i].flatten(), marker[i], label=labels[i])\n",
        "    plt.legend()\n",
        "    plt.xlim([time_steps[0], (future + 5) * 2])\n",
        "    plt.xlabel(\"Time-Step\")\n",
        "    plt.show()\n",
        "    return\n",
        "\n",
        "for x, y in dataset_test.take(5):\n",
        "    prediction = model.predict(x)\n",
        "    prediction = prediction[0]\n",
        "    print('prediction:', prediction)\n",
        "    print('truth:', y[0].numpy())\n",
        "    show_plot(\n",
        "        [x[0][:, 1].numpy(), y[0], prediction],\n",
        "        12,\n",
        "        \"Single Step Prediction\",\n",
        "    )"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "gmq8e4AQU9SH"
      },
      "source": [
        "## Convert to TF Lite"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "U4-ggEmAfMJl"
      },
      "outputs": [],
      "source": [
        "converter = tf.lite.TFLiteConverter.from_keras_model(model)\n",
        "tflite_model = converter.convert()\n",
        "\n",
        "with open('weather_forecast.tflite', 'wb') as f:\n",
        "  f.write(tflite_model)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "SwwaB09occEh"
      },
      "source": [
        "That gives us a basic TF Lite version, which can be useful for benchmarks, but we need it fully-quantized for compatibility with the Edge TPU..."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "86JNhH3z_Eky"
      },
      "source": [
        "### Create quantized version"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "7PwqHWWCfD1c"
      },
      "source": [
        "First, we need to update the model's input shape to a fixed size, as dictated by the [Edge TPU model requirements](https://coral.ai/docs/edgetpu/models-intro/#model-requirements)."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Ark015gJeExM"
      },
      "outputs": [],
      "source": [
        "batch_size = 1\n",
        "model.input.set_shape((batch_size,) + model.input.shape[1:])\n",
        "model.summary()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "x6QQiy8y5FvU"
      },
      "outputs": [],
      "source": [
        "# Our representative dataset is the same as the training dataset,\n",
        "# but the batch size must now be 1\n",
        "dataset_repr = keras.preprocessing.timeseries_dataset_from_array(\n",
        "    x_train,\n",
        "    y_train,\n",
        "    sequence_length=sequence_length,\n",
        "    sampling_rate=step,\n",
        "    batch_size=batch_size,\n",
        ")\n",
        "\n",
        "def representative_data_gen():\n",
        "  # To ensure full coverage of possible inputs, we use the whole train set\n",
        "  for input_data, _ in dataset_repr.take(int(len(x_train))):\n",
        "    input_data = tf.cast(input_data, dtype=tf.float32)\n",
        "    yield [input_data]"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "JepjY4VubkXQ"
      },
      "source": [
        "This part takes several minutes due to the size of the representative dataset:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "EBRDh9SZVBX1"
      },
      "outputs": [],
      "source": [
        "converter = tf.lite.TFLiteConverter.from_keras_model(model)\n",
        "# This enables quantization\n",
        "converter.optimizations = [tf.lite.Optimize.DEFAULT]\n",
        "# This sets the representative dataset for quantization\n",
        "converter.representative_dataset = representative_data_gen\n",
        "# This ensures that if any ops can't be quantized, the converter throws an error\n",
        "converter.target_spec.supported_ops = [tf.lite.OpsSet.TFLITE_BUILTINS_INT8]\n",
        "# For full integer quantization, though supported types defaults to int8 only, we explicitly declare it for clarity\n",
        "converter.target_spec.supported_types = [tf.int8]\n",
        "# These set the input and output tensors to int8\n",
        "converter.inference_input_type = tf.uint8\n",
        "converter.inference_output_type = tf.uint8\n",
        "tflite_model_quant = converter.convert()\n",
        "\n",
        "with open('weather_forecast_quant.tflite', 'wb') as f:\n",
        "  f.write(tflite_model_quant)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "kQJFmfHQma1n"
      },
      "source": [
        "### Try some TF Lite predictions"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "DOQjdt1_Oi0-"
      },
      "outputs": [],
      "source": [
        "import numpy as np\n",
        "\n",
        "def set_input_tensor(interpreter, input):\n",
        "  input_details = interpreter.get_input_details()[0]\n",
        "  tensor_index = input_details['index']\n",
        "  input_tensor = interpreter.tensor(tensor_index)()\n",
        "  # Inputs for the TFLite model must be uint8, so we quantize our input data.\n",
        "  scale, zero_point = input_details['quantization']\n",
        "  quantized_input = np.uint8(input / scale + zero_point)\n",
        "  input_tensor[:, :, :] = quantized_input\n",
        "\n",
        "def predict_weather(interpreter, input):\n",
        "  set_input_tensor(interpreter, input)\n",
        "  interpreter.invoke()\n",
        "  output_details = interpreter.get_output_details()[0]\n",
        "  output = interpreter.get_tensor(output_details['index'])\n",
        "  # Outputs from the TFLite model are uint8, so we dequantize the results:\n",
        "  scale, zero_point = output_details['quantization']\n",
        "  output = scale * (output - zero_point)\n",
        "  return output\n",
        "\n",
        "interpreter = tf.lite.Interpreter('weather_forecast_quant.tflite')\n",
        "interpreter.allocate_tensors()\n",
        "\n",
        "for x, y in dataset_test.take(5):\n",
        "  prediction = predict_weather(interpreter, x)\n",
        "  print('prediction:', prediction[0])\n",
        "  print('truth:', y[0].numpy())\n",
        "\n",
        "  show_plot(\n",
        "      [x[0][:, 1].numpy(), y[0], prediction[0]],\n",
        "      12,\n",
        "      \"Single Step Prediction (TF Lite)\",\n",
        "  )"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ip3NBzgViVun"
      },
      "source": [
        "If you compare these predictions from the quantized model to those we got from the float Keras model above, they're not very different."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "BdwfJCVDmEx6"
      },
      "source": [
        "## Compile for the Edge TPU"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "xHL3ByrR-5Qd"
      },
      "outputs": [],
      "source": [
        "! curl https://packages.cloud.google.com/apt/doc/apt-key.gpg | sudo apt-key add -\n",
        "\n",
        "! echo \"deb https://packages.cloud.google.com/apt coral-edgetpu-stable main\" | sudo tee /etc/apt/sources.list.d/coral-edgetpu.list\n",
        "\n",
        "! sudo apt-get update\n",
        "\n",
        "! sudo apt-get install edgetpu-compiler\t\n",
        "\n",
        "! edgetpu_compiler weather_forecast_quant.tflite"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "uDzC-ouVj32-"
      },
      "source": [
        "You can find the compiled model in the **Files** window on the left, named `weather_forecast_quant_edgetput.tflite`."
      ]
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "collapsed_sections": [
        "license"
      ],
      "name": "Train an LSTM weather forecasting model for the Coral Edge TPU",
      "provenance": [],
      "toc_visible": true
    },
    "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.0"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}