notebooks/community/model_garden/model_garden_mediapipe_image_classification.ipynb

{
  "cells": [
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ur8xi4C7S06n"
      },
      "outputs": [],
      "source": [
        "# Copyright 2023 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": "TirJ-SGQseby"
      },
      "source": [
        "# Vertex AI Model Garden MediaPipe with image classification\n",
        "\n",
        "<table align=\"left\">\n",
        "  <td>\n",
        "    <a href=\"https://colab.research.google.com/github/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/community/model_garden/model_garden_mediapipe_image_classification.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",
        "\n",
        "  <td>\n",
        "    <a href=\"https://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/community/model_garden/model_garden_mediapipe_image_classification.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/notebooks/deploy-notebook?download_url=https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/main/notebooks/community/model_garden/model_garden_mediapipe_image_classification.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>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "dwGLvtIeECLK"
      },
      "source": [
        "**_NOTE_**: This notebook has been tested in the following environment:\n",
        "\n",
        "* Python version = 3.9\n",
        "\n",
        "**NOTE**: The checkpoint and the dataset linked in this Colab are not owned or distributed by Google, and are made available by third parties. Please review the terms and conditions made available by the third parties before using the checkpoint and data."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "tvgnzT1CKxrO"
      },
      "source": [
        "## Overview\n",
        "\n",
        "This notebook demonstrates how to use [MediaPipe Model Maker](https://developers.google.com/mediapipe/solutions/model_maker) to train an on-device image classification model in Vertex AI Model Garden.\n",
        "\n",
        "### Objective\n",
        "\n",
        "* Train new models\n",
        "  * Convert input data to training formats\n",
        "  * Create [custom jobs](https://cloud.google.com/vertex-ai/docs/training/create-custom-job) to train new models\n",
        "  * Export models\n",
        "\n",
        "* Cleanup resources\n",
        "\n",
        "### 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": "KEukV6uRk_S3"
      },
      "source": [
        "## Before you begin"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "z__i0w0lCAsW"
      },
      "source": [
        "### Colab only\n",
        "Run the following commands to install dependencies and to authenticate with Google Cloud if running on Colab."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Jvqs-ehKlaYh"
      },
      "outputs": [],
      "source": [
        "! pip3 install --upgrade pip\n",
        "\n",
        "import sys\n",
        "\n",
        "if \"google.colab\" in sys.modules:\n",
        "    ! pip3 install --upgrade google-cloud-aiplatform\n",
        "\n",
        "    # Automatically restart kernel after installs\n",
        "    import IPython\n",
        "\n",
        "    app = IPython.Application.instance()\n",
        "    app.kernel.do_shutdown(True)\n",
        "\n",
        "    from google.colab import auth as google_auth\n",
        "\n",
        "    google_auth.authenticate_user()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "WReHDGG5g0XY"
      },
      "source": [
        "#### Set your project ID\n",
        "\n",
        "**If you don't know your project ID**, 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": "tTy1gX11kCJY"
      },
      "outputs": [],
      "source": [
        "REGION = \"us-central1\"  # @param {type: \"string\"}\n",
        "REGION_PREFIX = REGION.split(\"-\")[0]\n",
        "assert REGION_PREFIX in (\n",
        "    \"us\",\n",
        "    \"europe\",\n",
        "    \"asia\",\n",
        "), f'{REGION} is not supported. It must be prefixed by \"us\", \"asia\", or \"europe\".'"
      ]
    },
    {
      "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": [
        "**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": {
        "id": "NIq7R4HZCfIc"
      },
      "outputs": [],
      "source": [
        "! gsutil mb -l {REGION} -p {PROJECT_ID} {BUCKET_URI}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "960505627ddf"
      },
      "source": [
        "### Import libraries"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "PyQmSRbKA8r-"
      },
      "outputs": [],
      "source": [
        "import json\n",
        "import os\n",
        "from datetime import datetime\n",
        "\n",
        "import tensorflow\n",
        "from google.cloud import aiplatform"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "init_aip:mbsdk,all"
      },
      "source": [
        "### Initialize Vertex AI SDK for Python\n",
        "\n",
        "Initialize the Vertex AI SDK for Python for your project."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "9wExiMUxFk91"
      },
      "outputs": [],
      "source": [
        "now = datetime.now().strftime(\"%Y%m%d-%H%M%S\")\n",
        "\n",
        "STAGING_BUCKET = os.path.join(BUCKET_URI, \"temp/%s\" % now)\n",
        "\n",
        "EVALUATION_RESULT_OUTPUT_DIRECTORY = os.path.join(STAGING_BUCKET, \"evaluation\")\n",
        "EVALUATION_RESULT_OUTPUT_FILE = os.path.join(\n",
        "    EVALUATION_RESULT_OUTPUT_DIRECTORY, \"evaluation.json\"\n",
        ")\n",
        "\n",
        "EXPORTED_MODEL_OUTPUT_DIRECTORY = os.path.join(STAGING_BUCKET, \"model\")\n",
        "EXPORTED_MODEL_OUTPUT_FILE = os.path.join(\n",
        "    EXPORTED_MODEL_OUTPUT_DIRECTORY, \"model.tflite\"\n",
        ")\n",
        "\n",
        "aiplatform.init(project=PROJECT_ID, location=REGION, staging_bucket=STAGING_BUCKET)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "n6IFz75WGCam"
      },
      "source": [
        "### Define training machine specs"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "riG_qUokg0XZ"
      },
      "outputs": [],
      "source": [
        "TRAINING_JOB_DISPLAY_NAME = \"mediapipe_image_classifier_%s\" % now\n",
        "TRAINING_CONTAINER = f\"{REGION_PREFIX}-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/mediapipe-train\"\n",
        "TRAINING_MACHINE_TYPE = \"n1-highmem-16\"\n",
        "TRAINING_ACCELERATOR_TYPE = \"NVIDIA_TESLA_V100\"\n",
        "TRAINING_ACCELERATOR_COUNT = 2"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "-rsdAcBV-vlf"
      },
      "source": [
        "## Train your customized models"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "zgPO1eR3CYjk"
      },
      "source": [
        "### Prepare input data for training\n",
        "\n",
        "Finetuning a model for image classification requires a dataset that includes all kinds of items, or classes, that you want the completed model to be able to identify. You can do this by trimming down a public dataset to only the classes that are relevant to your usecase, compiling your own data, or some combination of both. The dataset can be significantly smaller than what would be required to train a new model from scratch. For example, the [ImageNet](https://www.image-net.org/) dataset used to train many reference models contains millions of images with thousands of categories. Transfer learning with Model Maker can finetune an existing model with a smaller dataset and still perform well, depending on your inference accuracy goals.\n",
        "\n",
        "You can re-use an existing dataset such as `gs://cloud-samples-data-us-central1/vision/automl_classification/flowers` to finetune the model or you can upload your own dataset to GCS. If you are using your own dataset, ensure that your image directory contains several subdirectories, each corresponding to specific class labels. Your training data should also follow this pattern: <image_path>/<label_name>/<image_names>.*."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "IndQ_m6ddUEM"
      },
      "outputs": [],
      "source": [
        "training_data_path = \"gs://cloud-samples-data-us-central1/vision/automl_classification/flowers\"  # @param {type:\"string\"}\n",
        "split_ratio = \"0.8,0.1,0.1\"  # @param {type:\"string\"}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "aaff6f5be7f6"
      },
      "source": [
        "### Set fine-tuning options\n",
        "\n",
        "You can pick between different model architectures to further customize your training:\n",
        "\n",
        "*   MobileNet-V2\n",
        "*   EfficientNet-Lite0\n",
        "*   EfficientNet-Lite2\n",
        "*   EfficientNet-Lite4\n",
        "\n",
        "To set the model architecture and other training parameters, adjust the following values:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "um_XKbmpTaHx"
      },
      "outputs": [],
      "source": [
        "model_architecture = \"mobilenet_v2\"  # @param [\"mobilenet_v2\", \"efficientnet_lite0\", \"efficientnet_lite2\", \"efficientnet_lite4\"]\n",
        "\n",
        "# The learning rate to use for gradient descent training.\n",
        "learning_rate: float = 0.01  # @param {type:\"number\"}\n",
        "# Batch size for training.\n",
        "batch_size: int = 2  # @param {type:\"number\"}\n",
        "# Number of training iterations over the dataset.\n",
        "epochs: int = 10  # @param {type:\"slider\", min:0, max:100, step:1}\n",
        "# If true, the base module is trained together with the classification layer on\n",
        "# top.\n",
        "do_fine_tuning: bool = False  # @param {type:\"boolean\"}\n",
        "# A regularizer that applies a L1 regularization penalty.\n",
        "l1_regularizer: float = 0.0  # @param {type:\"number\"}\n",
        "# A regularizer that applies a L2 regularization penalty.\n",
        "l2_regularizer: float = 0.0001  # @param {type:\"number\"}\n",
        "# Amount of label smoothing to apply. See tf.keras.losses for more details.\n",
        "label_smoothing: float = 0.1  # @param {type:\"number\"}\n",
        "# A boolean controlling whether the training dataset is augmented by randomly\n",
        "# distorting input images, including random cropping, flipping, etc. See\n",
        "# utils.image_preprocessing documentation for details.\n",
        "do_data_augmentation: bool = True  # @param {type:\"boolean\"}\n",
        "# Number of training samples used to calculate the decay steps\n",
        "# and create the training optimizer.\n",
        "decay_samples: int = 2560000  # @param {type:\"number\"}\n",
        "# Number of warmup steps for a linear increasing warmup schedule on learning\n",
        "# rate. Used to set up warmup schedule by model_util.WarmUp.\n",
        "warmup_epochs: int = 2  # @param {type:\"number\"}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "HwcCjwlBTQIz"
      },
      "source": [
        "### Run fine-tuning\n",
        "With your training dataset and fine-tuning options prepared, you are ready to start the fine-tuning process. This process is resource intensive and can take a few minutes to a few hours depending on your available compute resources. On Vertex AI with GPU processing, the example fine-tuning below takes between 4-6 minutes to train on approximately 3700 images.\n",
        "\n",
        "To begin the fine-tuning process, use the following code:\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "aec22792ee84"
      },
      "outputs": [],
      "source": [
        "model_export_path = EXPORTED_MODEL_OUTPUT_DIRECTORY\n",
        "evaluation_result_path = EVALUATION_RESULT_OUTPUT_DIRECTORY\n",
        "\n",
        "worker_pool_specs = [\n",
        "    {\n",
        "        \"machine_spec\": {\n",
        "            \"machine_type\": TRAINING_MACHINE_TYPE,\n",
        "            \"accelerator_type\": TRAINING_ACCELERATOR_TYPE,\n",
        "            \"accelerator_count\": TRAINING_ACCELERATOR_COUNT,\n",
        "        },\n",
        "        \"replica_count\": 1,\n",
        "        \"container_spec\": {\n",
        "            \"image_uri\": TRAINING_CONTAINER,\n",
        "            \"command\": [],\n",
        "            \"args\": [\n",
        "                \"--task_name=image_classifier\",\n",
        "                \"--training_data_path=%s\" % training_data_path,\n",
        "                \"--model_export_path=%s\" % model_export_path,\n",
        "                \"--evaluation_result_path=%s\" % evaluation_result_path,\n",
        "                \"--split_ratio=%s\" % split_ratio,\n",
        "                \"--model_architecture=%s\" % model_architecture,\n",
        "                \"--hparams=%s\"\n",
        "                % json.dumps(\n",
        "                    {\n",
        "                        \"learning_rate\": learning_rate,\n",
        "                        \"batch_size\": batch_size,\n",
        "                        \"epochs\": epochs,\n",
        "                        \"do_fine_tuning\": do_fine_tuning,\n",
        "                        \"l1_regularizer\": l1_regularizer,\n",
        "                        \"l2_regularizer\": l2_regularizer,\n",
        "                        \"label_smoothing\": label_smoothing,\n",
        "                        \"do_data_augmentation\": do_data_augmentation,\n",
        "                        \"decay_samples\": decay_samples,\n",
        "                        \"warmup_epochs\": warmup_epochs,\n",
        "                    }\n",
        "                ),\n",
        "            ],\n",
        "        },\n",
        "    }\n",
        "]\n",
        "\n",
        "training_job = aiplatform.CustomJob(\n",
        "    display_name=TRAINING_JOB_DISPLAY_NAME,\n",
        "    project=PROJECT_ID,\n",
        "    worker_pool_specs=worker_pool_specs,\n",
        "    staging_bucket=STAGING_BUCKET,\n",
        ")\n",
        "\n",
        "training_job.run()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "rXMF2tnV_WS0"
      },
      "source": [
        "## Evaluate and export model"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "mV-Djz-frBni"
      },
      "source": [
        "### Evaluate performance\n",
        "\n",
        "After fine-tuning the model, we evaluate the training result on a test dataset, which is typically a portion of your original dataset not used during training. Accuracy levels between 0.8 and 0.9 are generally considered very good, but your use case requirements may differ. You should also consider how fast the model can produce an inference. Higher accuracy frequently comes at the cost of longer inference times.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "09Rz1AYspK19"
      },
      "outputs": [],
      "source": [
        "def get_evaluation_result(evaluation_result_path):\n",
        "    try:\n",
        "        with tensorflow.io.gfile.GFile(evaluation_result_path, \"r\") as input_file:\n",
        "            evalutation_result = json.loads(input_file.read())\n",
        "        return evalutation_result[\"accuracy\"], evalutation_result[\"loss\"]\n",
        "    except:\n",
        "        print(\n",
        "            \"Evaluation result not found. Your test dataset is likely \"\n",
        "            + \"empty. You can adjust the size of your test dataset or adjust \"\n",
        "            + \"how you split your dataset.\"\n",
        "        )\n",
        "        return None\n",
        "\n",
        "\n",
        "evaluation_result = get_evaluation_result(EVALUATION_RESULT_OUTPUT_FILE)\n",
        "\n",
        "if evaluation_result is not None:\n",
        "    print(\"Accuracy:\", evaluation_result[0])\n",
        "    print(\"Loss:\", evaluation_result[1])"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "g0BGaofgsMsy"
      },
      "source": [
        "### Export model\n",
        "After finetuning and evaluating the model, you can save the Tensorflow Lite model, try it out in the [Image Classification](https://mediapipe-studio.webapps.google.com/demo/image_classifier) demo in MediaPipe Studio or integrate it with your on-device application by following the [Image classification task guide](https://developers.google.com/mediapipe/solutions/vision/image_classifier). The exported model contains the generates required model metadata, as well as a classification label file."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "NYuQowyZEtxK"
      },
      "outputs": [],
      "source": [
        "import sys\n",
        "\n",
        "def copy_model(model_source, model_dest):\n",
        "    ! gsutil cp {model_source} {model_dest}\n",
        "\n",
        "copy_model(EXPORTED_MODEL_OUTPUT_FILE, \"image_classification_model.tflite\")\n",
        "\n",
        "if \"google.colab\" in sys.modules:\n",
        "    from google.colab import files\n",
        "\n",
        "    files.download(\"image_classification_model.tflite\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "kkH2nrpdp4sp"
      },
      "source": [
        "## Clean up"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Ax6vQVZhp9pR"
      },
      "outputs": [],
      "source": [
        "# Delete training data and jobs.\n",
        "if training_job.list(filter=f'display_name=\"{TRAINING_JOB_DISPLAY_NAME}\"'):\n",
        "    training_job.delete()\n",
        "\n",
        "!gsutil rm -r {STAGING_BUCKET}"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "name": "model_garden_mediapipe_image_classification.ipynb",
      "toc_visible": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}