notebooks/official/pipelines/custom_model_training_and_batch_prediction.ipynb

{
  "cells": [
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "b0b4f2bf"
      },
      "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": "c5c0d0bbf74d"
      },
      "source": [
        "# Vertex AI Pipelines: Custom training with pre-built Google Cloud Pipeline Components\n",
        "\n",
        "<table align=\"left\">\n",
        "\n",
        "  <td>\n",
        "    <a href=\"https://colab.research.google.com/github/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/official/pipelines/custom_model_training_and_batch_prediction.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/vertex-ai-samples/blob/main/notebooks/official/pipelines/custom_model_training_and_batch_prediction.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",
        "  <td>\n",
        "    <a href=\"https://console.cloud.google.com/vertex-ai/workbench/deploy-notebook?download_url=https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/main/notebooks/official/pipelines/custom_model_training_and_batch_prediction.ipynb\">\n",
        "      <img src=\"https://lh3.googleusercontent.com/UiNooY4LUgW_oTvpsNhPpQzsstV5W8F7rYgxgGBD85cWJoLmrOzhVs_ksK_vgx40SHs7jCqkTkCk=e14-rj-sc0xffffff-h130-w32\" alt=\"Vertex AI logo\">\n",
        "      Open in Vertex AI Workbench\n",
        "    </a>\n",
        "  </td>                                                                                               \n",
        "</table>\n",
        "<br/><br/><br/>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "d79790f56d30"
      },
      "source": [
        "## Overview\n",
        "\n",
        "\n",
        "This tutorial demonstrates how to use Vertex AI Pipelines with pre-built Google Cloud Pipeline Components for custom training.\n",
        "\n",
        "Learn more about [Vertex AI Pipelines](https://cloud.google.com/vertex-ai/docs/pipelines/introduction) and [Custom training components](https://cloud.google.com/vertex-ai/docs/training/create-training-pipeline)."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "b87eb4840013"
      },
      "source": [
        "### Objective\n",
        "\n",
        "In this tutorial, you learn to use Vertex AI Pipelines and Google Cloud Pipeline components to build a custom model.\n",
        "\n",
        "\n",
        "This tutorial uses the following Google Cloud ML services:\n",
        "\n",
        "- Vertex AI Pipelines\n",
        "- Google Cloud Pipeline components\n",
        "- Vertex AI Training\n",
        "- Vertex AI model resource\n",
        "- Vertex AI endpoint resource\n",
        "\n",
        "The steps performed include:\n",
        "\n",
        "- Create a KFP pipeline:\n",
        "    - Train a custom model.\n",
        "    - Upload the trained model as a model resource.\n",
        "    - Create an endpoint resource.\n",
        "    - Deploy the model resource to the endpoint resource.\n",
        "    - Make a batch prediction request.\n",
        "\n",
        "Learn more about [Google Cloud Pipeline components](https://cloud.google.com/vertex-ai/docs/pipelines/build-pipeline)."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "57139e75264f"
      },
      "source": [
        "### Dataset\n",
        "\n",
        "The dataset used for this tutorial is the [CIFAR10 dataset](https://www.tensorflow.org/datasets/catalog/cifar10) from [TensorFlow Datasets](https://www.tensorflow.org/datasets/catalog/overview). The version of the dataset you'll use is built into TensorFlow. The trained model predicts which type of class an image is from ten classes: airplane, automobile, bird, cat, deer, dog, frog, horse, ship, or truck."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "181d4dfbf917"
      },
      "source": [
        "### Costs\n",
        "\n",
        "This tutorial uses billable components of Google Cloud:\n",
        "\n",
        "* Vertex AI\n",
        "* Cloud Storage\n",
        "\n",
        "Learn about [Vertex AI\n",
        "pricing](https://cloud.google.com/vertex-ai/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": {
        "id": "install_aip:mbsdk"
      },
      "source": [
        "## Installation\n",
        "\n",
        "Install the packages required for executing this notebook."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "install_aip:mbsdk"
      },
      "outputs": [],
      "source": [
        "! pip3 install --upgrade --quiet google-cloud-aiplatform \\\n",
        "                                 google-cloud-storage \\\n",
        "                                 kfp \\\n",
        "                                 google-cloud-pipeline-components\n",
        "\n",
        "\n",
        "! pip3 install --upgrade --force-reinstall tensorflow kfp google-cloud-aiplatform google-cloud-storage google-cloud-pipeline-components -q"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "58707a750154"
      },
      "source": [
        "### Colab only: Uncomment the following cell to restart the kernel."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "f200f10a1da3"
      },
      "outputs": [],
      "source": [
        "# Automatically restart kernel after installs so that your environment can access the new packages\n",
        "# import IPython\n",
        "\n",
        "# app = IPython.Application.instance()\n",
        "# app.kernel.do_shutdown(True)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "BF1j6f9HApxa"
      },
      "source": [
        "## Before you begin\n",
        "\n",
        "### Set up your Google Cloud project\n",
        "\n",
        "**The following steps are required, regardless of your notebook environment.**\n",
        "\n",
        "1. [Select or create a Google Cloud project](https://console.cloud.google.com/cloud-resource-manager). When you first create an account, you get a $300 free credit towards your compute/storage costs.\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 Vertex AI API](https://console.cloud.google.com/flows/enableapi?apiid=aiplatform.googleapis.com )\n",
        "\n",
        "4. If you are running this notebook locally, you need to install the [Cloud SDK](https://cloud.google.com/sdk)."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "WReHDGG5g0XY"
      },
      "source": [
        "#### Set your project ID\n",
        "\n",
        "**If you don't know your project ID**, try the following:\n",
        "* Run `gcloud config list`.\n",
        "* Run `gcloud projects list`.\n",
        "* See the support page: [Locate the project ID](https://support.google.com/googleapi/answer/7014113)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "oM1iC_MfAts1"
      },
      "outputs": [],
      "source": [
        "PROJECT_ID = \"[your-project-id]\"  # @param {type:\"string\"}\n",
        "\n",
        "# Set the project id\n",
        "! gcloud config set project {PROJECT_ID}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "region"
      },
      "source": [
        "#### Region\n",
        "\n",
        "You can also change the `REGION` variable used by Vertex AI. Learn more about [Vertex AI regions](https://cloud.google.com/vertex-ai/docs/general/locations)."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "region"
      },
      "outputs": [],
      "source": [
        "REGION = \"us-central1\"  # @param {type: \"string\"}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "sBCra4QMA2wR"
      },
      "source": [
        "### Authenticate your Google Cloud account\n",
        "\n",
        "Depending on your Jupyter environment, you may have to manually authenticate. Follow the relevant instructions below.\n",
        "\n",
        "**1. Vertex AI Workbench**\n",
        "* Do nothing as you are already authenticated.\n",
        "\n",
        "**2. Local JupyterLab instance, uncomment and run:**"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "254614fa0c46"
      },
      "outputs": [],
      "source": [
        "# ! gcloud auth login"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ef21552ccea8"
      },
      "source": [
        "**3. Colab, uncomment and run:**"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "603adbbf0532"
      },
      "outputs": [],
      "source": [
        "# from google.colab import auth\n",
        "# auth.authenticate_user()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "f6b2ccc891ed"
      },
      "source": [
        "**4. Service account or other**\n",
        "* See how to grant Cloud Storage permissions to your service account at https://cloud.google.com/storage/docs/gsutil/commands/iam#ch-examples."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "zgPO1eR3CYjk"
      },
      "source": [
        "### Create a Cloud Storage bucket\n",
        "\n",
        "Create a storage bucket to store intermediate artifacts such as datasets."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "MzGDU7TWdts_"
      },
      "outputs": [],
      "source": [
        "BUCKET_URI = f\"gs://your-bucket-name-{PROJECT_ID}-unique\"  # @param {type:\"string\"}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "-EcIXiGsCePi"
      },
      "source": [
        "**If your bucket doesn't already exist**: Run the following cell to create your Cloud Storage bucket."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "NIq7R4HZCfIc"
      },
      "outputs": [],
      "source": [
        "! gsutil mb -l {REGION} -p {PROJECT_ID} {BUCKET_URI}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "set_service_account"
      },
      "source": [
        "#### Service Account\n",
        "\n",
        "**If you don't know your service account**, try to get your service account using `gcloud` command by executing the second cell below."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "set_service_account"
      },
      "outputs": [],
      "source": [
        "SERVICE_ACCOUNT = \"[your-service-account]\"  # @param {type:\"string\"}"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "autoset_service_account"
      },
      "outputs": [],
      "source": [
        "import sys\n",
        "\n",
        "IS_COLAB = \"google.colab\" in sys.modules\n",
        "if (\n",
        "    SERVICE_ACCOUNT == \"\"\n",
        "    or SERVICE_ACCOUNT is None\n",
        "    or SERVICE_ACCOUNT == \"[your-service-account]\"\n",
        "):\n",
        "    # Get your service account from gcloud\n",
        "    if not IS_COLAB:\n",
        "        shell_output = !gcloud auth list 2>/dev/null\n",
        "        SERVICE_ACCOUNT = shell_output[2].replace(\"*\", \"\").strip()\n",
        "\n",
        "    if IS_COLAB:\n",
        "        shell_output = ! gcloud projects describe  $PROJECT_ID\n",
        "        project_number = shell_output[-1].split(\":\")[1].strip().replace(\"'\", \"\")\n",
        "        SERVICE_ACCOUNT = f\"{project_number}-compute@developer.gserviceaccount.com\"\n",
        "\n",
        "    print(\"Service Account:\", SERVICE_ACCOUNT)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "set_service_account:pipelines"
      },
      "source": [
        "#### Set service account access for Vertex AI Pipelines\n",
        "\n",
        "Run the following commands to grant your service account access to read and write pipeline artifacts in the bucket that you created in the previous step -- you only need to run these once per service account."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "set_service_account:pipelines"
      },
      "outputs": [],
      "source": [
        "! gsutil iam ch serviceAccount:{SERVICE_ACCOUNT}:roles/storage.objectCreator $BUCKET_URI\n",
        "\n",
        "! gsutil iam ch serviceAccount:{SERVICE_ACCOUNT}:roles/storage.objectViewer $BUCKET_URI"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "setup_vars"
      },
      "source": [
        "### Set up variables\n",
        "\n",
        "Next, set up some variables used throughout the tutorial.\n",
        "### Import libraries and define constants"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "import_aip:mbsdk"
      },
      "outputs": [],
      "source": [
        "import google.cloud.aiplatform as aip\n",
        "from google_cloud_pipeline_components.v1.custom_job import utils\n",
        "from kfp import compiler, dsl\n",
        "from kfp.dsl import component"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "pipeline_constants"
      },
      "source": [
        "#### Vertex AI Pipelines constants\n",
        "\n",
        "Setup up the following constants for Vertex AI Pipelines:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "pipeline_constants"
      },
      "outputs": [],
      "source": [
        "PIPELINE_ROOT = \"{}/pipeline_root/bikes_weather\".format(BUCKET_URI)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "init_aip:mbsdk"
      },
      "source": [
        "## Initialize Vertex AI SDK for Python\n",
        "\n",
        "Initialize the Vertex AI SDK for Python for your project and corresponding bucket."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "yX-aimhGRRRl"
      },
      "outputs": [],
      "source": [
        "aip.init(project=PROJECT_ID, staging_bucket=BUCKET_URI)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "accelerators:training,cpu,prediction,cpu,mbsdk"
      },
      "source": [
        "#### Set hardware accelerators\n",
        "\n",
        "You can set hardware accelerators for training and prediction.\n",
        "\n",
        "Set the variables `TRAIN_GPU/TRAIN_NGPU` and `DEPLOY_GPU/DEPLOY_NGPU` to use a container image supporting a GPU and the number of GPUs allocated to the virtual machine (VM) instance. For example, to use a GPU container image with 4 Nvidia Telsa K80 GPUs allocated to each VM, you would specify:\n",
        "\n",
        "    (aip.gapic.AcceleratorType.NVIDIA_TESLA_K80, 4)\n",
        "\n",
        "\n",
        "Otherwise specify `(None, None)` to use a container image to run on a CPU.\n",
        "\n",
        "Learn more about [ hardware accelerator support for your region](https://cloud.google.com/vertex-ai/docs/general/locations#accelerators).\n",
        "\n",
        "*Note*: TF releases before 2.3 for GPU support will fail to load the custom model in this tutorial. It's a known issue and fixed in TF 2.3. This is caused by static graph ops that are generated in the serving function. If you encounter this issue on your own custom models, use a container image for TF 2.3 with GPU support."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ipG9uBUDRRRm"
      },
      "outputs": [],
      "source": [
        "TRAIN_GPU, TRAIN_NGPU = (None, None)\n",
        "\n",
        "DEPLOY_GPU, DEPLOY_NGPU = (None, None)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "container:training,prediction"
      },
      "source": [
        "#### Set pre-built containers\n",
        "\n",
        "Set the pre-built Docker container image for training and prediction.\n",
        "\n",
        "\n",
        "For the latest list, see [Pre-built containers for training](https://cloud.google.com/ai-platform-unified/docs/training/pre-built-containers).\n",
        "\n",
        "\n",
        "For the latest list, see [Pre-built containers for prediction](https://cloud.google.com/ai-platform-unified/docs/predictions/pre-built-containers)."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "WrKTNnVLRRRm"
      },
      "outputs": [],
      "source": [
        "TF = \"2-5\"\n",
        "\n",
        "if TRAIN_GPU:\n",
        "    TRAIN_VERSION = \"tf-gpu.{}\".format(TF)\n",
        "else:\n",
        "    TRAIN_VERSION = \"tf-cpu.{}\".format(TF)\n",
        "if DEPLOY_GPU:\n",
        "    DEPLOY_VERSION = \"tf2-gpu.{}\".format(TF)\n",
        "else:\n",
        "    DEPLOY_VERSION = \"tf2-cpu.{}\".format(TF)\n",
        "\n",
        "TRAIN_IMAGE = \"gcr.io/cloud-aiplatform/training/{}:latest\".format(TRAIN_VERSION)\n",
        "DEPLOY_IMAGE = \"gcr.io/cloud-aiplatform/prediction/{}:latest\".format(DEPLOY_VERSION)\n",
        "\n",
        "print(\"Training:\", TRAIN_IMAGE, TRAIN_GPU, TRAIN_NGPU)\n",
        "print(\"Deployment:\", DEPLOY_IMAGE, DEPLOY_GPU, DEPLOY_NGPU)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "machine:training,prediction"
      },
      "source": [
        "#### Set machine type\n",
        "\n",
        "Next, set the machine type to use for training and prediction.\n",
        "\n",
        "- Set the variables `TRAIN_COMPUTE` and `DEPLOY_COMPUTE` to configure  the compute resources for the VMs you'll use for for training and prediction.\n",
        " - `machine type`\n",
        "     - `n1-standard`: 3.75GB of memory per vCPU.\n",
        "     - `n1-highmem`: 6.5GB of memory per vCPU\n",
        "     - `n1-highcpu`: 0.9 GB of memory per vCPU\n",
        " - `vCPUs`: number of \\[2, 4, 8, 16, 32, 64, 96 \\]\n",
        "\n",
        "*Note: The following is not supported for training:*\n",
        "\n",
        " - `standard`: 2 vCPUs\n",
        " - `highcpu`: 2, 4 and 8 vCPUs\n",
        "\n",
        "*Note: You may also use n2 and e2 machine types for training and deployment, but they do not support GPUs*."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "K4RhNsDzRRRn"
      },
      "outputs": [],
      "source": [
        "MACHINE_TYPE = \"n1-standard\"\n",
        "\n",
        "VCPU = \"4\"\n",
        "TRAIN_COMPUTE = MACHINE_TYPE + \"-\" + VCPU\n",
        "print(\"Train machine type\", TRAIN_COMPUTE)\n",
        "\n",
        "VCPU = \"4\"\n",
        "DEPLOY_COMPUTE = MACHINE_TYPE + \"-\" + VCPU\n",
        "print(\"Deploy machine type\", DEPLOY_COMPUTE)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "tutorial_start:custom"
      },
      "source": [
        "# Tutorial\n",
        "\n",
        "Now you are ready to start creating your own custom model and training for CIFAR10."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "examine_training_package"
      },
      "source": [
        "### Examine the training package\n",
        "\n",
        "#### Package layout\n",
        "\n",
        "Before you start the training, you will look at how a Python package is assembled for a custom training job. When unarchived, the package contains the following directory/file layout.\n",
        "\n",
        "- PKG-INFO\n",
        "- README.md\n",
        "- setup.cfg\n",
        "- setup.py\n",
        "- trainer\n",
        "  - \\_\\_init\\_\\_.py\n",
        "  - task.py\n",
        "\n",
        "The files `setup.cfg` and `setup.py` are the instructions for installing the package into the operating environment of the Docker image.\n",
        "\n",
        "The file `trainer/task.py` is the Python script for executing the custom training job. *Note*, when we referred to it in the worker pool specification, we replace the directory slash with a dot (`trainer.task`) and dropped the file suffix (`.py`).\n",
        "\n",
        "#### Package Assembly\n",
        "\n",
        "In the following cells, you assemble the training package."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "YpA6MFcLRRRn"
      },
      "outputs": [],
      "source": [
        "# Make folder for Python training script\n",
        "! rm -rf custom\n",
        "! mkdir custom\n",
        "\n",
        "# Add package information\n",
        "! touch custom/README.md\n",
        "\n",
        "setup_cfg = \"[egg_info]\\n\\ntag_build =\\n\\ntag_date = 0\"\n",
        "! echo \"$setup_cfg\" > custom/setup.cfg\n",
        "\n",
        "setup_py = \"import setuptools\\n\\nsetuptools.setup(\\n\\n    install_requires=[\\n\\n        'tensorflow_datasets==1.3.0',\\n\\n    ],\\n\\n    packages=setuptools.find_packages())\"\n",
        "! echo \"$setup_py\" > custom/setup.py\n",
        "\n",
        "pkg_info = \"Metadata-Version: 1.0\\n\\nName: CIFAR10 image classification\\n\\nVersion: 0.0.0\\n\\nSummary: Demostration training script\\n\\nHome-page: www.google.com\\n\\nAuthor: Google\\n\\nAuthor-email: aferlitsch@google.com\\n\\nLicense: Public\\n\\nDescription: Demo\\n\\nPlatform: Vertex\"\n",
        "! echo \"$pkg_info\" > custom/PKG-INFO\n",
        "\n",
        "# Make the training subfolder\n",
        "! mkdir custom/trainer\n",
        "! touch custom/trainer/__init__.py"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "86f3ab08b20a"
      },
      "source": [
        "### Create a custom component for training the custom model\n",
        "\n",
        "Next, create a lightweight Python function component for training the CIFAR10 image classification model."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "2etUCVVMRRRo"
      },
      "outputs": [],
      "source": [
        "# Single, Mirror and Multi-Machine Distributed Training for CIFAR-10\n",
        "\n",
        "\n",
        "@component(\n",
        "    base_image=\"tensorflow/tensorflow:latest\",\n",
        "    packages_to_install=[\"tensorflow_datasets\", \"opencv-python-headless\"],\n",
        ")\n",
        "def custom_train_model(\n",
        "    model_dir: str,\n",
        "    lr: float = 0.01,\n",
        "    epochs: int = 10,\n",
        "    steps: int = 200,\n",
        "    distribute: str = \"single\",\n",
        "):\n",
        "\n",
        "    import faulthandler\n",
        "    import os\n",
        "    import sys\n",
        "\n",
        "    import tensorflow as tf\n",
        "    import tensorflow_datasets as tfds\n",
        "    from tensorflow.python.client import device_lib\n",
        "\n",
        "    faulthandler.enable()\n",
        "    tfds.disable_progress_bar()\n",
        "\n",
        "    print(\"Component start\")\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",
        "    print(\"DEVICES\", device_lib.list_local_devices())\n",
        "\n",
        "    # Single Machine, single compute device\n",
        "    if 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 distribute == \"mirror\":\n",
        "        strategy = tf.distribute.MirroredStrategy()\n",
        "    # Multiple Machine, multiple compute device\n",
        "    elif 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",
        "    # Preparing dataset\n",
        "    BUFFER_SIZE = 10000\n",
        "    BATCH_SIZE = 64\n",
        "\n",
        "    def make_datasets_unbatched():\n",
        "\n",
        "        # Scaling CIFAR10 data from (0, 255] to (0., 1.]\n",
        "        def scale(image, label):\n",
        "            image = tf.cast(image, tf.float32)\n",
        "            image /= 255.0\n",
        "            return image, label\n",
        "\n",
        "        datasets, info = tfds.load(name=\"cifar10\", with_info=True, as_supervised=True)\n",
        "        return datasets[\"train\"].map(scale).cache().shuffle(BUFFER_SIZE).repeat()\n",
        "\n",
        "    # Build the Keras model\n",
        "    def build_and_compile_cnn_model(lr: int = 0.01):\n",
        "        model = tf.keras.Sequential(\n",
        "            [\n",
        "                tf.keras.layers.Conv2D(\n",
        "                    32, 3, activation=\"relu\", input_shape=(32, 32, 3)\n",
        "                ),\n",
        "                tf.keras.layers.MaxPooling2D(),\n",
        "                tf.keras.layers.Conv2D(32, 3, activation=\"relu\"),\n",
        "                tf.keras.layers.MaxPooling2D(),\n",
        "                tf.keras.layers.Flatten(),\n",
        "                tf.keras.layers.Dense(10, activation=\"softmax\"),\n",
        "            ]\n",
        "        )\n",
        "        model.compile(\n",
        "            loss=tf.keras.losses.sparse_categorical_crossentropy,\n",
        "            optimizer=tf.keras.optimizers.SGD(learning_rate=lr),\n",
        "            metrics=[\"accuracy\"],\n",
        "        )\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",
        "    GLOBAL_BATCH_SIZE = BATCH_SIZE * NUM_WORKERS\n",
        "    train_dataset = make_datasets_unbatched().batch(GLOBAL_BATCH_SIZE)\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_cnn_model(lr)\n",
        "\n",
        "    model.fit(x=train_dataset, epochs=epochs, steps_per_epoch=steps)\n",
        "    model.save(model_dir)\n",
        "\n",
        "    model_path_to_deploy = model_dir\n",
        "\n",
        "    # Load the saved model\n",
        "    local_model = tf.keras.models.load_model(model_dir)\n",
        "\n",
        "    # Load evaluation data\n",
        "    import numpy as np\n",
        "    from tensorflow.keras.datasets import cifar10\n",
        "\n",
        "    (_, _), (x_test, y_test) = cifar10.load_data()\n",
        "    x_test = (x_test / 255.0).astype(np.float32)\n",
        "\n",
        "    print(x_test.shape, y_test.shape)\n",
        "\n",
        "    # Perform the model evaluation\n",
        "    local_model.evaluate(x_test, y_test)\n",
        "\n",
        "    # Serving function for image data\n",
        "    CONCRETE_INPUT = \"numpy_inputs\"\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=(32, 32))\n",
        "        return resized\n",
        "\n",
        "    @tf.function(input_signature=[tf.TensorSpec([None], tf.string)])\n",
        "    def preprocess_fn(bytes_inputs):\n",
        "        decoded_images = tf.map_fn(\n",
        "            _preprocess, bytes_inputs, dtype=tf.float32, back_prop=False\n",
        "        )\n",
        "        return {\n",
        "            CONCRETE_INPUT: decoded_images\n",
        "        }  # User needs to make sure the key matches model's input\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",
        "    m_call = tf.function(local_model.call).get_concrete_function(\n",
        "        [tf.TensorSpec(shape=[None, 32, 32, 3], dtype=tf.float32, name=CONCRETE_INPUT)]\n",
        "    )\n",
        "\n",
        "    tf.saved_model.save(\n",
        "        local_model, model_path_to_deploy, signatures={\"serving_default\": serving_fn}\n",
        "    )\n",
        "\n",
        "    # Get the serving function signature\n",
        "    loaded = tf.saved_model.load(model_path_to_deploy)\n",
        "\n",
        "    serving_input = list(\n",
        "        loaded.signatures[\"serving_default\"].structured_input_signature[1].keys()\n",
        "    )[0]\n",
        "    print(\"Serving function input:\", serving_input)\n",
        "\n",
        "    # Get test items\n",
        "    test_image_1 = x_test[0]\n",
        "    test_image_2 = x_test[1]\n",
        "    print(test_image_1.shape)\n",
        "\n",
        "    BUCKET_URI = model_dir + \"/test\"\n",
        "\n",
        "    import cv2\n",
        "\n",
        "    cv2.imwrite(\"tmp1.jpg\", (test_image_1 * 255).astype(np.uint8))\n",
        "    cv2.imwrite(\"tmp2.jpg\", (test_image_2 * 255).astype(np.uint8))\n",
        "\n",
        "    print(\"Writing jpg files\")\n",
        "\n",
        "    # Copy test item(s)\n",
        "    # For the batch prediction, copy the test items over to your Cloud Storage bucket.\n",
        "    test_item_1 = BUCKET_URI + \"/tmp1.jpg\"\n",
        "    test_item_2 = BUCKET_URI + \"/tmp2.jpg\"\n",
        "\n",
        "    with tf.io.gfile.GFile(test_item_1, \"wb\") as w:\n",
        "        with tf.io.gfile.GFile(\"tmp1.jpg\", \"rb\") as r:\n",
        "            bytes = r.read()\n",
        "            w.write(bytes)\n",
        "\n",
        "    with tf.io.gfile.GFile(test_item_2, \"wb\") as w:\n",
        "        with tf.io.gfile.GFile(\"tmp2.jpg\", \"rb\") as r:\n",
        "            bytes = r.read()\n",
        "            w.write(bytes)\n",
        "\n",
        "    # Make the batch input file\n",
        "    import base64\n",
        "    import json\n",
        "\n",
        "    gcs_input_uri = BUCKET_URI + \"/\" + \"test.jsonl\"\n",
        "    with tf.io.gfile.GFile(gcs_input_uri, \"w\") as f:\n",
        "        bytes = tf.io.read_file(test_item_1)\n",
        "        b64str = base64.b64encode(bytes.numpy()).decode(\"utf-8\")\n",
        "        data = {serving_input: {\"b64\": b64str}}\n",
        "        f.write(json.dumps(data) + \"\\n\")\n",
        "        bytes = tf.io.read_file(test_item_2)\n",
        "        b64str = base64.b64encode(bytes.numpy()).decode(\"utf-8\")\n",
        "        data = {serving_input: {\"b64\": b64str}}\n",
        "        f.write(json.dumps(data) + \"\\n\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "1abdbe48"
      },
      "source": [
        "### Convert the component to a Vertex AI Custom Job\n",
        "\n",
        "Next, use the `create_custom_training_job_op_from_component` method to convert the custom component into a Vertex AI custom job pre-built component.\n",
        "\n",
        "**replica_count :**  The  number of machine replicas the batch operation may be scaled to. Only used if machine_type is set. Default is 10."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "deb6c1cc"
      },
      "outputs": [],
      "source": [
        "custom_job_distributed_training_op = utils.create_custom_training_job_op_from_component(\n",
        "    custom_train_model, replica_count=1\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "e57e7311"
      },
      "source": [
        "### Define the pipeline for the custom training job\n",
        "\n",
        "Next, define the pipeline job, consisting of the tasks:\n",
        "\n",
        "- Train the custom model.\n",
        "- Upload the model to a Verex AI Model resource.\n",
        "- Execute a batch prediction."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "5R23d2J_HJr-"
      },
      "outputs": [],
      "source": [
        "MODEL_DIR = BUCKET_URI + \"/model\"\n",
        "\n",
        "\n",
        "@dsl.pipeline(name=\"custom-model-training-sample-pipeline\")\n",
        "def pipeline(\n",
        "    model_dir: str = MODEL_DIR,\n",
        "    lr: float = 0.01,\n",
        "    epochs: int = 10,\n",
        "    steps: int = 200,\n",
        "    distribute: str = \"single\",\n",
        "):\n",
        "    from google_cloud_pipeline_components.types import artifact_types\n",
        "    from google_cloud_pipeline_components.v1.batch_predict_job import \\\n",
        "        ModelBatchPredictOp\n",
        "    from google_cloud_pipeline_components.v1.model import ModelUploadOp\n",
        "    from kfp.dsl import importer_node\n",
        "\n",
        "    custom_producer_task = custom_job_distributed_training_op(\n",
        "        model_dir=model_dir,\n",
        "        lr=lr,\n",
        "        epochs=epochs,\n",
        "        steps=steps,\n",
        "        distribute=distribute,\n",
        "        project=PROJECT_ID,\n",
        "        location=REGION,\n",
        "        base_output_directory=PIPELINE_ROOT,\n",
        "    )\n",
        "\n",
        "    unmanaged_model_importer = importer_node.importer(\n",
        "        artifact_uri=model_dir,\n",
        "        artifact_class=artifact_types.UnmanagedContainerModel,\n",
        "        metadata={\n",
        "            \"containerSpec\": {\n",
        "                \"imageUri\": \"us-docker.pkg.dev/cloud-aiplatform/prediction/tf2-cpu.2-3:latest\"\n",
        "            }\n",
        "        },\n",
        "    )\n",
        "\n",
        "    model_upload_op = ModelUploadOp(\n",
        "        project=PROJECT_ID,\n",
        "        display_name=\"model_display_name\",\n",
        "        unmanaged_container_model=unmanaged_model_importer.outputs[\"artifact\"],\n",
        "    )\n",
        "    model_upload_op.after(custom_producer_task)\n",
        "\n",
        "    batch_predict_op = ModelBatchPredictOp(\n",
        "        project=PROJECT_ID,\n",
        "        job_display_name=\"batch_predict_job\",\n",
        "        model=model_upload_op.outputs[\"model\"],\n",
        "        gcs_source_uris=[MODEL_DIR + \"/test/test.jsonl\"],\n",
        "        gcs_destination_output_uri_prefix=PIPELINE_ROOT,\n",
        "        instances_format=\"jsonl\",\n",
        "        predictions_format=\"jsonl\",\n",
        "        model_parameters={},\n",
        "        machine_type=DEPLOY_COMPUTE,\n",
        "        starting_replica_count=1,\n",
        "        max_replica_count=1,\n",
        "    )\n",
        "\n",
        "    batch_predict_op.after(model_upload_op)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "e584790f"
      },
      "source": [
        "### Compile and run the pipeline\n",
        "Next, compile the pipeline into a DAG and then exeute it."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "6be28881"
      },
      "outputs": [],
      "source": [
        "compiler.Compiler().compile(\n",
        "    pipeline_func=pipeline, package_path=\"custom_model_training_spec.yaml\"\n",
        ")\n",
        "\n",
        "DISPLAY_NAME = \"cifar10\"\n",
        "\n",
        "job = aip.PipelineJob(\n",
        "    display_name=DISPLAY_NAME,\n",
        "    template_path=\"custom_model_training_spec.yaml\",\n",
        "    pipeline_root=PIPELINE_ROOT,\n",
        ")\n",
        "\n",
        "job.run(service_account=SERVICE_ACCOUNT)\n",
        "\n",
        "! rm custom_model_training_spec.json"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "87774222c338"
      },
      "source": [
        "### View custom training pipeline results\n",
        "\n",
        "Finally, you view the artifact outputs of each task in the pipeline."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ac26e1314c80"
      },
      "outputs": [],
      "source": [
        "import json\n",
        "\n",
        "import tensorflow as tf\n",
        "\n",
        "PROJECT_NUMBER = job.gca_resource.name.split(\"/\")[1]\n",
        "print(PROJECT_NUMBER)\n",
        "\n",
        "\n",
        "def print_pipeline_output(job, output_task_name):\n",
        "    JOB_ID = job.name\n",
        "    print(JOB_ID)\n",
        "    for _ in range(len(job.gca_resource.job_detail.task_details)):\n",
        "        TASK_ID = job.gca_resource.job_detail.task_details[_].task_id\n",
        "        EXECUTE_OUTPUT = (\n",
        "            PIPELINE_ROOT\n",
        "            + \"/\"\n",
        "            + PROJECT_NUMBER\n",
        "            + \"/\"\n",
        "            + JOB_ID\n",
        "            + \"/\"\n",
        "            + output_task_name\n",
        "            + \"_\"\n",
        "            + str(TASK_ID)\n",
        "            + \"/executor_output.json\"\n",
        "        )\n",
        "        GCP_RESOURCES = (\n",
        "            PIPELINE_ROOT\n",
        "            + \"/\"\n",
        "            + PROJECT_NUMBER\n",
        "            + \"/\"\n",
        "            + JOB_ID\n",
        "            + \"/\"\n",
        "            + output_task_name\n",
        "            + \"_\"\n",
        "            + str(TASK_ID)\n",
        "            + \"/gcp_resources\"\n",
        "        )\n",
        "        EVAL_METRICS = (\n",
        "            PIPELINE_ROOT\n",
        "            + \"/\"\n",
        "            + PROJECT_NUMBER\n",
        "            + \"/\"\n",
        "            + JOB_ID\n",
        "            + \"/\"\n",
        "            + output_task_name\n",
        "            + \"_\"\n",
        "            + str(TASK_ID)\n",
        "            + \"/evaluation_metrics\"\n",
        "        )\n",
        "        if tf.io.gfile.exists(EXECUTE_OUTPUT):\n",
        "            ! gsutil cat $EXECUTE_OUTPUT\n",
        "            return EXECUTE_OUTPUT\n",
        "        elif tf.io.gfile.exists(GCP_RESOURCES):\n",
        "            ! gsutil cat $GCP_RESOURCES\n",
        "            return GCP_RESOURCES\n",
        "        elif tf.io.gfile.exists(EVAL_METRICS):\n",
        "            ! gsutil cat $EVAL_METRICS\n",
        "            return EVAL_METRICS\n",
        "\n",
        "    return None\n",
        "\n",
        "\n",
        "print(\"model-upload\")\n",
        "artifacts = print_pipeline_output(job, \"model-upload\")\n",
        "print(\"\\n\")\n",
        "output = !gsutil cat $artifacts\n",
        "print(output)\n",
        "output = json.loads(output[0])\n",
        "model_id = output[\"artifacts\"][\"model\"][\"artifacts\"][0][\"metadata\"][\"resourceName\"]\n",
        "print(\"model-batch-predict\")\n",
        "artifacts = print_pipeline_output(job, \"model-batch-predict\")\n",
        "print(\"\\n\")\n",
        "output = !gsutil cat $artifacts\n",
        "output = json.loads(output[0])\n",
        "batch_job_id = output[\"artifacts\"][\"batchpredictionjob\"][\"artifacts\"][0][\"metadata\"][\n",
        "    \"resourceName\"\n",
        "]"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "a2cd7ee5cc21"
      },
      "source": [
        "## Cleaning up\n",
        "\n",
        "To clean up all Google Cloud resources used in this project, you can [delete the Google Cloud\n",
        "project](https://cloud.google.com/resource-manager/docs/creating-managing-projects#shutting_down_projects) you used for the tutorial.\n",
        "\n",
        "Otherwise, you can delete the individual resources you created in this tutorial."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "c29a5f39e211"
      },
      "outputs": [],
      "source": [
        "import os\n",
        "\n",
        "job.delete()\n",
        "\n",
        "model = aip.Model(model_id)\n",
        "model.delete()\n",
        "\n",
        "batch_job = aip.BatchPredictionJob(batch_job_id)\n",
        "batch_job.delete()\n",
        "\n",
        "\n",
        "delete_bucket = False\n",
        "if delete_bucket or os.getenv(\"IS_TESTING\"):\n",
        "    ! gsutil rm -rf {BUCKET_URI}"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "collapsed_sections": [],
      "name": "custom_model_training_and_batch_prediction.ipynb",
      "toc_visible": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}