notebooks/community/sdk/SDK_Explainable_AI_Custom_Tabular.ipynb

{
  "cells": [
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Snm88Z0sROQB"
      },
      "outputs": [],
      "source": [
        "# Copyright 2021 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": {
        "id": "ThTHoy7DRSye"
      },
      "source": [
        "# Feedback or issues?\n",
        "\n",
        "For any feedback or questions, please open an [issue](https://github.com/googleapis/python-aiplatform/issues)."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "acg-u47CJ7-T"
      },
      "source": [
        "# Explainable AI via MB SDK on Custom Tabular model\n",
        "\n",
        "To use this Jupyter notebook, copy the notebook to a Google Cloud Notebooks instance with Tensorflow installed and open it. You can run each step, or cell, and see its results. To run a cell, use Shift+Enter. Jupyter automatically displays the return value of the last line in each cell. For more information about running notebooks in Google Cloud Notebook, see the [Google Cloud Notebook guide](https://cloud.google.com/vertex-ai/docs/general/notebooks).\n",
        "\n",
        "\n",
        "This notebook demonstrate how to create an Custom Tabular Model and how to serve the model for online prediction with Explainability.\n",
        "\n",
        "Note: you may incur charges for training, prediction, storage or usage of other GCP products in connection with testing this SDK."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "-VL5avdUJ7-V"
      },
      "source": [
        "### Install Vertex SDK for Python\n",
        "\n",
        "\n",
        "After the SDK installation the kernel will be automatically restarted."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "5-6l0OdRJ7-V"
      },
      "outputs": [],
      "source": [
        "%%capture\n",
        "!pip3 uninstall -y google-cloud-aiplatform\n",
        "!pip3 install -y google-cloud-aiplatform tabulate\n",
        "import IPython\n",
        "\n",
        "app = IPython.Application.instance()\n",
        "app.kernel.do_shutdown(True)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "KjLd9x994wm3"
      },
      "outputs": [],
      "source": [
        "import sys\n",
        "\n",
        "if \"google.colab\" in sys.modules:\n",
        "    from google.colab import auth\n",
        "\n",
        "    auth.authenticate_user()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "amZxHeNzRhgN"
      },
      "source": [
        "### Enter Your Project and GCS Bucket\n",
        "\n",
        "Enter your Project Id in the cell below. Then run the cell to make sure the Cloud SDK uses the right project for all the commands in this notebook."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "GtIjHUhlJ7-W"
      },
      "outputs": [],
      "source": [
        "MY_PROJECT = \"YOUR PROJECT ID\"\n",
        "MY_STAGING_BUCKET = \"gs://YOUR BUCKET\"  # bucket should be in same region as ucaip"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "itgJdQPGJ7-W"
      },
      "source": [
        "## Set up SDK workspace"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "o_wnT10RJ7-W"
      },
      "outputs": [],
      "source": [
        "import uuid\n",
        "\n",
        "import tensorflow as tf\n",
        "from google.cloud import aiplatform\n",
        "from tabulate import tabulate"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "O8XJZB3gR8eL"
      },
      "source": [
        "## Initialize Vertex SDK for Python\n",
        "\n",
        "Initialize the *client* for Vertex AI."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Wrlk2B2nJ7-X"
      },
      "outputs": [],
      "source": [
        "aiplatform.init(project=MY_PROJECT, staging_bucket=MY_STAGING_BUCKET)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "rQMGw0XlPh-o"
      },
      "source": [
        "## Create Training Script that saves Explainable model"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "jNBmoLbhJ7-X"
      },
      "outputs": [],
      "source": [
        "%%writefile training_script.py\n",
        "# Single, Mirror and Multi-Machine Distributed Training for CIFAR-10\n",
        "\n",
        "from explainable_ai_sdk.metadata.tf.v2 import SavedModelMetadataBuilder\n",
        "\n",
        "from tensorflow.python.client import device_lib\n",
        "import tensorflow_datasets as tfds\n",
        "import tensorflow as tf\n",
        "\n",
        "import numpy as np\n",
        "import tempfile\n",
        "import argparse\n",
        "import sys\n",
        "import os\n",
        "\n",
        "tfds.disable_progress_bar()\n",
        "\n",
        "parser = argparse.ArgumentParser()\n",
        "parser.add_argument('--model-dir', dest='model_dir',\n",
        "                    default=os.getenv('AIP_MODEL_DIR'), type=str, help='Model dir.')\n",
        "parser.add_argument('--lr', dest='lr',\n",
        "                    default=0.001, type=float,\n",
        "                    help='Learning rate.')\n",
        "parser.add_argument('--epochs', dest='epochs',\n",
        "                    default=20, type=int,\n",
        "                    help='Number of epochs.')\n",
        "parser.add_argument('--steps', dest='steps',\n",
        "                    default=100, type=int,\n",
        "                    help='Number of steps per epoch.')\n",
        "parser.add_argument('--distribute', dest='distribute', type=str, default='single',\n",
        "                    help='distributed training strategy')\n",
        "args = parser.parse_args()\n",
        "\n",
        "print('Python Version = {}'.format(sys.version))\n",
        "print('TensorFlow Version = {}'.format(tf.__version__))\n",
        "print('TF_CONFIG = {}'.format(os.environ.get('TF_CONFIG', 'Not found')))\n",
        "\n",
        "# Single Machine, single compute device\n",
        "if args.distribute == 'single':\n",
        "    if tf.test.is_gpu_available():\n",
        "        strategy = tf.distribute.OneDeviceStrategy(device=\"/gpu:0\")\n",
        "    else:\n",
        "        strategy = tf.distribute.OneDeviceStrategy(device=\"/cpu:0\")\n",
        "# Single Machine, multiple compute device\n",
        "elif args.distribute == 'mirror':\n",
        "    strategy = tf.distribute.MirroredStrategy()\n",
        "# Multiple Machine, multiple compute device\n",
        "elif args.distribute == 'multi':\n",
        "    strategy = tf.distribute.experimental.MultiWorkerMirroredStrategy()\n",
        "\n",
        "# Multi-worker configuration\n",
        "print('num_replicas_in_sync = {}'.format(strategy.num_replicas_in_sync))\n",
        "\n",
        "def make_dataset():\n",
        "  # Scaling Boston Housing data features\n",
        "  def scale(feature):\n",
        "    max = np.max(feature)\n",
        "    feature = (feature / max).astype(np.float)\n",
        "    return feature\n",
        "\n",
        "  (x_train, y_train), (x_test, y_test) = tf.keras.datasets.boston_housing.load_data(\n",
        "    path=\"boston_housing.npz\", test_split=0.2, seed=113\n",
        "  )\n",
        "  for _ in range(13):\n",
        "    x_train[_] = scale(x_train[_])\n",
        "    x_test[_] = scale(x_test[_])\n",
        "  return (x_train, y_train), (x_test, y_test)\n",
        "\n",
        "# Build the Keras model\n",
        "def build_and_compile_dnn_model():\n",
        "  model = tf.keras.Sequential([\n",
        "      tf.keras.layers.Dense(128, activation='relu', input_shape=(13,)),\n",
        "      tf.keras.layers.Dense(128, activation='relu'),\n",
        "      tf.keras.layers.Dense(1, activation='linear')\n",
        "  ])\n",
        "  model.compile(\n",
        "      loss='mse',\n",
        "      optimizer=tf.keras.optimizers.RMSprop(learning_rate=args.lr))\n",
        "  return model\n",
        "\n",
        "# Train the model\n",
        "NUM_WORKERS = strategy.num_replicas_in_sync\n",
        "# Here the batch size scales up by number of workers since\n",
        "# `tf.data.Dataset.batch` expects the global batch size.\n",
        "BATCH_SIZE = 16\n",
        "GLOBAL_BATCH_SIZE = BATCH_SIZE * NUM_WORKERS\n",
        "\n",
        "with strategy.scope():\n",
        "  # Creation of dataset, and model building/compiling need to be within\n",
        "  # `strategy.scope()`.\n",
        "  model = build_and_compile_dnn_model()\n",
        "\n",
        "# Train the model\n",
        "(x_train, y_train), (x_test, y_test) = make_dataset()\n",
        "model.fit(x_train, y_train, epochs=args.epochs, batch_size=GLOBAL_BATCH_SIZE)\n",
        "\n",
        "tmpdir = tempfile.mkdtemp()\n",
        "\n",
        "model.save(tmpdir)\n",
        "\n",
        "# Save TF Model with Explainable metadata to GCS\n",
        "builder = SavedModelMetadataBuilder(tmpdir)\n",
        "builder.save_model_with_metadata(args.model_dir)\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Yp2clkOJSDhR"
      },
      "source": [
        "## Launch a Training Job and Create a Model on Vertex AI"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "FTtte5rGSGhc"
      },
      "source": [
        "### Config a Training Job"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "OwmqEIKYJ7-Y"
      },
      "outputs": [],
      "source": [
        "job = aiplatform.CustomTrainingJob(\n",
        "    display_name=f\"temp-mbsdk-explainable-ai-custom-tabular-nb-{uuid.uuid4()}\",\n",
        "    script_path=\"training_script.py\",\n",
        "    container_uri=\"gcr.io/cloud-aiplatform/training/tf-gpu.2-1:latest\",\n",
        "    requirements=[\n",
        "        \"tensorflow_datasets\",\n",
        "        \"explainable-ai-sdk\",\n",
        "    ],\n",
        "    model_serving_container_image_uri=\"gcr.io/cloud-aiplatform/prediction/tf2-gpu.2-1:latest\",\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "FdqpORPhSLbx"
      },
      "source": [
        "### Run the Training Job"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "oVTORjQpJ7-Y"
      },
      "outputs": [],
      "source": [
        "model = job.run(\n",
        "    model_display_name=\"temp-boston-housing-mbsdk-explainable-tabular-model\",\n",
        "    replica_count=1,\n",
        "    machine_type=\"n1-standard-4\",\n",
        "    accelerator_type=\"NVIDIA_TESLA_K80\",\n",
        "    accelerator_count=1,\n",
        "    args=[\"--epochs=50\", \"--distribute=single\"],\n",
        ")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "uMnxLeMrJ7-Y"
      },
      "outputs": [],
      "source": [
        "# Get info about the Custom Job\n",
        "print(\n",
        "    f\"Display Name:\\t{job.display_name}\\n\"\n",
        "    f\"Resource Name:\\t{job.resource_name}\\n\"\n",
        "    f\"Current State:\\t{job.state.name}\\n\"\n",
        ")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "2Udmt7tpJ7-Y"
      },
      "outputs": [],
      "source": [
        "# Get path to saved model in GCS\n",
        "output_dir = model._gca_resource.artifact_uri"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "vKlHRCOPJ7-Z"
      },
      "source": [
        "## Build the Explanation Metadata and Parameters"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "6z9zrDmTJ7-Z"
      },
      "outputs": [],
      "source": [
        "loaded = tf.keras.models.load_model(output_dir)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "jbp3HRc7J7-Z"
      },
      "outputs": [],
      "source": [
        "serving_input = list(\n",
        "    loaded.signatures[\"serving_default\"].structured_input_signature[1].keys()\n",
        ")[0]\n",
        "serving_output = list(loaded.signatures[\"serving_default\"].structured_outputs.keys())[0]\n",
        "feature_names = [\n",
        "    \"crim\",\n",
        "    \"zn\",\n",
        "    \"indus\",\n",
        "    \"chas\",\n",
        "    \"nox\",\n",
        "    \"rm\",\n",
        "    \"age\",\n",
        "    \"dis\",\n",
        "    \"rad\",\n",
        "    \"tax\",\n",
        "    \"ptratio\",\n",
        "    \"b\",\n",
        "    \"lstat\",\n",
        "]"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "xpYFtjRbJ7-Z"
      },
      "outputs": [],
      "source": [
        "explain_params = aiplatform.explain.ExplanationParameters(\n",
        "    {\"sampled_shapley_attribution\": {\"path_count\": 10}}\n",
        ")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "AtLFH_FeJ7-a"
      },
      "outputs": [],
      "source": [
        "input_metadata = {\n",
        "    \"input_tensor_name\": serving_input,\n",
        "    \"encoding\": \"BAG_OF_FEATURES\",\n",
        "    \"modality\": \"numeric\",\n",
        "    \"index_feature_mapping\": feature_names,\n",
        "}\n",
        "output_metadata = {\"output_tensor_name\": serving_output}\n",
        "\n",
        "input_metadata = aiplatform.explain.ExplanationMetadata.InputMetadata(input_metadata)\n",
        "output_metadata = aiplatform.explain.ExplanationMetadata.OutputMetadata(output_metadata)\n",
        "\n",
        "explain_metadata = aiplatform.explain.ExplanationMetadata(\n",
        "    inputs={\"features\": input_metadata}, outputs={\"medv\": output_metadata}\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "744xmYcKJ7-a"
      },
      "source": [
        "## Deploy the model with model explanations enabled"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "IByGRREtJ7-a"
      },
      "outputs": [],
      "source": [
        "endpoint = model.deploy(\n",
        "    machine_type=\"n1-standard-4\",\n",
        "    accelerator_type=\"NVIDIA_TESLA_K80\",\n",
        "    accelerator_count=1,\n",
        "    explanation_metadata=explain_metadata,\n",
        "    explanation_parameters=explain_params,\n",
        ")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Gx0EJFHXxtqv"
      },
      "outputs": [],
      "source": [
        "print(f\"Endpoint resource name: {endpoint.resource_name}\")\n",
        "print(\n",
        "    f\"\\nTo use this endpoint in the future:\\nendpoint = aiplatform.Endpoint('{endpoint.resource_name}')\"\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "E3K2TSrPJ7-a"
      },
      "source": [
        "## Fetch test data to use"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "UYrdURLsJ7-b"
      },
      "outputs": [],
      "source": [
        "import numpy as np\n",
        "from tensorflow.keras.datasets import boston_housing\n",
        "\n",
        "(_, _), (x_test, y_test) = boston_housing.load_data(\n",
        "    path=\"boston_housing.npz\", test_split=0.2, seed=113\n",
        ")\n",
        "\n",
        "\n",
        "def scale(feature):\n",
        "    max = np.max(feature)\n",
        "    feature = (feature / max).astype(np.float32)\n",
        "    return feature\n",
        "\n",
        "\n",
        "for _ in range(13):\n",
        "    x_test[_] = scale(x_test[_])\n",
        "x_test = x_test.astype(np.float32)\n",
        "\n",
        "print(x_test.shape, x_test.dtype, y_test.shape)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "s_fNdLGPJ7-b"
      },
      "source": [
        "## Get predictions with explanations on our deployed tabular model"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "AhEjJQCsJ7-b"
      },
      "outputs": [],
      "source": [
        "response = endpoint.explain(\n",
        "    instances=[{\"dense_input\": s.tolist()} for s in [x_test[0]]]\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "8FBZDKVsJ7-b"
      },
      "source": [
        "## Check out feature attributions"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "urngajW6J7-b"
      },
      "outputs": [],
      "source": [
        "test_data = x_test[0]\n",
        "attributions = response.explanations[0].attributions[0].feature_attributions\n",
        "\n",
        "rows = []\n",
        "for i, val in enumerate(feature_names):\n",
        "    rows.append([val, test_data[i], attributions[val][0]])\n",
        "print(tabulate(rows, headers=[\"Feature name\", \"Feature value\", \"Attribution value\"]))"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "collapsed_sections": [],
      "name": "AI_Platform_(Unified)_SDK_Explainable_AI_Custom_Tabular.ipynb",
      "toc_visible": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}