ai-explanations-local-experience.ipynb

{
  "cells": [
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "# Copyright 2020 Google LLC\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",
        "#     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": {},
      "source": [
        "# AI Explanations: Local Experience\n",
        "\n",
        "<table align=\"left\">\n",
        "  <td>\n",
        "    <a href=\"https://colab.research.google.com/github/GoogleCloudPlatform/ai-platform-samples/tree/main/notebooks/samples/explanations/tf2/ai-explanations-local-experience.ipynb\">\n",
        "      <img src=\"https://cloud.google.com/ml-engine/images/colab-logo-32px.png\" alt=\"Colab logo\"> Run in Colab\n",
        "    </a>\n",
        "  </td>\n",
        "  <td>\n",
        "    <a href=\"https://github.com/GoogleCloudPlatform/ai-platform-samples/tree/main/notebooks/samples/explanations/tf2/ai-explanations-local-experience.ipynb\">\n",
        "      <img src=\"https://cloud.google.com/ml-engine/images/github-logo-32px.png\" alt=\"GitHub logo\">\n",
        "      View on GitHub\n",
        "    </a>\n",
        "  </td>\n",
        "</table>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Overview\n",
        "\n",
        "This tutorial shows how to generate explanations for a model deployed locally using the AI Platform Explanations service. The AI Platform Explanations service is used to get feature attributions on your deployed model --explanations on how the features that contributed to the prediction.\n",
        "\n",
        "The local experience feature is only accessible from an AI Platform notebook with Tensorflow 2.3 framework. It will not run in Colab or your laptop, or on n AIP instance with an earlier version of Tensorflow."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "### Dataset\n",
        "\n",
        "The dataset used for this tutorial is the [flowers dataset](https://www.tensorflow.org/datasets/catalog/tf_flowers) from [TensorFlow Datasets](https://www.tensorflow.org/datasets/catalog/overview)."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "### Objective\n",
        "\n",
        "The goal of this tutorial is to undetstand how to generate explanations for a model deployed locally using the AI Explanations service to obtain explanations for prediction from an image (flower) classification model. For image models, AI Explanations returns an image with the pixels highlighted that signaled your model's prediction the most.\n",
        "\n",
        "This tutorial focuses more on deploying the model to AI Platform with Explanations than on the design of the model itself. "
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "### Costs\n",
        "\n",
        "This tutorial uses billable components of Google Cloud Platform (GCP):\n",
        "\n",
        "* AI Platform for:\n",
        "  * Prediction\n",
        "  * Explanation: AI Explanations comes at no extra charge to prediction prices. However, explanation requests take longer to process than normal predictions, so heavy usage of Explanations along with auto-scaling may result in more nodes being started and thus more charges\n",
        "* Cloud Storage for:\n",
        "  * Storing model files for deploying to Cloud AI Platform\n",
        "\n",
        "Learn about [AI Platform\n",
        "pricing](https://cloud.google.com/ml-engine/docs/pricing) and [Cloud Storage\n",
        "pricing](https://cloud.google.com/storage/pricing), and use the [Pricing\n",
        "Calculator](https://cloud.google.com/products/calculator/)\n",
        "to generate a cost estimate based on your projected usage."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Before you begin\n",
        "\n",
        "**Make sure you're running this notebook in a GPU runtime if you have that option. In Colab, select Runtime --> Change runtime type**\n",
        "\n",
        "This tutorial assumes you are running the notebook either in **Colab** or **Cloud AI Platform Notebooks**."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "### Set up your GCP project\n",
        "\n",
        "**The following steps are required, regardless of your notebook environment.**\n",
        "\n",
        "1. [Select or create a GCP project.](https://console.cloud.google.com/cloud-resource-manager)\n",
        "\n",
        "2. [Make sure that billing is enabled for your project.](https://cloud.google.com/billing/docs/how-to/modify-project)\n",
        "\n",
        "3. [Enable the AI Platform Training & Prediction and Compute Engine APIs.](https://console.cloud.google.com/flows/enableapi?apiid=ml.googleapis.com,compute_component)\n",
        "\n",
        "4. Enter your project ID in the cell below. Then run the  cell to make sure the\n",
        "Cloud SDK uses the right project for all the commands in this notebook.\n",
        "\n",
        "**Note**: Jupyter runs lines prefixed with `!` as shell commands, and it interpolates Python variables prefixed with `$` into these commands."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "#### Project ID\n",
        "\n",
        "**If you don't know your project ID.**\n",
        "\n",
        "You might able to get your project ID using `gcloud` command, by executing the second code cell below."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "PROJECT_ID = \"[your-project-id]\" #@param {type:\"string\"}"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "if PROJECT_ID == \"\" or PROJECT_ID is None or PROJECT_ID == \"[your-project-id]\":\n",
        "    # Get your GCP project id from gcloud\n",
        "    shell_output = !gcloud config list --format 'value(core.project)' 2>/dev/null\n",
        "    PROJECT_ID = shell_output[0]\n",
        "    print(\"Project ID:\", PROJECT_ID)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "#### Region\n",
        "\n",
        "You can also change the `REGION` variable, which is used for operations\n",
        "throughout the rest of this notebook. Make sure to [choose a region where Cloud\n",
        "AI Platform services are\n",
        "available](https://cloud.google.com/ml-engine/docs/tensorflow/regions). You can\n",
        "not use a Multi-Regional Storage bucket for training with AI Platform."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "REGION = 'us-central1' #@param {type: \"string\"}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "#### Timestamp\n",
        "\n",
        "If you are in a live tutorial session, you might be using a shared test account or project. To avoid name collisions between users on resources created, you create a timestamp for each instance session, and append onto the name of resources which will be created in this tutorial."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "from datetime import datetime\n",
        "\n",
        "TIMESTAMP = datetime.now().strftime(\"%Y%m%d%H%M%S\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "### Create a Cloud Storage bucket\n",
        "\n",
        "**The following steps are required, regardless of your notebook environment.**\n",
        "\n",
        "When you submit a training job using the Cloud SDK, you upload a Python package\n",
        "containing your training code to a Cloud Storage bucket. AI Platform runs\n",
        "the code from this package. In this tutorial, AI Platform also saves the\n",
        "trained model that results from your job in the same bucket. You can then\n",
        "create an AI Platform model version based on this output in order to serve\n",
        "online predictions.\n",
        "\n",
        "Set the name of your Cloud Storage bucket below. It must be unique across all\n",
        "Cloud Storage buckets. "
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "BUCKET_NAME = \"[your-bucket-name]\" #@param {type:\"string\"}"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "if BUCKET_NAME == \"\" or BUCKET_NAME is None or BUCKET_NAME == \"[your-bucket-name]\":\n",
        "    BUCKET_NAME = PROJECT_ID + \"_xai_metadata_\" + TIMESTAMP"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "**Only if your bucket doesn't already exist**: Run the following cell to create your Cloud Storage bucket."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "! gsutil mb -l $REGION gs://$BUCKET_NAME"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "### Import libraries\n",
        "\n",
        "Import the libraries for this tutorial."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "import os\n",
        "import json\n",
        "import random\n",
        "import numpy as np\n",
        "import PIL\n",
        "import tensorflow as tf\n",
        "import googleapiclient\n",
        "\n",
        "from matplotlib import pyplot as plt\n",
        "from base64 import b64encode\n",
        "\n",
        "# should be >= 2.1\n",
        "print(\"Tensorflow version \" + tf.__version__)\n",
        "if tf.__version__ < \"2.1\":\n",
        "    raise Exception(\"TF 2.1 or greater is required\")\n",
        "\n",
        "AUTO = tf.data.experimental.AUTOTUNE\n",
        "print(\"AUTO\", AUTO)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Import the Explanability SDK."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "import explainable_ai_sdk as xai"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Download and preprocess the data\n",
        "\n",
        "This section shows how to download the flower images, use the `tf.data` API to create a data input pipeline, and split the data into training and validation sets."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "# GCS location of TFRecords for flowers dataset\n",
        "GCS_PATTERN = 'gs://flowers-public/tfrecords-jpeg-192x192-2/*.tfrec'\n",
        "\n",
        "# Input size (Height, Width) for the Model\n",
        "IMAGE_SIZE = [192, 192]\n",
        "# Class labels for prediction\n",
        "CLASSES = ['daisy', 'dandelion', 'roses', 'sunflowers', 'tulips']  # do not change, maps to the labels in the data (folder names)\n",
        "\n",
        "# Batch size for training\n",
        "BATCH_SIZE = 32\n",
        "# Validation split for training\n",
        "VALIDATION_SPLIT = 0.2\n",
        "\n",
        "\n",
        "# Split data files between training and validation\n",
        "filenames = tf.io.gfile.glob(GCS_PATTERN)\n",
        "random.shuffle(filenames)\n",
        "split = int(len(filenames) * VALIDATION_SPLIT)\n",
        "training_filenames = filenames[split:]\n",
        "validation_filenames = filenames[:split]\n",
        "print(\"Pattern matches {} data files. Splitting dataset into {} training files and {} validation files\".format(len(filenames), len(training_filenames), len(validation_filenames)))\n",
        "validation_steps = int(3670 // len(filenames) * len(validation_filenames)) // BATCH_SIZE\n",
        "steps_per_epoch = int(3670 // len(filenames) * len(training_filenames)) // BATCH_SIZE\n",
        "print(\"With a batch size of {}, there will be {} batches per training epoch and {} batch(es) per validation run.\".format(BATCH_SIZE, steps_per_epoch, validation_steps))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "The following cell contains some image visualization utility functions. This code isn't essential to training or deploying the model. \n",
        "\n",
        "If you're running this from Colab the cell is hidden. You can look at the code by right clicking on the cell --> \"Form\" --> \"Show form\" if you'd like to see it."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "# @title display utilities [RUN ME]\n",
        "\n",
        "def dataset_to_numpy_util(dataset, N):\n",
        "    dataset = dataset.batch(N)\n",
        "\n",
        "    if tf.executing_eagerly():\n",
        "        #  In eager mode, iterate in the Dataset directly.\n",
        "        for images, labels in dataset:\n",
        "            numpy_images = images.numpy()\n",
        "            numpy_labels = labels.numpy()\n",
        "            break\n",
        "    else:\n",
        "        # In non-eager mode, must get the TF note that\n",
        "        # yields the nextitem and run it in a tf.Session.\n",
        "        get_next_item = dataset.make_one_shot_iterator().get_next()\n",
        "        with tf.Session() as ses:\n",
        "            numpy_images, numpy_labels = ses.run(get_next_item)\n",
        "\n",
        "    return numpy_images, numpy_labels\n",
        "\n",
        "\n",
        "def title_from_label_and_target(label, correct_label):\n",
        "    label = np.argmax(label, axis=-1)  # one-hot to class number\n",
        "    correct_label = np.argmax(correct_label, axis=-1)  # one-hot to class number\n",
        "    correct = (label == correct_label)\n",
        "    return \"{} [{}{}{}]\".format(CLASSES[label], str(correct), ', shoud be ' if not correct else '',\n",
        "                                CLASSES[correct_label] if not correct else ''), correct\n",
        "\n",
        "\n",
        "def display_one_flower(image, title, subplot, red=False):\n",
        "    plt.subplot(subplot)\n",
        "    plt.axis('off')\n",
        "    plt.imshow(image)\n",
        "    plt.title(title, fontsize=16, color='red' if red else 'black')\n",
        "    return subplot + 1\n",
        "\n",
        "\n",
        "def display_9_images_from_dataset(dataset):\n",
        "    subplot = 331\n",
        "    plt.figure(figsize=(13, 13))\n",
        "    images, labels = dataset_to_numpy_util(dataset, 9)\n",
        "    for i, image in enumerate(images):\n",
        "        title = CLASSES[np.argmax(labels[i], axis=-1)]\n",
        "        subplot = display_one_flower(image, title, subplot)\n",
        "        if i >= 8:\n",
        "            break\n",
        "\n",
        "    plt.tight_layout()\n",
        "    plt.subplots_adjust(wspace=0.1, hspace=0.1)\n",
        "    plt.show()\n",
        "\n",
        "\n",
        "def display_9_images_with_predictions(images, predictions, labels):\n",
        "    subplot = 331\n",
        "    plt.figure(figsize=(13, 13))\n",
        "    for i, image in enumerate(images):\n",
        "        title, correct = title_from_label_and_target(predictions[i], labels[i])\n",
        "        subplot = display_one_flower(image, title, subplot, not correct)\n",
        "        if i >= 8:\n",
        "            break\n",
        "\n",
        "    plt.tight_layout()\n",
        "    plt.subplots_adjust(wspace=0.1, hspace=0.1)\n",
        "    plt.show()\n",
        "\n",
        "\n",
        "def display_training_curves(training, validation, title, subplot):\n",
        "    if subplot % 10 == 1:  # set up the subplots on the first call\n",
        "        plt.subplots(figsize=(10, 10), facecolor='#F0F0F0')\n",
        "        plt.tight_layout()\n",
        "    ax = plt.subplot(subplot)\n",
        "    ax.set_facecolor('#F8F8F8')\n",
        "    ax.plot(training)\n",
        "    ax.plot(validation)\n",
        "    ax.set_title('model ' + title)\n",
        "    ax.set_ylabel(title)\n",
        "    ax.set_xlabel('epoch')\n",
        "    ax.legend(['train', 'valid.'])"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "### Read images and labels from TFRecords\n",
        "\n",
        "In this dataset each image is stored as a`tf.Example` in a `TFRecords` file, which is a serialized format for storing and retrieving training data from a disk storage. To learn more about `tf.Example` and `TFRecords`, you can view the Tensorflow documentation [here](https://www.tensorflow.org/tutorials/load_data/tfrecord).\n",
        "\n",
        "For our purpose, each image and corresponding label is encoded as a `tf.Example` as:\n",
        "\n",
        "- image: JPG compressed image bytes\n",
        "- class: the label (class) as a scalar (integer) value\n",
        "- one_hot_class: the label (class) as a one-hot encoded vector\n",
        "\n",
        "Our helper function `read_tfrecord` takes an individual `tf.Example`, deserializes using the specification defined in the variable `features`, extracts the image bytes and corresponding label, and JPG decodes (uncompresses) the image bytes.\n",
        "\n",
        "Our helper function `load_datasets` creates a `tf.data.Dataset` from the corresponding `TFRecords` specified by the parameter `filenames`."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "def read_tfrecord(example):\n",
        "    features = {\n",
        "        \"image\": tf.io.FixedLenFeature([], tf.string),  # tf.string means bytestring\n",
        "        \"class\": tf.io.FixedLenFeature([], tf.int64),  # shape [] means scalar\n",
        "        \"one_hot_class\": tf.io.VarLenFeature(tf.float32),\n",
        "    }\n",
        "    example = tf.io.parse_single_example(example, features)\n",
        "    image = tf.image.decode_jpeg(example['image'], channels=3)\n",
        "    image = tf.cast(image, tf.float32) / 255.0  # convert image to floats in [0, 1] range\n",
        "    image = tf.reshape(image, [*IMAGE_SIZE, 3])  # explicit size will be needed for TPU\n",
        "    one_hot_class = tf.sparse.to_dense(example['one_hot_class'])\n",
        "    one_hot_class = tf.reshape(one_hot_class, [5])\n",
        "    return image, one_hot_class\n",
        "\n",
        "\n",
        "def load_dataset(filenames):\n",
        "    # Read data from TFRecords\n",
        "\n",
        "    dataset = tf.data.Dataset.from_tensor_slices(filenames)\n",
        "    dataset = dataset.interleave(tf.data.TFRecordDataset, cycle_length=16, num_parallel_calls=AUTO)  # faster\n",
        "    dataset = dataset.map(read_tfrecord, num_parallel_calls=AUTO)\n",
        "    return dataset"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Use the visualization utility function provided earlier to preview flower images with their labels."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "display_9_images_from_dataset(load_dataset(training_filenames))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "### Create training and validation datasets\n",
        "\n",
        "Next, we will use our helper functions from the previous code cell to load the training and validation data into tf.data.Datasets, and set dataset properties for:\n",
        "\n",
        "- cache(): Reads the dataset once from disk, which is then held in memory during training.\n",
        "- repeat(): Enables multiple passes over the dataset for training more than a single epoch.\n",
        "- shuffle(): The number of examples to random shuffle at a time.\n",
        "- batch(): The batch size for training.\n",
        "- prefetch(): With AUTO setting, the dataset iterator will auto determine when/how many examples to prefetch in parallel while feeding examples for training.\n",
        "\n",
        "Finally, we will prefetch a batch of examples from the validation dataset to use later for prediction and explainability --which we store as `some_flowers` and `some_labels`."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "def get_batched_dataset(filenames):\n",
        "    dataset = load_dataset(filenames)\n",
        "    dataset = dataset.cache()  # This dataset fits in RAM\n",
        "    dataset = dataset.repeat()\n",
        "    dataset = dataset.shuffle(2048)\n",
        "    dataset = dataset.batch(BATCH_SIZE)\n",
        "    dataset = dataset.prefetch(AUTO)  # prefetch next batch while training (autotune prefetch buffer size)\n",
        "    # For proper ordering of map/batch/repeat/prefetch, see Dataset performance guide: https://www.tensorflow.org/guide/performance/datasets\n",
        "    return dataset\n",
        "\n",
        "\n",
        "def get_training_dataset():\n",
        "    return get_batched_dataset(training_filenames)\n",
        "\n",
        "\n",
        "def get_validation_dataset():\n",
        "    return get_batched_dataset(validation_filenames)\n",
        "\n",
        "\n",
        "some_flowers, some_labels = dataset_to_numpy_util(load_dataset(validation_filenames), BATCH_SIZE)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Build, train, and evaluate the model\n",
        "\n",
        "This section shows how to build, train, evaluate, and get local predictions from a model by using the TF.Keras [Sequential API](https://www.tensorflow.org/guide/keras/sequential_model).\n",
        "\n",
        "Note that when we loaded and deserialized the `tf.Example` records, we extracted the one-hot-encoded version of the label. So when we `compile()` the model, we use `categorical_crossentropy` for the loss function."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "from tensorflow.keras import Sequential\n",
        "from tensorflow.keras.layers import Conv2D, MaxPooling2D, Dense, GlobalAveragePooling2D, BatchNormalization\n",
        "from tensorflow.keras.optimizers import Adam\n",
        "\n",
        "model = Sequential([\n",
        "    # Stem\n",
        "    Conv2D(kernel_size=3, filters=16, padding='same', activation='relu', input_shape=[*IMAGE_SIZE, 3]),\n",
        "    BatchNormalization(),\n",
        "    Conv2D(kernel_size=3, filters=32, padding='same', activation='relu'),\n",
        "    BatchNormalization(),\n",
        "    MaxPooling2D(pool_size=2),\n",
        "\n",
        "    # Conv Group\n",
        "    Conv2D(kernel_size=3, filters=64, padding='same', activation='relu'),\n",
        "    BatchNormalization(),\n",
        "    MaxPooling2D(pool_size=2),\n",
        "    Conv2D(kernel_size=3, filters=96, padding='same', activation='relu'),\n",
        "    BatchNormalization(),\n",
        "    MaxPooling2D(pool_size=2),\n",
        "\n",
        "    # Conv Group\n",
        "    Conv2D(kernel_size=3, filters=128, padding='same', activation='relu'),\n",
        "    BatchNormalization(),\n",
        "    MaxPooling2D(pool_size=2),\n",
        "    Conv2D(kernel_size=3, filters=128, padding='same', activation='relu'),\n",
        "    BatchNormalization(),\n",
        "\n",
        "    # 1x1 Reduction\n",
        "    Conv2D(kernel_size=1, filters=32, padding='same', activation='relu'),\n",
        "    BatchNormalization(),\n",
        "\n",
        "    # Classifier\n",
        "    GlobalAveragePooling2D(),\n",
        "    Dense(5, activation='softmax')\n",
        "])\n",
        "\n",
        "model.compile(\n",
        "    optimizer=Adam(lr=0.005, decay=0.98),\n",
        "    loss='categorical_crossentropy',\n",
        "    metrics=['accuracy'])\n",
        "model.summary()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "### Train the model\n",
        "\n",
        "Train this on a GPU if you have access (in Colab, from the menu select **Runtime** --> **Change runtime type**). On a CPU, it'll take ~30 minutes to run training. On a GPU, it takes ~5 minutes.\n",
        "\n",
        "After the model is trained, we will print out the validation accuracy per epoch from the model's training `history`."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "EPOCHS = 20  # Train for 60 epochs for higher accuracy, 20 should get you ~75%\n",
        "\n",
        "history = model.fit(get_training_dataset(), steps_per_epoch=steps_per_epoch, epochs=EPOCHS,\n",
        "                    validation_data=get_validation_dataset(), validation_steps=validation_steps)\n",
        "\n",
        "print(history.history['val_accuracy'])"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "### Visualize local predictions\n",
        "\n",
        "Get predictions on your local model and visualize the images with their predicted labels, using the visualization utility function provided earlier."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "# Randomize the input so that you can execute multiple times to change results\n",
        "permutation = np.random.permutation(BATCH_SIZE)\n",
        "some_flowers, some_labels = (some_flowers[permutation], some_labels[permutation])\n",
        "\n",
        "# Get predictions from our sample batch\n",
        "predictions = model.predict(some_flowers, batch_size=16)\n",
        "print('Predictions', np.array(CLASSES)[np.argmax(predictions, axis=-1)].tolist(), '\\n')\n",
        "\n",
        "# Get evaluations (true labels) from our sample batch\n",
        "evaluations = model.evaluate(some_flowers, some_labels, batch_size=16)\n",
        "print('Evaluation [val_loss, val_acc]', evaluations)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "display_9_images_with_predictions(some_flowers, predictions, some_labels)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Export the model as a TF 2.x SavedModel\n",
        "\n",
        "When using TensorFlow 2.x, you export the model as a `SavedModel` and load it into Cloud Storage. During export, you need to define a serving function to convert data to the format your model expects. If you send encoded data to AI Platform, your serving function ensures that the data is decoded on the model server before it is passed as input to your model.\n",
        "\n",
        "### Serving function for image data\n",
        "\n",
        "Sending base 64 encoded image data to AI Platform is more space efficient. Since this deployed model expects input data as raw bytes, you need to ensure that the b64 encoded data gets converted back to raw bytes before it is passed as input to the deployed model.\n",
        "\n",
        "To resolve this, define a serving function (`serving_fn`) and attach it to the model as a preprocessing step. Add a `@tf.function` decorator so the serving function is part of the model's graph (instead of upstream on a CPU).\n",
        "\n",
        "When you send a prediction or explanation request, the request goes to the serving function (`serving_fn`), which preprocesses the b64 bytes into raw numpy bytes (`preprocess_fn`). At this point, the data can be passed to the model (`m_call`)."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "export_path = 'gs://' + BUCKET_NAME + '/explanations/mymodel'\n",
        "\n",
        "\n",
        "def _preprocess(bytes_input):\n",
        "    decoded = tf.io.decode_jpeg(bytes_input, channels=3)\n",
        "    decoded = tf.image.convert_image_dtype(decoded, tf.float32)\n",
        "    resized = tf.image.resize(decoded, size=(192, 192))\n",
        "    return resized\n",
        "\n",
        "\n",
        "@tf.function(input_signature=[tf.TensorSpec([None], tf.string)])\n",
        "def preprocess_fn(bytes_inputs):\n",
        "    with tf.device(\"cpu:0\"):\n",
        "        decoded_images = tf.map_fn(_preprocess, bytes_inputs, dtype=tf.float32)\n",
        "    return {\"numpy_inputs\": decoded_images}  # User needs to make sure the key matches model's input\n",
        "\n",
        "\n",
        "m_call = tf.function(model.call).get_concrete_function([tf.TensorSpec(shape=[None, 192, 192, 3], dtype=tf.float32, name=\"numpy_inputs\")])\n",
        "\n",
        "\n",
        "@tf.function(input_signature=[tf.TensorSpec([None], tf.string)])\n",
        "def serving_fn(bytes_inputs):\n",
        "    images = preprocess_fn(bytes_inputs)\n",
        "    prob = m_call(**images)\n",
        "    return prob\n",
        "\n",
        "\n",
        "tf.saved_model.save(model, export_path, signatures={\n",
        "    'serving_default': serving_fn,\n",
        "    'xai_preprocess': preprocess_fn,  # Required for XAI\n",
        "    'xai_model': m_call  # Required for XAI\n",
        "})"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Get signatures\n",
        "\n",
        "You can get the signatures of your model's input and output layers by reloading the model into memory, and querying it for the signatures corresponding to each layer.\n",
        "\n",
        "When you subsequently generate metadata for explanations, you will need the signature of the input layer to override the metadata builder defaulting to using the serving function's input signature."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "loaded = tf.saved_model.load(export_path)\n",
        "\n",
        "input_name = list(loaded.signatures['xai_model'].structured_input_signature[1].keys())[0]\n",
        "print(input_name)\n",
        "\n",
        "output_name = list(loaded.signatures['xai_model'].structured_outputs.keys())[0]\n",
        "print(output_name)\n",
        "\n",
        "preprocess_name = list(loaded.signatures['xai_preprocess'].structured_input_signature[1].keys())[0]\n",
        "print(preprocess_name)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Generate explanation metadata\n",
        "\n",
        "In order to deploy this model to AI Platform Explanations, we need to create an explanation_metadata.json file with information about our model inputs, outputs, and baseline. We can automatically generate and export this file by using Metadata builder from the AI Platform Explainability SDK.\n",
        "\n",
        "### Input baseline\n",
        "\n",
        "For image models, using [0,1] as your input baseline represents black and white images. In this case we're using np.random to generate the baseline because our training images contain a lot of black and white (i.e. daisy petals).\n",
        "\n",
        "Let's start by first creating the input baseline."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "random_baseline = np.random.rand(192, 192, 3)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "### Metadata builder\n",
        "\n",
        "Next, you will instantiate an instance of SavedModelMetadatabuilder, with the parameters:\n",
        "\n",
        "- model_path: path to the SavedModel format stored model. In this case, the path is the model stored in your Cloud Storage bucket, specified by `export_path`.\n",
        "- signature_name: overrides looking at the default signature to find the input and output layers for explanations."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "from explainable_ai_sdk.metadata.tf.v2 import SavedModelMetadataBuilder\n",
        "# By default models inputs and outputs are added to metadata.\n",
        "builder = xai.metadata.tf.v2.SavedModelMetadataBuilder(export_path, signature_name='xai_model')"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Next, we set some parameters for generating the metadata with the method `set_image_metadata`, with the parameters:\n",
        "\n",
        "- input_name: the signature of the input layer to override using the serving function.\n",
        "- input_baseline: the input baseline.\n",
        "- visualization: how the explanation data should be visualized."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "# We add actual float image to explain as b64 string can't be explained.\n",
        "builder.set_image_metadata(input_name, input_baselines=[random_baseline.tolist()], visualization={'type': 'Pixels', 'overlay_type': 'grayscale'})"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "### Upload the metadata\n",
        "\n",
        "The last step in generating the metadata for explanations, is to upload the generated metadata json to where the stored model using the method `save_model_with_metadata`."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "# Save the model with metadata to our GCS export path.\n",
        "export_path = builder.save_model_with_metadata(export_path)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Local experience - predictions and explanations\n",
        "\n",
        "Now that you have a AI Platform model with the corresponding predictions, we can export the model to a local instance of an AI Plarform notebook and make local predictions with the corresponding explanations.\n",
        "\n",
        "The first step is to export the model locally. We will use the AI Platform Explanability SDK method `load_model_from_local_path()`, with the parameters:\n",
        "\n",
        "- `model_path`: The GCS location of the AI Platform model.\n",
        "- `config`: The type of XAI explanations to instrument the local model for."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "lm = xai.load_model_from_local_path(export_path, xai.XraiConfig())"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "### Get and prepare test images\n",
        "\n",
        "To prepare the test images:\n",
        "\n",
        "- Download a small sample of images from the flowers dataset -- just enough for a multiple instance prediction (3 in total).\n",
        "- Resize the images to match the input shape (192, 192) of the model.\n",
        "- Save the resized images back to your bucket."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "# Download test flowers from public bucket\n",
        "! mkdir flowers\n",
        "! gsutil -m cp gs://flowers_model/test_flowers/* ./flowers\n",
        "    \n",
        "# Resize the images to what your model is expecting (192,192)\n",
        "test_filenames = []\n",
        "\n",
        "for i in os.listdir('flowers'):\n",
        "    img_path = 'flowers/' + i\n",
        "    with PIL.Image.open(img_path) as ex_img:\n",
        "        resize_img = ex_img.resize([192, 192])\n",
        "        resize_img.save(img_path)\n",
        "        test_filenames.append(img_path)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "### Format your explanation request\n",
        "\n",
        "Prepare a batch of instances in JSONL format."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "# Prepare your prediction JSON to send to your Cloud model\n",
        "instances = []\n",
        "\n",
        "for image_path in test_filenames:\n",
        "    img_bytes = tf.io.read_file(image_path)\n",
        "    b64str = b64encode(img_bytes.numpy()).decode('utf-8')\n",
        "    instances.append({preprocess_name: {'b64': b64str}})"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "### Make prediction with explanation\n",
        "\n",
        "Now that you have a local version of your model on your AI Platform notebook, we can make a prediction from the three test images and get corresponding explanations using the `explain()` method for your model."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "explanations = lm.explain(instances)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "### Visualize the predictions and explanations\n",
        "\n",
        "Let's now visualize the explanations for the three test images from the response object to the `explain()` call. This response object will return a list, where each element is the explanation data for the corresponding test image. You will use the method `visualize_attributions()` to visualize the explanation."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "explanations[0].visualize_attributions()\n",
        "explanations[1].visualize_attributions()\n",
        "explanations[2].visualize_attributions()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Deploy model to AI Explanations\n",
        "\n",
        "This section shows how to use `gcloud` to deploy the model to AI Explanations, using two different explanation methods for image models. After the model is deployed, we will export a local copy of the deployed model to your AI Platform notebook instance and repeat making a prediction and corresponding explanations.\n",
        "\n",
        "### Create the model"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "import datetime\n",
        "MODEL = 'flowers' + TIMESTAMP\n",
        "print(MODEL)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "# Create the model if it doesn't exist yet (you only need to run this once)\n",
        "! gcloud ai-platform models create $MODEL --enable-logging --regions=us-central1"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "### Create explainable model versions\n",
        "\n",
        "For image models, we offer two choices for explanation methods: \n",
        "* Integrated Gradients (IG)\n",
        "* XRAI \n",
        "\n",
        "You can find more info on each method in the [documentation](TODO). In this tutorial, you will deploy for the XRAI explanation method.\n",
        "\n",
        "Creating the version will take ~5-10 minutes. Note that your first deploy may take longer.\n",
        "\n",
        "#### Deploy an XRAI model"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "XRAI_VERSION='xrai'\n",
        "\n",
        "# Create the XRAI version with gcloud\n",
        "! gcloud beta ai-platform versions create $XRAI_VERSION \\\n",
        "--model $MODEL \\\n",
        "--origin $export_path \\\n",
        "--runtime-version 2.2 \\\n",
        "--framework TENSORFLOW \\\n",
        "--python-version 3.7 \\\n",
        "--machine-type n1-standard-4 \\\n",
        "--explanation-method xrai \\\n",
        "--num-integral-steps 25"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "# Make sure the XRAI model deployed correctly. State should be `READY` in the following log\n",
        "! gcloud ai-platform versions describe $XRAI_VERSION --model $MODEL"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "### Export deployed model with explanability to your local notebook instance\n",
        "\n",
        "Finally, you will repeat the same explanability task, but this time the exported model will be the deployed model (vs. previously was the model resource). You export the deployed model using the `load_model_from_ai_platform()` method of the AI Explainability SDK."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "lm = xai.load_model_from_ai_platform(PROJECT_ID, MODEL, XRAI_VERSION)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "### Make prediction with explanation\n",
        "\n",
        "Now that you have a local version of your deployed model on your AI Platform notebook, we can repeat making a prediction from the three test images and get corresponding explanations using the `explain()` method for your model."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "explanations = lm.explain(instances)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "### Visualize the predictions and explanations\n",
        "\n",
        "Let's now visualize the explanations for the three test images from the response object to the `explain()` call. This response object will return a list, where each element is the explanation data for the corresponding test image. You will use the method `visualize_attributions()` to visualize the explanation."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "explanations[0].visualize_attributions()\n",
        "explanations[1].visualize_attributions()\n",
        "explanations[2].visualize_attributions()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Cleaning up\n",
        "\n",
        "To clean up all GCP resources used in this project, you can [delete the GCP\n",
        "project](https://cloud.google.com/resource-manager/docs/creating-managing-projects#shutting_down_projects) you used for the tutorial.\n",
        "\n",
        "Alternatively, you can clean up individual resources by running the following\n",
        "commands:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "# Delete model version resource\n",
        "! gcloud ai-platform versions delete $IG_VERSION --quiet --model $MODEL\n",
        "! gcloud ai-platform versions delete $XRAI_VERSION --quiet --model $MODEL\n",
        "\n",
        "# Delete model resource\n",
        "! gcloud ai-platform models delete $MODEL --quiet\n",
        "\n",
        "# Delete Cloud Storage objects that were created\n",
        "! gsutil -m rm -r gs://$BUCKET_NAME"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "If your Cloud Storage bucket doesn't contain any other objects and you would like to delete it, run `gsutil rm -r gs://$BUCKET_NAME`."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## What's next?\n",
        "\n",
        "To learn more about AI Explanations, check out the resources here.\n",
        "\n",
        "* [AI Explanations documentation](https://cloud.google.com/ai-platform/prediction/docs/ai-explanations/overview)\n",
        "* [AI Explanations whitepaper](https://storage.googleapis.com/cloud-ai-whitepapers/AI%20Explainability%20Whitepaper.pdf)\n",
        "* [Integrated gradients paper](https://arxiv.org/abs/1703.01365)\n",
        "* [XRAI paper](https://arxiv.org/abs/1906.02825)"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "name": "ai-explanations-local-experience.ipynb",
      "provenance": [],
      "toc_visible": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}