notebooks/community/model_garden/model_garden_pytorch_open_clip.ipynb

{
  "cells": [
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "7d9bbf86da5e"
      },
      "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": "99c1c3fc2ca5"
      },
      "source": [
        "# Vertex AI Model Garden - OpenCLIP\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_pytorch_open_clip.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/community/model_garden/model_garden_pytorch_open_clip.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_pytorch_open_clip.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": "3de7470326a2"
      },
      "source": [
        "## Overview\n",
        "\n",
        "This notebook demonstrates finetuning [OpenCLIP](https://github.com/mlfoundations/open_clip) with [CC3M](https://ai.google.com/research/ConceptualCaptions/download) dataset and deploying it on Vertex AI for online prediction.\n",
        "\n",
        "### Objective\n",
        "\n",
        "- Finetune the OpenCLIP model with [Vertex AI custom training](https://cloud.google.com/vertex-ai/docs/training/overview).\n",
        "- Upload the model to [Vertex AI Model Registry](https://cloud.google.com/vertex-ai/docs/model-registry/introduction).\n",
        "- Deploy the model to a [Vertex AI Endpoint resource](https://cloud.google.com/vertex-ai/docs/predictions/using-private-endpoints).\n",
        "- Run online predictions for zero-shot image classification.\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 pricing](https://cloud.google.com/vertex-ai/pricing) and [Cloud Storage pricing](https://cloud.google.com/storage/pricing), and use the [Pricing Calculator](https://cloud.google.com/products/calculator/) to generate a cost estimate based on your projected usage."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "264c07757582"
      },
      "source": [
        "## Setup environment\n",
        "\n",
        "**NOTE**: Jupyter runs lines prefixed with `!` as shell commands, and it interpolates Python variables prefixed with `$` into these commands."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "d73ffa0c0b83"
      },
      "source": [
        "### Colab only"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "2707b02ef5df"
      },
      "outputs": [],
      "source": [
        "if \"google.colab\" in str(get_ipython()):\n",
        "    ! pip3 install --upgrade google-cloud-aiplatform\n",
        "    from google.colab import auth as google_auth\n",
        "\n",
        "    google_auth.authenticate_user()\n",
        "\n",
        "    # Restart the notebook kernel after installs.\n",
        "    import IPython\n",
        "\n",
        "    app = IPython.Application.instance()\n",
        "    app.kernel.do_shutdown(True)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "bb7adab99e41"
      },
      "source": [
        "### Setup Google Cloud project\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",
        "\n",
        "1. [Create a Cloud Storage bucket](https://cloud.google.com/storage/docs/creating-buckets) for storing experiment outputs.\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`,  `Storage Object Admin`, and `GCS Storage Bucket Owner roles` roles for deploying fine tuned model to Vertex AI endpoint."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "6c460088b873"
      },
      "source": [
        "Fill the following variables for experiments environment:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "a12c23679315"
      },
      "outputs": [],
      "source": [
        "# Cloud project id.\n",
        "PROJECT_ID = \"\"  # @param {type:\"string\"}\n",
        "\n",
        "# The region you want to launch jobs in.\n",
        "REGION = \"\"  # @param {type:\"string\"}\n",
        "\n",
        "# The GCS bucket for storing experiments output. Fill it without the 'gs://' prefix.\n",
        "GCS_BUCKET = \"\"  # @param {type:\"string\"}\n",
        "\n",
        "# The service account for deploying fine tuned model.\n",
        "# The service account looks like:\n",
        "# '<account_name>@<project>.iam.gserviceaccount.com'\n",
        "# Follow step 5 above to create this account.\n",
        "SERVICE_ACCOUNT = \"\"  # @param {type:\"string\"}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "12c249e14a5d"
      },
      "source": [
        "### Download data to Google Cloud Storage"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "fb0259bfd059"
      },
      "outputs": [],
      "source": [
        "# Install the library for downloading training data.\n",
        "!pip install img2dataset==1.41.0"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "1d6b99305e7e"
      },
      "outputs": [],
      "source": [
        "import os\n",
        "\n",
        "!gcloud storage cp gs://gcc-data/Validation/GCC-1.1.0-Validation.tsv ./data.tsv  # Download list of URLs.\n",
        "!sed -i '1s/^/caption\\turl\\n/' data.tsv  # Add column name."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "f36f82f74dde"
      },
      "outputs": [],
      "source": [
        "# Download images from URLs. It takes around 2min.\n",
        "output_folder = \"data\"\n",
        "!img2dataset --url_list data.tsv --input_format \"tsv\"\\\n",
        "    --output_folder {output_folder}\\\n",
        "    --url_col \"url\" --caption_col \"caption\" --output_format webdataset\\\n",
        "    --processes_count {os.cpu_count()} --thread_count {os.cpu_count()*4}\\\n",
        "    --image_size 256"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "514ee0208d47"
      },
      "outputs": [],
      "source": [
        "GCS_data_dir = f\"{GCS_BUCKET}/CC3M-val-wds\""
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "b90bbd80476a"
      },
      "outputs": [],
      "source": [
        "# Upload data to GCS.\n",
        "!gcloud storage cp -r data gs://{GCS_data_dir}"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "7ffc8c1aed83"
      },
      "outputs": [],
      "source": [
        "# Get total number of samples, which is required by OpenCLIP training.\n",
        "import json\n",
        "\n",
        "n_samples = 0\n",
        "for filename in [f for f in os.listdir(output_folder) if f.endswith(\"_stats.json\")]:\n",
        "    with open(os.path.join(output_folder, filename)) as f:\n",
        "        n_samples += json.load(f)[\"successes\"]\n",
        "print(n_samples)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "e828eb320337"
      },
      "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": "12cd25839741"
      },
      "outputs": [],
      "source": [
        "from google.cloud import aiplatform\n",
        "\n",
        "aiplatform.init(project=PROJECT_ID, location=REGION, staging_bucket=GCS_BUCKET)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "2cc825514deb"
      },
      "source": [
        "### Define constants"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "b42bd4fa2b2d"
      },
      "outputs": [],
      "source": [
        "# The pre-built training docker image. It contains training scripts and models.\n",
        "TRAIN_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-open-clip-train\"\n",
        "\n",
        "# The pre-built serving docker image. It contains serving scripts and models.\n",
        "SERVE_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-open-clip-serve\""
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "0c250872074f"
      },
      "source": [
        "### Define common functions for job operations and test data preparations"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "354da31189dc"
      },
      "outputs": [],
      "source": [
        "import base64\n",
        "from datetime import datetime\n",
        "from io import BytesIO\n",
        "\n",
        "import matplotlib.pyplot as plt\n",
        "import requests\n",
        "from PIL import Image\n",
        "\n",
        "\n",
        "def create_job_name(prefix):\n",
        "    \"\"\"Create a job name string with a prefix.\"\"\"\n",
        "    user = os.environ.get(\"USER\")\n",
        "    now = datetime.now().strftime(\"%Y%m%d_%H%M%S\")\n",
        "    job_name = f\"{prefix}-{user}-{now}\"\n",
        "    return job_name\n",
        "\n",
        "\n",
        "def deploy_model(checkpoint, model, precision, task):\n",
        "    \"\"\"Deploy a model to Vertex AI endpoint.\"\"\"\n",
        "    model_name = \"openclip\"\n",
        "    endpoint = aiplatform.Endpoint.create(display_name=f\"{model_name}-{task}-endpoint\")\n",
        "    serving_env = {\n",
        "        \"CHECKPOINT\": checkpoint,\n",
        "        \"MODEL_ID\": model,\n",
        "        \"PRECISION\": precision,\n",
        "        \"TASK\": task,\n",
        "        \"DEPLOY_SOURCE\": \"notebook\",\n",
        "    }\n",
        "    model = aiplatform.Model.upload(\n",
        "        display_name=model_name,\n",
        "        serving_container_image_uri=SERVE_DOCKER_URI,\n",
        "        serving_container_ports=[7080],\n",
        "        serving_container_predict_route=\"/predictions/transformers_serving\",\n",
        "        serving_container_health_route=\"/ping\",\n",
        "        serving_container_environment_variables=serving_env,\n",
        "    )\n",
        "    model.deploy(\n",
        "        endpoint=endpoint,\n",
        "        machine_type=\"n1-standard-4\",\n",
        "        accelerator_type=\"NVIDIA_TESLA_V100\",\n",
        "        accelerator_count=1,\n",
        "        deploy_request_timeout=1800,\n",
        "        service_account=SERVICE_ACCOUNT,\n",
        "    )\n",
        "    return model, endpoint\n",
        "\n",
        "\n",
        "def download_image(url):\n",
        "    response = requests.get(url)\n",
        "    return Image.open(BytesIO(response.content)).convert(\"RGB\")\n",
        "\n",
        "\n",
        "def image_to_base64(image):\n",
        "    buffer = BytesIO()\n",
        "    image.save(buffer, format=\"PNG\")\n",
        "    image_str = base64.b64encode(buffer.getvalue()).decode(\"utf-8\")\n",
        "    return image_str\n",
        "\n",
        "\n",
        "def plot_images(images, rows=1, cols=None):\n",
        "    fig, axes = plt.subplots(rows, cols if cols is not None else len(images))\n",
        "    for ax, img in zip(axes, images):\n",
        "        ax.imshow(img)\n",
        "        ax.axis(\"off\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "e70e3519ff8b"
      },
      "source": [
        "## Fine tune the model\n",
        "\n",
        "This section fine tunes the OpenCLIP model with [CC3M](https://ai.google.com/research/ConceptualCaptions/download) dataset validation split for the image-text pre-training. It loads the pretrained checkpoint by OpenAI. You use a small model of ***RN50*** here; check [this list](https://github.com/mlfoundations/open_clip#pretrained-model-interface) for other options supported by OpenAI checkpoints, such as ***RN50, ViT-B-32, etc***.\n",
        "\n",
        "One `n1-standard-4` machine with 1 `NVIDIA_TESLA_V100` is required to run the fine-tuning job. The fine-tuning job takes about 3min to complete training for 2 epochs.\n",
        "\n",
        "The fine-tuned model will be saved after the job finishs and it can then be loaded for inference."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "55dabb1b02e3"
      },
      "outputs": [],
      "source": [
        "machine_type = \"n1-standard-4\"\n",
        "gpu_type = \"NVIDIA_TESLA_V100\"\n",
        "num_gpus = 1\n",
        "\n",
        "job_name = create_job_name(\"openclip\")\n",
        "\n",
        "model_name = \"RN50\"\n",
        "precision = \"amp\"\n",
        "\n",
        "job = aiplatform.CustomContainerTrainingJob(\n",
        "    display_name=job_name, container_uri=TRAIN_DOCKER_URI, command=[\"torchrun\"]\n",
        ")\n",
        "\n",
        "job.run(\n",
        "    args=[\n",
        "        f\"--nproc_per_node={num_gpus}\",\n",
        "        \"-m\",\n",
        "        \"training.main\",\n",
        "        f\"--name={job_name}\",\n",
        "        f\"--logs=/gcs/{GCS_BUCKET}\",\n",
        "        f\"--train-data=/gcs/{GCS_data_dir}/{{00000..00001}}.tar\",\n",
        "        f\"--train-num-samples={n_samples}\",\n",
        "        \"--dataset-type=webdataset\",\n",
        "        \"--batch-size=32\",\n",
        "        \"--precision=amp\",\n",
        "        \"--workers=8\",\n",
        "        \"--dataset-resampled\",\n",
        "        \"--save-frequency=2\",\n",
        "        \"--epochs=2\",\n",
        "        f\"--model={model_name}\",\n",
        "        \"--pretrained=openai\",\n",
        "        \"--save-most-recent\",\n",
        "    ],\n",
        "    boot_disk_size_gb=600,\n",
        "    replica_count=1,\n",
        "    machine_type=machine_type,\n",
        "    accelerator_type=gpu_type,\n",
        "    accelerator_count=num_gpus,\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "bf7f82732e61"
      },
      "source": [
        "## Upload and Deploy models\n",
        "\n",
        "This section uploads the fine-tuned model to Model Registry and deploys it on the Endpoint.\n",
        "\n",
        "One `n1-standard-4` machine with 1 `NVIDIA_TESLA_V100` is required to deploy OpenCLIP model.\n",
        "\n",
        "The model deployment step will take ~20 minutes to complete."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "6a2cf6e84b10"
      },
      "outputs": [],
      "source": [
        "# Prepare image samples\n",
        "img_diagram = download_image(\n",
        "    \"https://raw.githubusercontent.com/mlfoundations/open_clip/main/docs/CLIP.png\"\n",
        ")\n",
        "img_cat = download_image(\n",
        "    \"https://upload.wikimedia.org/wikipedia/commons/thumb/9/9d/Calicocats2222.jpg/220px-Calicocats2222.jpg\"\n",
        ")\n",
        "plot_images([img_diagram, img_cat])"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "5d63ffdd5f09"
      },
      "source": [
        "#### Zero-shot image classification"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "bf55e38815dc"
      },
      "outputs": [],
      "source": [
        "model, endpoint = deploy_model(\n",
        "    checkpoint=f\"gs://{GCS_BUCKET}/{job_name}/checkpoints/epoch_latest.pt\",\n",
        "    model=model_name,\n",
        "    precision=precision,\n",
        "    task=\"zero-shot-image-classification\",\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "80b3fd2ace09"
      },
      "source": [
        "NOTE: The model weights are downloaded after the deployment succeeds. An additional 5 minutes of waiting time is needed **after** the above model deployment step succeeds and before you run the next step below. Otherwise you might see a `ServiceUnavailable: 503 502:Bad Gateway` error when you send requests to the endpoint."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "002f438ecfea"
      },
      "outputs": [],
      "source": [
        "instances = [\n",
        "    {\n",
        "        \"text\": [\"a diagram\", \"a dog\", \"a cat\"],\n",
        "        \"image\": image_to_base64(img_diagram),\n",
        "    },\n",
        "    {\n",
        "        \"text\": [\"a diagram\", \"a dog\", \"two cats\", \"calico cat\"],\n",
        "        \"image\": image_to_base64(img_cat),\n",
        "    },\n",
        "]\n",
        "response = endpoint.predict(instances=instances).predictions\n",
        "response"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "af21a3cff1e0"
      },
      "source": [
        "Clean up resources:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "911406c1561e"
      },
      "outputs": [],
      "source": [
        "# Undeploy model and delete endpoint.\n",
        "endpoint.delete(force=True)\n",
        "\n",
        "# Delete models.\n",
        "model.delete()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "4fb87c6aeecf"
      },
      "source": [
        "#### Image/text feature embedding"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "bf55e38815dc"
      },
      "outputs": [],
      "source": [
        "model, endpoint = deploy_model(\n",
        "    checkpoint=f\"gs://{GCS_BUCKET}/{job_name}/checkpoints/epoch_latest.pt\",\n",
        "    model=model_name,\n",
        "    precision=precision,\n",
        "    task=\"feature-embedding\",\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "80b3fd2ace09"
      },
      "source": [
        "NOTE: The model weights are downloaded after the deployment succeeds. An additional 5 minutes of waiting time is needed **after** the above model deployment step succeeds and before you run the next step below. Otherwise you might see a `ServiceUnavailable: 503 502:Bad Gateway` error when you send requests to the endpoint."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "56d5c001575d"
      },
      "outputs": [],
      "source": [
        "import numpy as np\n",
        "\n",
        "instances = [\n",
        "    {\n",
        "        \"text\": [\"a diagram\", \"a dog\", \"a cat\"],\n",
        "        \"image\": image_to_base64(img_diagram),\n",
        "    },\n",
        "    {\n",
        "        \"image\": image_to_base64(img_cat),\n",
        "    },\n",
        "    {\n",
        "        \"text\": [\"a diagram\", \"a dog\", \"two cats\", \"calico cat\"],\n",
        "    },\n",
        "    {\n",
        "        \"text\": \"a single value\",\n",
        "    },\n",
        "]\n",
        "response = endpoint.predict(instances=instances).predictions\n",
        "for pred in response:\n",
        "    for k, v in pred.items():\n",
        "        print(k, np.array(v).shape)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "af21a3cff1e0"
      },
      "source": [
        "Clean up resources:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "911406c1561e"
      },
      "outputs": [],
      "source": [
        "# Undeploy model and delete endpoint.\n",
        "endpoint.delete(force=True)\n",
        "\n",
        "# Delete models.\n",
        "model.delete()"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "name": "model_garden_pytorch_open_clip.ipynb",
      "toc_visible": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}