diff --git a/notebooks/community/model_garden/model_garden_pytorch_sam.ipynb b/notebooks/community/model_garden/model_garden_pytorch_sam.ipynb index e8e922595..61ceda2bd 100644 --- a/notebooks/community/model_garden/model_garden_pytorch_sam.ipynb +++ b/notebooks/community/model_garden/model_garden_pytorch_sam.ipynb @@ -4,11 +4,12 @@ "cell_type": "code", "execution_count": null, "metadata": { + "cellView": "form", "id": "6ad30fe2-1fc1-47e3-8a9f-624170b5aae6" }, "outputs": [], "source": [ - "# Copyright 2023 Google LLC\n", + "# 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", @@ -31,26 +32,18 @@ "source": [ " # Vertex AI Model Garden - Segment Anything Model (SAM) Serving on Vertex AI\n", "\n", - "\n", - "
\n", - " \n", - " \"Colab Run in Colab\n", + "\n", + " \n", - " \n", - " \n", - "
\n", + " \n", + " \"Google
Run in Colab Enterprise\n", " \n", " 		\n", + " \n", " \n", - " \"GitHub\n", - " View on GitHub\n", + " \"GitHub
View on GitHub\n", " \n", " 		\n", - " \n", - " \"Vertex\n", - "Open in Vertex AI Workbench\n", - " \n", - " (a Python-3 CPU notebook is recommended)\n", - "
" + "
" ] }, { @@ -81,7 +74,7 @@ "* 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." + "Learn about [Vertex AI pricing](https://cloud.google.com/vertex-ai/pricing), [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." ] }, { @@ -90,223 +83,211 @@ "id": "78f72e0a-52e5-4de5-ac0f-2171b3493825" }, "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": "G__krby2Mqmh" - }, - "source": [ - "### Colab only" + "## Before you begin" ] }, { "cell_type": "code", "execution_count": null, "metadata": { - "id": "RS5_6QVFMyc-" + "cellView": "form", + "id": "m-Ql0m9edvcA" }, "outputs": [], "source": [ - "!pip3 install --upgrade google-cloud-aiplatform" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "id": "Li2aj1n9NDYl" - }, - "outputs": [], - "source": [ - "from google.colab import auth as google_auth\n", + "# @title Setup Google Cloud project\n", "\n", - "google_auth.authenticate_user()" - ] - }, - { - "cell_type": "markdown", - "metadata": { - "id": "f6c1bc20-3495-448a-b242-01930ba8153c" - }, - "source": [ - "### Setup Google Cloud project\n", + "# @markdown 1. [Make sure that billing is enabled for your project](https://cloud.google.com/billing/docs/how-to/modify-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", + "# @markdown 2. For finetuning, **[click here](https://console.cloud.google.com/iam-admin/quotas?location=us-central1&metric=aiplatform.googleapis.com%2Frestricted_image_training_nvidia_a100_80gb_gpus)** to check if your project already has the required 8 Nvidia A100 80 GB GPUs in the us-central1 region. If yes, then run this notebook in the us-central1 region. If you do not have 8 Nvidia A100 80 GPUs or have more GPU requirements than this, then schedule your job with Nvidia H100 GPUs via Dynamic Workload Scheduler using [these instructions](https://cloud.google.com/vertex-ai/docs/training/schedule-jobs-dws). For Dynamic Workload Scheduler, check the [us-central1](https://console.cloud.google.com/iam-admin/quotas?location=us-central1&metric=aiplatform.googleapis.com%2Fcustom_model_training_preemptible_nvidia_h100_gpus) or [europe-west4](https://console.cloud.google.com/iam-admin/quotas?location=europe-west4&metric=aiplatform.googleapis.com%2Fcustom_model_training_preemptible_nvidia_h100_gpus) quota for Nvidia H100 GPUs. If you do not have enough GPUs, then you can follow [these instructions](https://cloud.google.com/docs/quotas/view-manage#viewing_your_quota_console) to request quota.\n", "\n", - "1. [Make sure that billing is enabled for your project](https://cloud.google.com/billing/docs/how-to/modify-project).\n", + "# @markdown 3. For serving, **[click here](https://console.cloud.google.com/iam-admin/quotas?location=us-central1&metric=aiplatform.googleapis.com%2Fcustom_model_serving_nvidia_l4_gpus)** to check if your project already has the required 1 L4 GPU in the us-central1 region. If yes, then run this notebook in the us-central1 region. If you need more L4 GPUs for your project, then you can follow [these instructions](https://cloud.google.com/docs/quotas/view-manage#viewing_your_quota_console) to request more. Alternatively, if you want to run predictions with A100 80GB or H100 GPUs, we recommend using the regions listed below. **NOTE:** Make sure you have associated quota in selected regions. Click the links to see your current quota for each GPU type: [Nvidia A100 80GB](https://console.cloud.google.com/iam-admin/quotas?metric=aiplatform.googleapis.com%2Fcustom_model_serving_nvidia_a100_80gb_gpus), [Nvidia H100 80GB](https://console.cloud.google.com/iam-admin/quotas?metric=aiplatform.googleapis.com%2Fcustom_model_serving_nvidia_h100_gpus).\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", + "# @markdown > | Machine Type | Accelerator Type | Recommended Regions |\n", + "# @markdown | ----------- | ----------- | ----------- |\n", + "# @markdown | a2-ultragpu-1g | 1 NVIDIA_A100_80GB | us-central1, us-east4, europe-west4, asia-southeast1, us-east4 |\n", + "# @markdown | a3-highgpu-2g | 2 NVIDIA_H100_80GB | us-west1, asia-southeast1, europe-west4 |\n", + "# @markdown | a3-highgpu-4g | 4 NVIDIA_H100_80GB | us-west1, asia-southeast1, europe-west4 |\n", + "# @markdown | a3-highgpu-8g | 8 NVIDIA_H100_80GB | us-central1, us-east5, europe-west4, us-west1, asia-southeast1 |\n", "\n", - "1. [Create a Cloud Storage bucket](https://cloud.google.com/storage/docs/creating-buckets) for storing experiment outputs." - ] - }, - { - "cell_type": "markdown", - "metadata": { - "id": "5b1f2c08-c84e-4158-a976-e6b59d1055ca" - }, - "source": [ - "It's highly recommended to run this notebook on [Vertex AI workbench](https://cloud.google.com/vertex-ai-workbench)." - ] - }, - { - "cell_type": "markdown", - "metadata": { - "id": "4fc1fc14-2d77-4bf7-8f6d-c1afc10c848a" - }, - "source": [ - "If you are running this notebook locally, you will need to install the [Cloud SDK](https://cloud.google.com/sdk) and [gsutil](https://cloud.google.com/storage/docs/gsutil_install)." - ] - }, - { - "cell_type": "markdown", - "metadata": { - "id": "e02b9811-b730-4573-83fa-9d47f2ce0436" - }, - "source": [ - "####Fill following variables for experiments environment:\n" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "id": "7c2ce2fa-5a9b-40f6-b99d-6c1325775b36" - }, - "outputs": [], - "source": [ - "# Cloud project id.\n", - "PROJECT_ID = \"\" # @param {type:\"string\"}\n", + "# @markdown 4. **[Optional]** [Create a Cloud Storage bucket](https://cloud.google.com/storage/docs/creating-buckets) for storing experiment outputs. Set the BUCKET_URI for the experiment environment. The specified Cloud Storage bucket (`BUCKET_URI`) should be located in the same region as where the notebook was launched. Note that a multi-region bucket (eg. \"us\") is not considered a match for a single region covered by the multi-region range (eg. \"us-central1\"). If not set, a unique GCS bucket will be created instead.\n", "\n", - "# The region you want to launch jobs in.\n", - "REGION = \"us-central1\" # @param {type:\"string\"}\n", + "BUCKET_URI = \"gs://\" # @param {type:\"string\"}\n", "\n", - "# The Cloud Storage bucket for storing experiments output. Fill it without the 'gs://' prefix.\n", - "GCS_BUCKET = \"\" # @param {type:\"string\"}" - ] - }, - { - "cell_type": "markdown", - "metadata": { - "id": "fwgeVkFXOS64" - }, - "source": [ - "*Initialize* Vertex AI API:" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "id": "MxqnJMahPK8r" - }, - "outputs": [], - "source": [ + "# @markdown 5. **[Optional]** Set region. If not set, the region will be set automatically according to Colab Enterprise environment.\n", + "\n", + "REGION = \"\" # @param {type:\"string\"}\n", + "\n", + "! git clone https://github.com/GoogleCloudPlatform/vertex-ai-samples.git\n", + "\n", + "import datetime\n", + "import importlib\n", + "import os\n", + "import uuid\n", + "\n", + "import matplotlib.pyplot as plt\n", + "import numpy as np\n", + "import pycocotools.mask as mask_util\n", "from google.cloud import aiplatform\n", "\n", - "aiplatform.init(project=PROJECT_ID, location=REGION, staging_bucket=GCS_BUCKET)" - ] - }, - { - "cell_type": "markdown", - "metadata": { - "id": "tSjMVyG8PYS-" - }, - "source": [ - "### Define constants" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "id": "H05BzPO5Pcnh" - }, - "outputs": [], - "source": [ - "# The pre-built serving docker image.\n", - "# The model artifacts are embedded within the container, except for model weights which will be downloaded during deployment.\n", - "SERVE_DOCKER_URI = (\n", - " \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/sam-serve\"\n", + "common_util = importlib.import_module(\n", + " \"vertex-ai-samples.community-content.vertex_model_garden.model_oss.notebook_util.common_util\"\n", ")\n", "\n", - "# The serving port.\n", - "SERVE_PORT = 7080\n", + "models, endpoints = {}, {}\n", + "\n", + "\n", + "# Get the default cloud project id.\n", + "PROJECT_ID = os.environ[\"GOOGLE_CLOUD_PROJECT\"]\n", + "\n", + "# Get the default region for launching jobs.\n", + "if not REGION:\n", + " REGION = os.environ[\"GOOGLE_CLOUD_REGION\"]\n", + "\n", + "# Enable the Vertex AI API and Compute Engine API, if not already.\n", + "print(\"Enabling Vertex AI API and Compute Engine API.\")\n", + "! gcloud services enable aiplatform.googleapis.com compute.googleapis.com\n", + "\n", + "# Cloud Storage bucket for storing the experiment artifacts.\n", + "# A unique GCS bucket will be created for the purpose of this notebook. If you\n", + "# prefer using your own GCS bucket, change the value yourself below.\n", + "now = datetime.datetime.now().strftime(\"%Y%m%d%H%M%S\")\n", + "BUCKET_NAME = \"/\".join(BUCKET_URI.split(\"/\")[:3])\n", + "\n", + "if BUCKET_URI is None or BUCKET_URI.strip() == \"\" or BUCKET_URI == \"gs://\":\n", + " BUCKET_URI = f\"gs://{PROJECT_ID}-tmp-{now}-{str(uuid.uuid4())[:4]}\"\n", + " BUCKET_NAME = \"/\".join(BUCKET_URI.split(\"/\")[:3])\n", + " ! gsutil mb -l {REGION} {BUCKET_URI}\n", + "else:\n", + " assert BUCKET_URI.startswith(\"gs://\"), \"BUCKET_URI must start with `gs://`.\"\n", + " shell_output = ! gsutil ls -Lb {BUCKET_NAME} | grep \"Location constraint:\" | sed \"s/Location constraint://\"\n", + " bucket_region = shell_output[0].strip().lower()\n", + " if bucket_region != REGION:\n", + " raise ValueError(\n", + " \"Bucket region %s is different from notebook region %s\"\n", + " % (bucket_region, REGION)\n", + " )\n", + "print(f\"Using this GCS Bucket: {BUCKET_URI}\")\n", "\n", - "# The serving route.\n", - "SERVE_ROUTE = \"/predictions/sam_serving\"" - ] - }, - { - "cell_type": "markdown", - "metadata": { - "id": "82058975-23b5-4b97-9e14-dd9a29c578ed" - }, - "source": [ - "### Define common utility functions" + "STAGING_BUCKET = os.path.join(BUCKET_URI, \"temporal\")\n", + "MODEL_BUCKET = os.path.join(BUCKET_URI, \"sam\")\n", + "\n", + "\n", + "# Initialize Vertex AI API.\n", + "print(\"Initializing Vertex AI API.\")\n", + "aiplatform.init(project=PROJECT_ID, location=REGION, staging_bucket=STAGING_BUCKET)\n", + "\n", + "# Gets the default SERVICE_ACCOUNT.\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", + "print(\"Using this default Service Account:\", SERVICE_ACCOUNT)\n", + "\n", + "\n", + "# Provision permissions to the SERVICE_ACCOUNT with the GCS bucket\n", + "! gsutil iam ch serviceAccount:{SERVICE_ACCOUNT}:roles/storage.admin $BUCKET_NAME\n", + "\n", + "! gcloud config set project $PROJECT_ID\n", + "! gcloud projects add-iam-policy-binding --no-user-output-enabled {PROJECT_ID} --member=serviceAccount:{SERVICE_ACCOUNT} --role=\"roles/storage.admin\"\n", + "! gcloud projects add-iam-policy-binding --no-user-output-enabled {PROJECT_ID} --member=serviceAccount:{SERVICE_ACCOUNT} --role=\"roles/aiplatform.user\"" ] }, { "cell_type": "code", "execution_count": null, "metadata": { - "id": "a78f988d-e5e2-4a57-ba0f-569e970514c0" + "cellView": "form", + "id": "ot2HhTqxRYri" }, "outputs": [], "source": [ - "import base64\n", - "from io import BytesIO\n", + "# @title Deploy\n", "\n", - "import matplotlib.patches as patches\n", - "import matplotlib.pyplot as plt\n", - "import numpy as np\n", - "import pycocotools.mask as mask_util\n", - "import requests\n", - "from PIL import Image\n", + "# @markdown This section deploys a pre-trained `sam-vit-large` model on Model Registry by using 1 L4 Machine.\n", "\n", + "# @markdown The model deploy step will take around 20 minutes to complete.\n", "\n", - "def download_image(url):\n", - " response = requests.get(url)\n", - " return Image.open(BytesIO(response.content)).convert(\"RGB\")\n", + "# The pre-built serving docker image.\n", + "# The model artifacts are embedded within the container, except for model weights which will be downloaded during deployment.\n", + "SERVE_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/sam-serve:public-image-20240121\"\n", "\n", + "# @markdown Set the accelerator type.\n", + "serve_accelerator_type = \"NVIDIA_L4\" # @param[\"NVIDIA_TESLA_V100\", \"NVIDIA_L4\"]\n", "\n", - "def image_to_base64(image):\n", - " buffer = BytesIO()\n", - " image.save(buffer, format=\"JPEG\")\n", - " image_str = base64.b64encode(buffer.getvalue()).decode(\"utf-8\")\n", - " return image_str\n", + "if serve_accelerator_type == \"NVIDIA_TESLA_V100\":\n", + " serve_machine_type = \"n1-standard-8\"\n", + " serve_accelerator_count = 1\n", + "elif serve_accelerator_type == \"NVIDIA_L4\":\n", + " serve_machine_type = \"g2-standard-12\"\n", + " serve_accelerator_count = 1\n", + "else:\n", + " print(f\"Unsupported accelerator type: {accelerator_type}\")\n", "\n", + "MODEL_ID = \"facebook/sam-vit-large\"\n", + "task = \"mask-generation\"\n", "\n", - "def base64_to_image(image_str):\n", - " image = Image.open(BytesIO(base64.b64decode(image_str)))\n", - " return image\n", "\n", + "def deploy_model(\n", + " task, display_name, model_id, machine_type, accelerator_type, accelerator_count\n", + "):\n", + " endpoint = aiplatform.Endpoint.create(\n", + " display_name=common_util.get_job_name_with_datetime(prefix=task)\n", + " )\n", + " serving_env = {\n", + " \"MODEL_ID\": model_id,\n", + " \"TASK\": task,\n", + " \"DEPLOY_SOURCE\": \"notebook\",\n", + " }\n", + " model = aiplatform.Model.upload(\n", + " display_name=task,\n", + " serving_container_image_uri=SERVE_DOCKER_URI,\n", + " serving_container_ports=[7080],\n", + " serving_container_predict_route=\"/predictions/sam_serving\",\n", + " serving_container_health_route=\"/ping\",\n", + " serving_container_environment_variables=serving_env,\n", + " )\n", + " model.deploy(\n", + " endpoint=endpoint,\n", + " machine_type=serve_machine_type,\n", + " accelerator_type=serve_accelerator_type,\n", + " accelerator_count=1,\n", + " deploy_request_timeout=1800,\n", + " system_labels={\"NOTEBOOK_NAME\": \"model_garden_pytorch_sam.ipynb\"},\n", + " )\n", + " return model, endpoint\n", "\n", - "def image_grid(imgs, rows=2, cols=2):\n", - " w, h = imgs[0].size\n", - " grid = Image.new(\"RGB\", size=(cols * w, rows * h))\n", - " for i, img in enumerate(imgs):\n", - " grid.paste(img, box=(i % cols * w, i // cols * h))\n", - " return grid\n", "\n", + "common_util.check_quota(\n", + " project_id=PROJECT_ID,\n", + " region=REGION,\n", + " accelerator_type=serve_accelerator_type,\n", + " accelerator_count=serve_accelerator_count,\n", + " is_for_training=False,\n", + ")\n", "\n", - "def draw_image_with_boxes(image, boxes):\n", - " fig, ax = plt.subplots()\n", - " ax.imshow(image)\n", - " for box in boxes:\n", - " x, y = box[\"xmin\"], box[\"ymin\"]\n", - " width, height = box[\"xmax\"] - x, box[\"ymax\"] - y\n", - " rect = patches.Rectangle(\n", - " (x, y), width, height, linewidth=2, edgecolor=\"yellow\", facecolor=\"none\"\n", - " )\n", - " ax.add_patch(rect)\n", - " plt.axis(\"off\")\n", - " plt.show()\n", + "models[\"sam_model\"], endpoints[\"sam_endpoint\"] = deploy_model(\n", + " task=task,\n", + " display_name=common_util.get_job_name_with_datetime(prefix=task),\n", + " model_id=\"facebook/sam-vit-large\",\n", + " machine_type=serve_machine_type,\n", + " accelerator_type=serve_accelerator_type,\n", + " accelerator_count=1,\n", + ")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "cellView": "form", + "id": "d6e51c57-b5e2-4ae7-888a-5391cceee5fb" + }, + "outputs": [], + "source": [ + "# @title Predict\n", + "\n", + "input_image1 = \"http://images.cocodataset.org/val2017/000000039769.jpg\" # @param {type:\"string\"}\n", + "input_image2 = \"http://images.cocodataset.org/val2017/000000000285.jpg\" # @param {type:\"string\"}\n", "\n", "\n", "def decode_rle_masks(pred_masks_rle):\n", @@ -344,74 +325,17 @@ " plt.show()\n", "\n", "\n", - "def deploy_model(task, model_id):\n", - " endpoint = aiplatform.Endpoint.create(display_name=f\"{task}-endpoint\")\n", - " serving_env = {\n", - " \"MODEL_ID\": model_id,\n", - " \"TASK\": task,\n", - " \"DEPLOY_SOURCE\": \"notebook\",\n", - " }\n", - " model = aiplatform.Model.upload(\n", - " display_name=task,\n", - " serving_container_image_uri=SERVE_DOCKER_URI,\n", - " serving_container_ports=[SERVE_PORT],\n", - " serving_container_predict_route=SERVE_ROUTE,\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-8\",\n", - " accelerator_type=\"NVIDIA_TESLA_V100\",\n", - " accelerator_count=1,\n", - " deploy_request_timeout=1800,\n", - " system_labels={\n", - " \"NOTEBOOK_NAME\": \"model_garden_pytorch_sam.ipynb\"\n", - " },\n", - " )\n", - " return endpoint, model" - ] - }, - { - "cell_type": "markdown", - "metadata": { - "id": "242fc8cf-5a0e-483b-8c9d-77a474cedc4b" - }, - "source": [ - "## Mask Generation with Vertex AI Endpoint" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "id": "ot2HhTqxRYri" - }, - "outputs": [], - "source": [ - "endpoint, model = deploy_model(\n", - " task=\"mask-generation\", model_id=\"facebook/sam-vit-large\"\n", - ")" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "id": "d6e51c57-b5e2-4ae7-888a-5391cceee5fb" - }, - "outputs": [], - "source": [ - "image1 = download_image(\"http://images.cocodataset.org/val2017/000000039769.jpg\")\n", - "image2 = download_image(\"http://images.cocodataset.org/val2017/000000000285.jpg\")\n", - "grid = image_grid([image1, image2], 1, 2)\n", + "image1 = common_util.download_image(input_image1)\n", + "image2 = common_util.download_image(input_image2)\n", + "grid = common_util.image_grid([image1, image2], 1, 2)\n", "display(grid)\n", "\n", "instances = [\n", - " {\"image\": image_to_base64(image1)},\n", - " {\"image\": image_to_base64(image2)},\n", + " {\"image\": common_util.image_to_base64(image1)},\n", + " {\"image\": common_util.image_to_base64(image2)},\n", "]\n", - "preds = endpoint.predict(instances=instances).predictions\n", + "\n", + "preds = endpoints[\"sam_endpoint\"].predict(instances=instances).predictions\n", "show_predictions(preds)" ] }, @@ -421,22 +345,32 @@ "id": "5UdquXMWR1E4" }, "source": [ - "## Clean up" + "## Clean up resources" ] }, { "cell_type": "code", "execution_count": null, "metadata": { + "cellView": "form", "id": "TGTWIJK8R136" }, "outputs": [], "source": [ + "# @markdown Delete the experiment models and endpoints to recycle the resources\n", + "# @markdown and avoid unnecessary continuous charges that may incur.\n", + "\n", "# Undeploy model and delete endpoint.\n", - "endpoint.delete(force=True)\n", + "for endpoint in endpoints.values():\n", + " endpoint.delete(force=True)\n", "\n", "# Delete models.\n", - "model.delete()" + "for model in models.values():\n", + " model.delete()\n", + "\n", + "delete_bucket = False # @param {type:\"boolean\"}\n", + "if delete_bucket:\n", + " ! gsutil -m rm -r $BUCKET_NAME" ] } ],