notebooks/community/model_garden/model_garden_movinet_action_recognition.ipynb

{
  "cells": [
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ur8xi4C7S06n"
      },
      "outputs": [],
      "source": [
        "# Copyright 2024 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."
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {
        "id": "TirJ-SGQseby"
      },
      "source": [
        "# Vertex AI Model Garden MoViNet video action recognition\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_movinet_action_recognition.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_movinet_action_recognition.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>                                                                                               <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_movinet_action_recognition.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>"
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {
        "id": "dwGLvtIeECLK"
      },
      "source": [
        "**_NOTE_**: This notebook has been tested in the following environment:\n",
        "\n",
        "* Python version = 3.9"
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {
        "id": "tvgnzT1CKxrO"
      },
      "source": [
        "## Overview\n",
        "\n",
        "This notebook demonstrates how to use [MoViNet](https://github.com/tensorflow/models/tree/master/official/projects/movinet) for video action recognition in Vertex AI Model Garden.\n",
        "\n",
        "### Objective\n",
        "\n",
        "* Train new models\n",
        "  * Convert input data to training formats\n",
        "  * Create [hyperparameter tuning jobs](https://cloud.google.com/vertex-ai/docs/training/hyperparameter-tuning-overview) to train new models\n",
        "  * Find and export best models\n",
        "\n",
        "* Test trained models\n",
        "  * Upload models to the [Vertex AI Model Registry](https://cloud.google.com/vertex-ai/docs/model-registry/introduction)\n",
        "  * Run batch predictions\n",
        "\n",
        "* Clean up 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."
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {
        "id": "KEukV6uRk_S3"
      },
      "source": [
        "## Before you begin"
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {
        "id": "z__i0w0lCAsW"
      },
      "source": [
        "### Colab Only\n",
        "Run the following commands for Colab or skip this section if you use Workbench."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Jvqs-ehKlaYh"
      },
      "outputs": [],
      "source": [
        "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()"
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {
        "id": "BF1j6f9HApxa"
      },
      "source": [
        "### 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",
        "1. [Make sure that billing is enabled for your project](https://cloud.google.com/billing/docs/how-to/modify-project).\n",
        "\n",
        "1. [Enable the Vertex AI API and Compute Engine API](https://console.cloud.google.com/flows/enableapi?apiid=aiplatform.googleapis.com,compute_component).\n",
        "1. If you are running this notebook locally, you will need to install the [Cloud SDK](https://cloud.google.com/sdk).\n",
        "\n",
        "1. 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",
        "1. [Create a service account](https://cloud.google.com/iam/docs/service-accounts-create#iam-service-accounts-create-console) with `Vertex AI User` and `Storage Object Admin` roles for running batch predictions with the fine tuned model.\n",
        "\n",
        "**Note**: Jupyter runs lines prefixed with `!` as shell commands, and it interpolates Python variables prefixed with `$` into these commands."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "9wExiMUxFk91"
      },
      "outputs": [],
      "source": [
        "import os\n",
        "\n",
        "from google.cloud import aiplatform\n",
        "\n",
        "# The GCP project ID for experiments.\n",
        "PROJECT_ID = \"\"  # @param {type:\"string\"}\n",
        "\n",
        "# Bucket URI with gs:// prefix.\n",
        "BUCKET_URI = \"\"  # @param {type:\"string\"}\n",
        "\n",
        "# You can choose a region from https://cloud.google.com/about/locations.\n",
        "# Only regions prefixed by \"us\", \"asia\", or \"europe\" are supported.\n",
        "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\".'\n",
        "\n",
        "! gcloud config set project $PROJECT_ID\n",
        "\n",
        "STAGING_BUCKET = os.path.join(BUCKET_URI, \"temporal\")\n",
        "CHECKPOINT_BUCKET = os.path.join(BUCKET_URI, \"ckpt\")\n",
        "\n",
        "aiplatform.init(project=PROJECT_ID, location=REGION, staging_bucket=STAGING_BUCKET)\n",
        "\n",
        "# Download config files.\n",
        "CONFIG_DIR = os.path.join(BUCKET_URI, \"config\")"
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {
        "id": "n6IFz75WGCam"
      },
      "source": [
        "### Define constants"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "riG_qUokg0XZ"
      },
      "outputs": [],
      "source": [
        "OBJECTIVE = \"var\"\n",
        "\n",
        "# Data converter constants.\n",
        "DATA_CONVERTER_JOB_PREFIX = \"data_converter\"\n",
        "DATA_CONVERTER_CONTAINER = f\"{REGION_PREFIX}-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/data-converter\"\n",
        "DATA_CONVERTER_MACHINE_TYPE = \"n1-highmem-8\"\n",
        "IMAGE_SIZES = {\n",
        "    \"a0\": 172,\n",
        "    \"a1\": 172,\n",
        "    \"a2\": 224,\n",
        "    \"a3\": 256,\n",
        "    \"a4\": 290,\n",
        "    \"a5\": 320,\n",
        "}\n",
        "\n",
        "# Training constants.\n",
        "TRAINING_JOB_PREFIX = \"train\"\n",
        "TRAIN_CONTAINER_URI = f\"{REGION_PREFIX}-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/movinet-train\"\n",
        "TRAIN_MACHINE_TYPE = \"n1-highmem-32\"\n",
        "TRAIN_ACCELERATOR_TYPE = \"NVIDIA_TESLA_V100\"\n",
        "TRAIN_NUM_GPU = 8\n",
        "\n",
        "# Evaluation constants.\n",
        "EVALUATION_METRIC = \"accuracy\"\n",
        "\n",
        "# Export constants.\n",
        "EXPORT_JOB_PREFIX = \"export\"\n",
        "EXPORT_CONTAINER_URI = f\"{REGION_PREFIX}-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/movinet-model-export\"\n",
        "EXPORT_MACHINE_TYPE = \"n1-highmem-8\"\n",
        "\n",
        "# Prediction constants.\n",
        "# You can adjust accelerator types and machine types to get faster predictions.\n",
        "PREDICTION_CONTAINER_URI = f\"{REGION_PREFIX}-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/movinet-serve\"\n",
        "PREDICTION_PORT = 8080\n",
        "PREDICTION_ACCELERATOR_COUNT = 1\n",
        "PREDICTION_ACCELERATOR_TYPE = \"NVIDIA_TESLA_T4\"\n",
        "PREDICTION_MACHINE_TYPE = \"n1-standard-4\"\n",
        "PREDICTION_JOB_PREFIX = \"predict\""
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {
        "id": "ZZFPe_GezXg8"
      },
      "source": [
        "### Define common helper functions"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "XcYUGwr-AJGY"
      },
      "outputs": [],
      "source": [
        "import json\n",
        "from datetime import datetime\n",
        "from typing import Any\n",
        "\n",
        "import tensorflow as tf\n",
        "import yaml\n",
        "\n",
        "\n",
        "def get_job_name_with_datetime(prefix: str) -> str:\n",
        "    \"\"\"Returns a timestamped job name with the given prefix.\"\"\"\n",
        "    return prefix + datetime.now().strftime(\"_%Y%m%d_%H%M%S\")\n",
        "\n",
        "\n",
        "def get_label_map(label_map_yaml_filepath: str) -> tuple[dict[int, str], int]:\n",
        "    \"\"\"Reads label map from a YAML file and returns the label map with the number of classes.\"\"\"\n",
        "    with tf.io.gfile.GFile(label_map_yaml_filepath, \"rb\") as input_file:\n",
        "        label_map = yaml.safe_load(input_file.read())[\"label_map\"]\n",
        "    num_classes = max(label_map.keys()) + 1\n",
        "    return label_map, num_classes\n",
        "\n",
        "\n",
        "def get_best_trial(\n",
        "    model_di: str, max_trial_count: int, evaluation_metric: str\n",
        ") -> tuple[str, Any]:\n",
        "    \"\"\"Finds the best trial directory and eval results from a hyperparameter tuning job.\"\"\"\n",
        "    best_trial_dir = \"\"\n",
        "    best_trial_evaluation_results = {}\n",
        "    best_performance = -1\n",
        "\n",
        "    for i in range(max_trial_count):\n",
        "        current_trial = i + 1\n",
        "        current_trial_dir = os.path.join(model_dir, \"trial_\" + str(current_trial))\n",
        "        current_trial_best_ckpt_dir = os.path.join(current_trial_dir, \"best_ckpt\")\n",
        "        current_trial_best_ckpt_evaluation_filepath = os.path.join(\n",
        "            current_trial_best_ckpt_dir, \"info.json\"\n",
        "        )\n",
        "        with tf.io.gfile.GFile(current_trial_best_ckpt_evaluation_filepath, \"rb\") as f:\n",
        "            eval_metric_results = json.load(f)\n",
        "            current_performance = eval_metric_results[evaluation_metric]\n",
        "            if current_performance > best_performance:\n",
        "                best_performance = current_performance\n",
        "                best_trial_dir = current_trial_dir\n",
        "                best_trial_evaluation_results = eval_metric_results\n",
        "    return best_trial_dir, best_trial_evaluation_results\n",
        "\n",
        "\n",
        "def print_response_instance(json_str: str, label_map: dict[int, str]):\n",
        "    \"\"\"Prints summary of a prediction JSON result from the model response.\"\"\"\n",
        "    json_obj = json.loads(json_str)\n",
        "    if \"prediction\" not in json_obj:\n",
        "        print(\"Error:\", json_str)\n",
        "        return\n",
        "    instance = json_obj[\"instance\"]\n",
        "    prediction = json_obj[\"prediction\"]\n",
        "    gcs_uri = instance[\"content\"]\n",
        "    time_start = instance.get(\"timeSegmentStart\", \"0.0s\")\n",
        "    time_end = instance.get(\"timeSegmentEnd\", \"Infinity\")\n",
        "    print(f\"---------- Predict {gcs_uri}, {time_start} to {time_end}:\")\n",
        "    for predicted in prediction:\n",
        "        time = predicted[\"timeSegmentStart\"]\n",
        "        label = label_map[predicted[\"label\"]]\n",
        "        confidence = predicted[\"confidence\"]\n",
        "        print(f\"At {time}, detected {label} with {confidence} confidence.\")\n",
        "\n",
        "\n",
        "def find_checkpoint_in_dir(checkpoint_dir: str) -> str:\n",
        "    \"\"\"Finds a checkpoint path relative to the directory.\"\"\"\n",
        "    for root, dirs, files in tf.io.gfile.walk(checkpoint_dir):\n",
        "        for file in files:\n",
        "            if file.endswith(\".index\"):\n",
        "                return os.path.join(root, os.path.splitext(file)[0])\n",
        "\n",
        "\n",
        "def upload_checkpoint_to_gcs(checkpoint_url: str) -> str:\n",
        "    \"\"\"Uploads a compressed .tar.gz checkpoint at the given URL to Cloud Storage.\"\"\"\n",
        "    filename = os.path.basename(checkpoint_url)\n",
        "    checkpoint_name = filename.replace(\".tar.gz\", \"\")\n",
        "    print(\"Download checkpoint from\", checkpoint_url, \"and store to\", CHECKPOINT_BUCKET)\n",
        "    ! wget $checkpoint_url -O $filename\n",
        "    ! mkdir -p $checkpoint_name\n",
        "    ! tar -xvzf $filename -C $checkpoint_name\n",
        "\n",
        "    checkpoint_path = find_checkpoint_in_dir(checkpoint_name)\n",
        "    checkpoint_path = os.path.relpath(checkpoint_path, checkpoint_name)\n",
        "\n",
        "    ! gsutil cp -r $checkpoint_name $CHECKPOINT_BUCKET/\n",
        "    checkpoint_uri = os.path.join(CHECKPOINT_BUCKET, checkpoint_name, checkpoint_path)\n",
        "    print(\"Checkpoint uploaded to\", checkpoint_uri)\n",
        "    return checkpoint_uri\n",
        "\n",
        "\n",
        "def upload_config_to_gcs(url: str) -> str:\n",
        "    \"\"\"Uploads a config file at the given URL to Cloud Storage.\"\"\"\n",
        "    filename = os.path.basename(url)\n",
        "    destination = os.path.join(CONFIG_DIR, filename)\n",
        "    print(\"Copy\", url, \"to\", destination)\n",
        "    ! wget \"$url\" -O \"$filename\"\n",
        "    ! gsutil cp \"$filename\" \"$destination\"\n",
        "    return destination"
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {
        "id": "RB_xY9ipr7ZU"
      },
      "source": [
        "## Train new models\n",
        "This section shows how to train new models.\n",
        "1. Convert input data to training formats\n",
        "2. Create hyperparameter tuning jobs to train new models\n",
        "3. Find and export best models\n",
        "\n",
        "If you already trained models, please go to the section `Test Trained models`.\n",
        "\n",
        "Please select a model:\n",
        "* `model_id`: MoViNet model variant ID, one of `a0`, `a1`, `a2`, `a3`, `a4`, `a5`. The model with a larger number requires more resources to train, and is expected to have a higher accuracy and latency. Here, we use `a3` for demonstration purpose. **`a0`, `a1`, and `a2` are not recommended for now as we are currently investigating some inference issues with them.**\n",
        "* `model_mode`: MoViNet model type, either `base` or `stream`. The base model has a slightly higher accuracy, while the streaming model is optimized for streaming and faster CPU inference. See [official MoViNet docs](https://github.com/tensorflow/models/tree/master/official/projects/movinet) for more information.\n",
        "\n",
        "**Note**: The prediction container only supports base model (non-streaming) for now. If you train a streaming model, you need to download the model and refer to the [MoViNet official guide](https://github.com/tensorflow/models/blob/master/official/projects/movinet/movinet_streaming_model_training_and_inference.ipynb) for running predictions locally."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "3Ry1mw6AHLTy"
      },
      "outputs": [],
      "source": [
        "model_id = \"a3\"  # @param [\"a0\", \"a1\", \"a2\", \"a3\", \"a4\", \"a5\"]\n",
        "model_mode = \"base\"  # @param [\"base\", \"stream\"]\n",
        "is_stream = model_mode == \"stream\"\n",
        "model_name = f\"movinet_{model_id}_{model_mode}\"\n",
        "image_size = IMAGE_SIZES[model_id]\n",
        "\n",
        "if is_stream:\n",
        "    export_container_args = {\n",
        "        \"conv_type\": \"2plus1d\",\n",
        "        \"se_type\": \"2plus3d\",\n",
        "        \"activation\": \"hard_swish\",\n",
        "        \"gating_activation\": \"hard_sigmoid\",\n",
        "        \"use_positional_encoding\": model_id in {\"a3\", \"a4\", \"a5\"},\n",
        "    }\n",
        "else:\n",
        "    export_container_args = {\n",
        "        \"conv_type\": \"3d\",\n",
        "        \"se_type\": \"3d\",\n",
        "        \"activation\": \"swish\",\n",
        "        \"gating_activation\": \"sigmoid\",\n",
        "        \"use_positional_encoding\": False,\n",
        "    }"
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {
        "id": "zgPO1eR3CYjk"
      },
      "source": [
        "### Prepare input data for training\n",
        "\n",
        "Prepare data in the format as described [here](https://cloud.google.com/vertex-ai/docs/video-data/action-recognition/prepare-data), and then convert them to the training formats by running the cell below:\n",
        "\n",
        "* `input_file_path`: The input file path to the prepared data.\n",
        "* `input_file_type`: The input file type, such as `csv` or `jsonl`.\n",
        "* `output_fps`: The sampling rate of the video; Frames per second.\n",
        "* `num_frames`: Number of frame to sample around keyframe inputs.\n",
        "* `min_duration_sec`: Minimum duration in seconds for sampling video clips around keyframe inputs. This is for validation purpose - an error will be thrown if there is not enough context around a keyframe.\n",
        "* `pos_neg_ratio`: Sampling ratio between positive and negative segments. For example, a pos_neg_ratio of 0.5 samples 1 negative instance every 2 positive instances.\n",
        "* `split_ratio`: Three comma separated floats indicating the proportion of data to split into train/validation/test. They must add up to 1.\n",
        "* `num_shard`: Three comma separated integers indicating the shards for train/validation/test.\n",
        "\n",
        "**Note**: For JSONL input, please use `aiplatform.googleapis.com/ml_use` instead of `ml_use` as the JSON key for ML use in `dataItemResourceLabels`. This is to be consistent with other objectives."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "IndQ_m6ddUEM"
      },
      "outputs": [],
      "source": [
        "# This job will convert input data as training format, with given split ratios\n",
        "# and number of shards on train/test/validation.\n",
        "\n",
        "data_converter_job_name = get_job_name_with_datetime(\n",
        "    DATA_CONVERTER_JOB_PREFIX + \"_\" + OBJECTIVE\n",
        ")\n",
        "\n",
        "input_file_path = \"\"  # @param {type:\"string\"}\n",
        "input_file_type = \"csv\"  # @param [\"csv\", \"jsonl\"]\n",
        "output_fps = 10  # @param {type:\"integer\"}\n",
        "num_frames = 32  # @param {type:\"integer\"}\n",
        "min_duration_sec = 1.0  # @param {type:\"number\"}\n",
        "pos_neg_ratio = 1.0  # @param {type:\"number\"}\n",
        "split_ratio = \"0.8,0.1,0.1\"\n",
        "num_shard = \"10,10,10\"\n",
        "data_converter_output_dir = os.path.join(BUCKET_URI, data_converter_job_name)\n",
        "\n",
        "worker_pool_specs = [\n",
        "    {\n",
        "        \"machine_spec\": {\n",
        "            \"machine_type\": DATA_CONVERTER_MACHINE_TYPE,\n",
        "        },\n",
        "        \"replica_count\": 1,\n",
        "        \"container_spec\": {\n",
        "            \"image_uri\": DATA_CONVERTER_CONTAINER,\n",
        "            \"command\": [],\n",
        "            \"args\": [\n",
        "                \"--input_file_path=%s\" % input_file_path,\n",
        "                \"--input_file_type=%s\" % input_file_type,\n",
        "                \"--objective=%s\" % OBJECTIVE,\n",
        "                \"--num_shard=%s\" % num_shard,\n",
        "                \"--split_ratio=%s\" % split_ratio,\n",
        "                \"--output_dir=%s\" % data_converter_output_dir,\n",
        "                \"--output_fps=%d\" % output_fps,\n",
        "                \"--num_frames=%d\" % num_frames,\n",
        "                \"--min_duration_sec=%f\" % min_duration_sec,\n",
        "                \"--pos_neg_ratio=%f\" % pos_neg_ratio,\n",
        "                \"--output_shape=%d,%d\" % (image_size, image_size),\n",
        "            ],\n",
        "        },\n",
        "    }\n",
        "]\n",
        "\n",
        "data_converter_custom_job = aiplatform.CustomJob(\n",
        "    display_name=data_converter_job_name,\n",
        "    project=PROJECT_ID,\n",
        "    worker_pool_specs=worker_pool_specs,\n",
        "    staging_bucket=STAGING_BUCKET,\n",
        ")\n",
        "\n",
        "data_converter_custom_job.run()\n",
        "\n",
        "input_train_data_path = os.path.join(data_converter_output_dir, \"train.tfrecord*\")\n",
        "input_validation_data_path = os.path.join(data_converter_output_dir, \"val.tfrecord*\")\n",
        "label_map_path = os.path.join(data_converter_output_dir, \"label_map.yaml\")\n",
        "print(\"input_train_data_path for training: \", input_train_data_path)\n",
        "print(\"input_validation_data_path for training: \", input_validation_data_path)\n",
        "print(\"label_map_path for prediction: \", label_map_path)"
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {
        "id": "aaff6f5be7f6"
      },
      "source": [
        "### Create a Vertex AI custom job with hyperparameter tuning\n",
        "\n",
        "You use the Vertex AI SDK to create and run the [hyperparameter tuning job](https://cloud.google.com/vertex-ai/docs/training/hyperparameter-tuning-overview) with Vertex AI Model Garden training docker images.\n",
        "\n",
        "#### Define the following specifications\n",
        "\n",
        "* `worker_pool_specs`: A list of dictionaries specifying the machine type and docker image. This example defines a single node cluster with one `n1-highmem-32` machine with 8 `NVIDIA_TESLA_V100` GPUs.\n",
        "\n",
        "  **Note**: We recommend using 8 GPUs for MoViNet-A2 and larger. Since loading video data requires a lot of GPU memory, it is recommended to experiment with a small batch size first.\n",
        "* `parameter_spec`: Dictionary specifying the parameters to optimize. The dictionary key is the string assigned to the command line argument for each hyperparameter in your training application code, and the dictionary value is the parameter specification. The parameter specification includes the type, min/max values, and scale for the hyperparameter.\n",
        "* `metric_spec`: Dictionary specifying the metric to optimize. The dictionary key is the `hyperparameter_metric_tag` that you set in your training application code, and the value is the optimization goal."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "um_XKbmpTaHx"
      },
      "outputs": [],
      "source": [
        "from google.cloud.aiplatform import hyperparameter_tuning as hpt\n",
        "\n",
        "# Input train and validation datasets can be found from the section above\n",
        "# `Prepare input data for training`.\n",
        "# Or, set prepared datasets paths if already exist.\n",
        "# input_train_data_path = \"\"\n",
        "# input_validation_data_path = \"\"\n",
        "# label_map_path = \"\"\n",
        "\n",
        "train_job_name = get_job_name_with_datetime(f\"{TRAINING_JOB_PREFIX}_{model_name}\")\n",
        "model_dir = os.path.join(BUCKET_URI, train_job_name)\n",
        "label_map, num_classes = get_label_map(label_map_path)\n",
        "\n",
        "# Uploads pretained checkpoint to GCS bucket.\n",
        "init_checkpoint = f\"https://storage.googleapis.com/tf_model_garden/vision/movinet/{model_name}_with_backbone.tar.gz\"\n",
        "init_checkpoint = upload_checkpoint_to_gcs(init_checkpoint)\n",
        "\n",
        "# Uploads config file according to model_id and streaming options.\n",
        "config_file = f\"{model_id}_stream\" if is_stream else model_id\n",
        "config_file = f\"https://raw.githubusercontent.com/tensorflow/models/master/official/projects/movinet/configs/yaml/movinet_{config_file}_gpu.yaml\"\n",
        "config_file = upload_config_to_gcs(config_file)\n",
        "\n",
        "# The parameters here are mainly for demonstration purpose. Please update them\n",
        "# for better performance.\n",
        "trainer_args = {\n",
        "    \"experiment\": \"movinet_kinetics600\",\n",
        "    \"config_file\": config_file,\n",
        "    \"input_train_data_path\": input_train_data_path,\n",
        "    \"input_validation_data_path\": input_validation_data_path,\n",
        "    \"init_checkpoint\": init_checkpoint,\n",
        "    \"model_dir\": model_dir,\n",
        "    \"num_classes\": num_classes,\n",
        "    \"global_batch_size\": 16,\n",
        "    \"prefetch_buffer_size\": 16,\n",
        "    \"shuffle_buffer_size\": 32,\n",
        "    \"train_steps\": 2000,\n",
        "}\n",
        "\n",
        "worker_pool_specs = [\n",
        "    {\n",
        "        \"machine_spec\": {\n",
        "            \"machine_type\": TRAIN_MACHINE_TYPE,\n",
        "            \"accelerator_type\": TRAIN_ACCELERATOR_TYPE,\n",
        "            # Each training job uses TRAIN_NUM_GPU GPUs.\n",
        "            \"accelerator_count\": TRAIN_NUM_GPU,\n",
        "        },\n",
        "        \"replica_count\": 1,\n",
        "        \"container_spec\": {\n",
        "            \"image_uri\": TRAIN_CONTAINER_URI,\n",
        "            \"args\": [\n",
        "                \"--mode=train_and_eval\",\n",
        "                \"--params_override=runtime.num_gpus=%d\" % TRAIN_NUM_GPU,\n",
        "            ]\n",
        "            + [\"--{}={}\".format(k, v) for k, v in trainer_args.items()],\n",
        "        },\n",
        "    }\n",
        "]\n",
        "\n",
        "metric_spec = {\"model_performance\": \"maximize\"}\n",
        "\n",
        "# These learning rates might not be optimal for your selected model type; To\n",
        "# tune learning rates, try hpt.DoubleParameterSpec with more trials.\n",
        "LEARNING_RATES = [1e-3, 3e-3]\n",
        "MAX_TRIAL_COUNT = len(LEARNING_RATES)\n",
        "parameter_spec = {\n",
        "    \"learning_rate\": hpt.DiscreteParameterSpec(values=LEARNING_RATES, scale=\"linear\"),\n",
        "}\n",
        "\n",
        "print(worker_pool_specs, metric_spec, parameter_spec)"
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {
        "id": "HwcCjwlBTQIz"
      },
      "source": [
        "#### Run the hyperparameter tuning job\n",
        "* `max_trial_count`: Sets an upper bound on the number of trials the service will run. The recommended practice is to start with a smaller number of trials and get a sense of how impactful your chosen hyperparameters are before scaling up.\n",
        "\n",
        "* `parallel_trial_count`:  If you use parallel trials, the service provisions multiple training processing clusters. The worker pool spec that you specify when creating the job is used for each individual training cluster.  Increasing the number of parallel trials reduces the amount of time the hyperparameter tuning job takes to run; however, it can reduce the effectiveness of the job overall. This is because the default tuning strategy uses results of previous trials to inform the assignment of values in subsequent trials.\n",
        "\n",
        "* `search_algorithm`: The available search algorithms are grid, random, or default (None). The default option applies Bayesian optimization to search the space of possible hyperparameter values and is the recommended algorithm.\n",
        "\n",
        "Click on the generated link in the output to see your run in the Cloud Console."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "aec22792ee84"
      },
      "outputs": [],
      "source": [
        "train_custom_job = aiplatform.CustomJob(\n",
        "    display_name=train_job_name,\n",
        "    project=PROJECT_ID,\n",
        "    worker_pool_specs=worker_pool_specs,\n",
        "    staging_bucket=STAGING_BUCKET,\n",
        ")\n",
        "\n",
        "train_hpt_job = aiplatform.HyperparameterTuningJob(\n",
        "    display_name=train_job_name,\n",
        "    custom_job=train_custom_job,\n",
        "    metric_spec=metric_spec,\n",
        "    parameter_spec=parameter_spec,\n",
        "    max_trial_count=MAX_TRIAL_COUNT,\n",
        "    parallel_trial_count=1,\n",
        "    project=PROJECT_ID,\n",
        "    search_algorithm=None,\n",
        ")\n",
        "\n",
        "train_hpt_job.run()\n",
        "\n",
        "print(\"model_dir is:\", model_dir)"
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {
        "id": "vugUfJEC2HrK"
      },
      "source": [
        "### Export model in Tensorflow SavedModel format"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "09Rz1AYspK19"
      },
      "outputs": [],
      "source": [
        "# This job will export models from TF checkpoints to TF saved model format.\n",
        "# model_dir is from the section above.\n",
        "best_trial_dir, best_trial_evaluation_results = get_best_trial(\n",
        "    model_dir, MAX_TRIAL_COUNT, EVALUATION_METRIC\n",
        ")\n",
        "best_checkpoint_path = find_checkpoint_in_dir(f\"{best_trial_dir}/best_ckpt/\")\n",
        "print(\"best_trial_dir: \", best_trial_dir)\n",
        "print(\"best_trial_evaluation_results: \", best_trial_evaluation_results)\n",
        "print(\"best_checkpoint: \", best_checkpoint_path)\n",
        "\n",
        "container_args = {\n",
        "    \"export_path\": f\"{model_dir}/best_model\",\n",
        "    \"model_id\": model_id,\n",
        "    \"num_classes\": num_classes,\n",
        "    \"causal\": is_stream,\n",
        "    \"checkpoint_path\": best_checkpoint_path,\n",
        "    \"assert_checkpoint_objects_matched\": False,\n",
        "    **export_container_args,\n",
        "}\n",
        "\n",
        "worker_pool_specs = [\n",
        "    {\n",
        "        \"machine_spec\": {\n",
        "            \"machine_type\": EXPORT_MACHINE_TYPE,\n",
        "        },\n",
        "        \"replica_count\": 1,\n",
        "        \"container_spec\": {\n",
        "            \"image_uri\": EXPORT_CONTAINER_URI,\n",
        "            \"args\": [\"--{}={}\".format(k, v) for k, v in container_args.items()],\n",
        "        },\n",
        "    }\n",
        "]\n",
        "\n",
        "model_export_job_name = get_job_name_with_datetime(EXPORT_JOB_PREFIX + \"_\" + OBJECTIVE)\n",
        "model_export_custom_job = aiplatform.CustomJob(\n",
        "    display_name=model_export_job_name,\n",
        "    project=PROJECT_ID,\n",
        "    worker_pool_specs=worker_pool_specs,\n",
        "    staging_bucket=STAGING_BUCKET,\n",
        ")\n",
        "\n",
        "model_export_custom_job.run()\n",
        "\n",
        "print(\"best model is saved to: \", container_args[\"export_path\"])"
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {
        "id": "g0BGaofgsMsy"
      },
      "source": [
        "## Test trained models\n",
        "This section shows the way to test with trained models.\n",
        "1. Upload and deploy models to the [Vertex AI Model Registry](https://cloud.google.com/vertex-ai/docs/model-registry/introduction)\n",
        "2. Run batch predictions\n",
        "\n",
        "**Note:** The prediction container only works with the base model. If you trained a streaming model, download the model from the exported path and refer to the [MoViNet official guide](https://github.com/tensorflow/models/blob/master/official/projects/movinet/movinet_streaming_model_training_and_inference.ipynb) for running predictions locally."
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {
        "id": "gdlca3BOypXU"
      },
      "source": [
        "### Upload model to Vertex AI Model Registry\n",
        "\n",
        "The following cell uploads the trained model to Vertex AI Model Registry. Skip it if you want to run batch predictions on an already uploaded model instead.\n",
        "\n",
        "#### Configurable environment variables\n",
        "\n",
        "* `MODEL_PATH`: Cloud Storage URI to the MoViNet model.\n",
        "* `BATCH_SIZE`: Batch size for inference. Use a larger value to accelerate GPU prediction.\n",
        "* `NUM_FRAMES`: Number of frames for a single prediction with the model.\n",
        "* `FPS`: Video sampling frame per second.\n",
        "* `OVERLAP_FRAMES`: Allowed overlapping frames between consecutive prediction windows. Set a smaller value for faster inference but less accurate."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "NYuQowyZEtxK"
      },
      "outputs": [],
      "source": [
        "serving_env = {\n",
        "    \"MODEL_ID\": \"tfvision-movinet-var\",\n",
        "    \"MODEL_PATH\": container_args[\"export_path\"],\n",
        "    \"BATCH_SIZE\": 1,\n",
        "    \"NUM_FRAMES\": 32,\n",
        "    \"FPS\": output_fps,\n",
        "    \"OVERLAP_FRAMES\": 24,\n",
        "    \"OBJECTIVE\": OBJECTIVE,\n",
        "    \"IMAGE_WIDTH\": image_size,\n",
        "    \"IMAGE_HEIGHT\": image_size,\n",
        "    \"DEPLOY_SOURCE\": \"notebook\",\n",
        "}\n",
        "\n",
        "model = aiplatform.Model.upload(\n",
        "    display_name=model_name,\n",
        "    serving_container_image_uri=PREDICTION_CONTAINER_URI,\n",
        "    serving_container_ports=[PREDICTION_PORT],\n",
        "    serving_container_predict_route=\"/predict\",\n",
        "    serving_container_health_route=\"/ping\",\n",
        "    serving_container_environment_variables=serving_env,\n",
        ")\n",
        "\n",
        "model.wait()\n",
        "\n",
        "print(\"The uploaded model name is: \", model_name)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "c2b47e629a01"
      },
      "source": [
        "Alternatively, uncomment the following cell to use an already uploaded model. Replace the model name string with that of the existing model."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "aa503a565b8f"
      },
      "outputs": [],
      "source": [
        "# model = aiplatform.Model(\"projects/123456789/locations/us-central1/models/12345678901234567890\")"
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {
        "id": "9SZsKGeS3x6S"
      },
      "source": [
        "### Run batch predictions\n",
        "\n",
        "We will now run batch predictions with the trained MoViNet action recognition model with [Vertex AI Batch Prediction](https://cloud.google.com/vertex-ai/docs/predictions/get-batch-predictions).\n",
        "\n",
        "Please prepare an input JSONL file where each line follows [this format](https://cloud.google.com/vertex-ai/docs/video-data/action-recognition/get-predictions?hl=en#input_data_requirements) and store it in a Cloud Storage bucket. The service account should have read access to the buckets containing the trained model and the input data. See [Service accounts overview](https://cloud.google.com/iam/docs/service-account-overview) for more information.\n",
        "\n",
        "The [Vertex AI Batch Prediction](https://cloud.google.com/vertex-ai/docs/predictions/get-batch-predictions) has a default timeout of 10 minutes. Therefore, please make sure the input video clip is around 5 minutes at 5~10 FPS or you may experience a timeout error. To use this model at a larger scale beyond this notebook demontration, you can try one of the following:\n",
        "\n",
        "- Pull the serving docker image to a VM or a local machine and send prediction requests directly.\n",
        "- To process more data concurrently, write a custom [DataFlow](https://cloud.google.com/dataflow) pipeline to send prediction requests to the movinet serving container.\n",
        "- Divide videos into 5-minute clips and run batch prediction with a small batch size."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "vbIW9me1F2RY"
      },
      "outputs": [],
      "source": [
        "# Path to the prediction input JSONL file.\n",
        "test_jsonl_path = \"\"  # @param {type:\"string\"}\n",
        "# Full service account name with the suffix `gserviceaccount.com`.\n",
        "batch_predict_service_account = \"\"  # @param {type:\"string\"}\n",
        "\n",
        "predict_job_name = get_job_name_with_datetime(f\"{PREDICTION_JOB_PREFIX}_{model_name}\")\n",
        "predict_destination_prefix = os.path.join(STAGING_BUCKET, predict_job_name)\n",
        "\n",
        "batch_prediction_job = model.batch_predict(\n",
        "    job_display_name=predict_job_name,\n",
        "    gcs_source=test_jsonl_path,\n",
        "    gcs_destination_prefix=predict_destination_prefix,\n",
        "    machine_type=PREDICTION_MACHINE_TYPE,\n",
        "    accelerator_count=PREDICTION_ACCELERATOR_COUNT,\n",
        "    accelerator_type=PREDICTION_ACCELERATOR_TYPE,\n",
        "    max_replica_count=1,\n",
        "    service_account=batch_predict_service_account,\n",
        ")\n",
        "\n",
        "batch_prediction_job.wait()\n",
        "\n",
        "print(batch_prediction_job.display_name)\n",
        "print(batch_prediction_job.resource_name)\n",
        "print(batch_prediction_job.state)"
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {
        "id": "ik-XPjfx9OCE"
      },
      "source": [
        "You can then read the prediction response JSONL files in the output directory:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "tdkW9e5B9OU1"
      },
      "outputs": [],
      "source": [
        "# The label map file was generated from the section above (`Prepare input data for training`).\n",
        "for file in tf.io.gfile.glob(os.path.join(predict_destination_prefix, \"*/*\")):\n",
        "    with tf.io.gfile.GFile(file, \"r\") as f:\n",
        "        for line in f:\n",
        "            print_response_instance(line, label_map)"
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {
        "id": "kkH2nrpdp4sp"
      },
      "source": [
        "## Clean up"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Ax6vQVZhp9pR"
      },
      "outputs": [],
      "source": [
        "# Delete the trained model.\n",
        "model.delete()\n",
        "# Delete custom and hpt jobs.\n",
        "if data_converter_custom_job.list(filter=f'display_name=\"{data_converter_job_name}\"'):\n",
        "    data_converter_custom_job.delete()\n",
        "if train_hpt_job.list(filter=f'display_name=\"{train_job_name}\"'):\n",
        "    train_hpt_job.delete()\n",
        "if model_export_custom_job.list(filter=f'display_name=\"{model_export_job_name}\"'):\n",
        "    model_export_custom_job.delete()\n",
        "if batch_prediction_job.list(filter=f'display_name=\"{predict_job_name}\"'):\n",
        "    batch_prediction_job.delete()"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "name": "model_garden_movinet_action_recognition.ipynb",
      "toc_visible": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}