mirror of
https://github.com/GoogleCloudPlatform/vertex-ai-samples.git
synced 2026-09-26 14:42:04 +00:00
29 KiB
29 KiB
In [ ]:
# Copyright 2025 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.In [ ]:
# @title Setup Google Cloud project
# @markdown 1. [Make sure that billing is enabled for your project](https://cloud.google.com/billing/docs/how-to/modify-project).
# @markdown 2. **[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.
BUCKET_URI = "gs://" # @param {type:"string"}
# @markdown 3. **[Optional]** Set region. If not set, the region will be set automatically according to Colab Enterprise environment.
REGION = "" # @param {type:"string"}
# @markdown 4. 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). You can request for quota following the instructions at ["Request a higher quota"](https://cloud.google.com/docs/quota/view-manage#requesting_higher_quota).
# @markdown | Machine Type | Accelerator Type | Recommended Regions |
# @markdown | ----------- | ----------- | ----------- |
# @markdown | a2-ultragpu-1g | 1 NVIDIA_A100_80GB | us-central1, us-east4, europe-west4, asia-southeast1, us-east4 |
# @markdown | a3-highgpu-2g | 2 NVIDIA_H100_80GB | us-west1, asia-southeast1, europe-west4 |
# @markdown | a3-highgpu-4g | 4 NVIDIA_H100_80GB | us-west1, asia-southeast1, europe-west4 |
# @markdown | a3-highgpu-8g | 8 NVIDIA_H100_80GB | us-central1, europe-west4, us-west1, asia-southeast1 |
# Configs for all notebooks.
! pip3 install --quiet keras-cv==0.9.0
! pip3 install --quiet keras-core==0.1.0
! git clone https://github.com/GoogleCloudPlatform/vertex-ai-samples.git
import base64
import datetime
import importlib
import io
import os
import tempfile
import uuid
from typing import Dict, List, Union
import keras_cv
import numpy as np
import tensorflow as tf
import yaml
from google.cloud import aiplatform
from google.protobuf import json_format
from google.protobuf.struct_pb2 import Value
from keras_cv import visualization
from PIL import Image
common_util = importlib.import_module(
"vertex-ai-samples.notebooks.community.model_garden.docker_source_codes.notebook_util.common_util"
)
models, endpoints = {}, {}
# Get the default cloud project id.
PROJECT_ID = os.environ["GOOGLE_CLOUD_PROJECT"]
# Get the default region for launching jobs.
if not REGION:
if not os.environ.get("GOOGLE_CLOUD_REGION"):
raise ValueError(
"REGION must be set. See"
" https://cloud.google.com/vertex-ai/docs/general/locations for"
" available cloud locations."
)
REGION = os.environ["GOOGLE_CLOUD_REGION"]
# Enable the Vertex AI API and Compute Engine API, if not already.
print("Enabling Vertex AI API and Compute Engine API.")
! gcloud services enable aiplatform.googleapis.com compute.googleapis.com
# Cloud Storage bucket for storing the experiment artifacts.
# A unique GCS bucket will be created for the purpose of this notebook. If you
# prefer using your own GCS bucket, change the value yourself below.
now = datetime.datetime.now().strftime("%Y%m%d%H%M%S")
BUCKET_NAME = "/".join(BUCKET_URI.split("/")[:3])
if BUCKET_URI is None or BUCKET_URI.strip() == "" or BUCKET_URI == "gs://":
BUCKET_URI = f"gs://{PROJECT_ID}-tmp-{now}-{str(uuid.uuid4())[:4]}"
BUCKET_NAME = "/".join(BUCKET_URI.split("/")[:3])
! gcloud storage buckets create --location={REGION} {BUCKET_URI}
else:
assert BUCKET_URI.startswith("gs://"), "BUCKET_URI must start with `gs://`."
shell_output = ! gcloud storage ls --full --buckets {BUCKET_NAME} | grep "Location Constraint:" | sed "s/Location Constraint://"
bucket_region = shell_output[0].strip().lower()
if bucket_region != REGION:
raise ValueError(
"Bucket region %s is different from notebook region %s"
% (bucket_region, REGION)
)
print(f"Using this GCS Bucket: {BUCKET_URI}")
STAGING_BUCKET = os.path.join(BUCKET_URI, "temporal")
MODEL_BUCKET = os.path.join(BUCKET_URI, "keras_yolov8")
# Initialize Vertex AI API.
print("Initializing Vertex AI API.")
aiplatform.init(project=PROJECT_ID, location=REGION, staging_bucket=STAGING_BUCKET)
# Gets the default SERVICE_ACCOUNT.
shell_output = ! gcloud projects describe $PROJECT_ID
project_number = shell_output[-1].split(":")[1].strip().replace("'", "")
SERVICE_ACCOUNT = f"{project_number}-compute@developer.gserviceaccount.com"
print("Using this default Service Account:", SERVICE_ACCOUNT)
# Provision permissions to the SERVICE_ACCOUNT with the GCS bucket
! gcloud storage buckets add-iam-policy-binding $BUCKET_NAME --member=serviceAccount:{SERVICE_ACCOUNT} --role=roles/storage.admin
! gcloud config set project $PROJECT_ID
! gcloud projects add-iam-policy-binding --no-user-output-enabled {PROJECT_ID} --member=serviceAccount:{SERVICE_ACCOUNT} --role="roles/storage.admin"
! gcloud projects add-iam-policy-binding --no-user-output-enabled {PROJECT_ID} --member=serviceAccount:{SERVICE_ACCOUNT} --role="roles/aiplatform.user"
TRAIN_CONTAINER_URI = (
"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/keras-yolov8-train"
)
SERVING_CONTAINER_URI = (
"us-docker.pkg.dev/vertex-ai-restricted/prediction/tf_opt-gpu.2-12:latest"
)
SERVING_CONTAINER_ARGS = ["--allow_precompilation", "--allow_compression"]
RESOLUTION = 512
def load_img(path):
"""Reads image from path and return PIL.Image instance."""
img = tf.io.read_file(path)
img = tf.image.decode_jpeg(img, channels=3)
return Image.fromarray(np.uint8(img)).convert("RGB")
def decode_image(image_str_tensor: tf.string) -> tf.float32:
"""Converts and resizes image bytes to image tensor."""
image = tf.io.decode_image(image_str_tensor, 3, expand_animations=False)
image = tf.image.resize(image, (RESOLUTION, RESOLUTION))
return image
def get_label_map(label_map_yaml_filepath):
"""Returns class id to label mapping given a filepath to the label map."""
temp_dir = tempfile.TemporaryDirectory()
label_map_yaml_filename = os.path.basename(label_map_yaml_filepath)
local_metrics_path = os.path.join(temp_dir.name, label_map_yaml_filename)
! gcloud storage cp $label_map_yaml_filepath $local_metrics_path
with open(local_metrics_path) as input_file:
label_map = yaml.safe_load(input_file.read())["label_map"]
temp_dir.cleanup()
return label_map
def get_prediction_instances(test_filepath, new_width=-1):
"""Generate instance from image path to pass to Vertex AI Endpoint for prediction."""
if new_width <= 0:
with tf.io.gfile.GFile(test_filepath, "rb") as input_file:
encoded_string = base64.b64encode(input_file.read()).decode("utf-8")
else:
img = load_img(test_filepath)
width, height = img.size
print("original input image size: ", width, " , ", height)
new_height = int(height * new_width / width)
new_img = img.resize((new_width, new_height))
print("resized input image size: ", new_width, " , ", new_height)
buffered = io.BytesIO()
new_img.save(buffered, format="JPEG")
encoded_string = base64.b64encode(buffered.getvalue()).decode("utf-8")
instances = [
{
"encoded_image": {"b64": encoded_string},
}
]
return instances
def predict_custom_trained_model(
project: str,
endpoint_id: str,
instances: Union[Dict, List[Dict]],
location: str = "us-central1",
):
# The AI Platform services require regional API endpoints.
client_options = {"api_endpoint": f"{location}-aiplatform.googleapis.com"}
# Initialize client that will be used to create and send requests.
# This client only needs to be created once, and can be reused for multiple requests.
client = aiplatform.gapic.PredictionServiceClient(client_options=client_options)
parameters_dict = {}
parameters = json_format.ParseDict(parameters_dict, Value())
endpoint = client.endpoint_path(
project=project, location=location, endpoint=endpoint_id
)
response = client.predict(
endpoint=endpoint, instances=instances, parameters=parameters
)
return response.predictions, response.deployed_model_idIn [ ]:
# @title Run local inferences with pretrained model
# @markdown This section shows how to run inferences locally with YOLOv8-M pretrained on PascalVOC 2012 object detection task, which consists of 20 classes.
test_filepath = "" # @param {type:"string"}
img_bytes = tf.io.read_file(test_filepath)
image = tf.expand_dims(decode_image(img_bytes), axis=0)
# Load model pretrained on PascalVOC 2012.
model_name = "yolo_v8_m_pascalvoc"
model = keras_cv.models.YOLOV8Detector.from_preset(
model_name,
bounding_box_format="xywh",
)
decoded = model.predict(image)
# Classes in PascalVOC 2012 dataset.
class_ids = [
"Aeroplane",
"Bicycle",
"Bird",
"Boat",
"Bottle",
"Bus",
"Car",
"Cat",
"Chair",
"Cow",
"Dining Table",
"Dog",
"Horse",
"Motorbike",
"Person",
"Potted Plant",
"Sheep",
"Sofa",
"Train",
"Tvmonitor",
"Total",
]
class_mapping = dict(zip(range(len(class_ids)), class_ids))
# Visualize the results.
visualization.plot_bounding_box_gallery(
image,
value_range=(0, 255),
rows=1,
cols=1,
y_pred=decoded,
scale=5,
font_scale=0.7,
bounding_box_format="xywh",
class_mapping=class_mapping,
)In [ ]:
# @title Finetune
# @markdown This section shows how to finetune the Keras YOLOv8 model and deploy to Vertex AI Endpoint resource.
# @markdown `input_csv_path` : The input dataset in CSV format. For further details, kindly check [AutoML Image Object Detection](https://cloud.google.com/vertex-ai/docs/image-data/object-detection/prepare-data).
input_csv_path = "gs://cloud-samples-data/vision/salads.csv" # @param {type:"string"}
# Hyperparameters
# @markdown `epochs`: Number of training epochs.
epochs = 10 # @param{type:"integer"}
# @markdown `learning_rate`: The learning rate of this training job.
learning_rate = 0.0005 # @param{type:"number"}
# @markdown `fpn_depth`: The depth of the CSP blocks in the Feature Pyramid Network. This is usually 1, 2, or 3, depending on the size of your YOLOV8Detector model. We recommend using 3 for 'yolo_v8_l_backbone' and 'yolo_v8_xl_backbone'.Defaults to 2.
fpn_depth = 3 # @param{type:"integer"}
# @markdown `confidence_threshold`: Only probabilities greater than this threshold will contribute to the final result
confidence_threshold = 0.02 # @param{type:"number"}
# @markdown `iou_threshold`: Intersection over Union (IoU) is a measure that shows how well the prediction bounding box aligns with the ground truth box.
iou_threshold = 0.3 # @param{type:"number"}
# @markdown `backbone`: The pretrained backbone. [Click here](https://keras.io/api/keras_cv/models/backbones/yolo_v8/) for the full list of available backbones.
backbone = "yolo_v8_xl_backbone_coco" # @param["yolo_v8_xs_backbone", "yolo_v8_s_backbone", "yolo_v8_m_backbone", "yolo_v8_l_backbone", "yolo_v8_xl_backbone", "yolo_v8_xs_backbone_coco", "yolo_v8_s_backbone_coco", "yolo_v8_m_backbone_coco", "yolo_v8_l_backbone_coco", "yolo_v8_xl_backbone_coco"]
MACHINE_TYPE = "n1-highmem-16"
ACCELERATOR_TYPE = "NVIDIA_TESLA_V100"
ACCELERATOR_COUNT = 2
train_job_name = common_util.get_job_name_with_datetime("train_yolov8")
model_dir = os.path.join(STAGING_BUCKET, train_job_name)
worker_pool_specs = [
{
"machine_spec": {
"machine_type": MACHINE_TYPE,
"accelerator_type": ACCELERATOR_TYPE,
"accelerator_count": ACCELERATOR_COUNT,
},
"replica_count": 1,
"disk_spec": {
"boot_disk_type": "pd-ssd",
"boot_disk_size_gb": 500,
},
"container_spec": {
"image_uri": TRAIN_CONTAINER_URI,
"command": [],
"env": [
{
"name": "RESOLUTION",
"value": f"{RESOLUTION}",
},
],
"args": [
f"--input_csv_path={input_csv_path}",
f"--output_model_dir={model_dir}",
f"--epochs={epochs}",
f"--pretrained_backbone={backbone}",
f"--fpn_depth={fpn_depth}",
f"--learning_rate={learning_rate}",
f"--confidence_threshold={confidence_threshold}",
f"--iou_threshold={iou_threshold}",
],
},
}
]
common_util.check_quota(
project_id=PROJECT_ID,
region=REGION,
accelerator_type=ACCELERATOR_TYPE,
accelerator_count=ACCELERATOR_COUNT,
is_for_training=True,
)
# Add labels for the finetuning job.
labels = {
"mg-source": "notebook",
"mg-notebook-name": "model_garden_keras_yolov8.ipynb".split(".")[0],
}
labels["mg-tune"] = "publishers-google-models-yolov8"
versioned_model_id = model_name.lower().replace("_", "-")
labels["versioned-mg-tune"] = f"{labels['mg-tune']}-{versioned_model_id}"
train_job = aiplatform.CustomJob(
display_name=train_job_name,
project=PROJECT_ID,
worker_pool_specs=worker_pool_specs,
staging_bucket=STAGING_BUCKET,
labels=labels,
)
train_job.run()
print("The trained model is saved in: ", model_dir)In [ ]:
# @title Upload model
upload_job_name = common_util.get_job_name_with_datetime("upload_yolov8")
common_util.check_quota(
project_id=PROJECT_ID,
region=REGION,
accelerator_type=ACCELERATOR_TYPE,
accelerator_count=ACCELERATOR_COUNT,
is_for_training=False,
)
serving_env = {
"MODEL_ID": "keras-yolov8",
"DEPLOY_SOURCE": "notebook",
}
models["yolov8"] = aiplatform.Model.upload(
display_name=upload_job_name,
artifact_uri=model_dir,
serving_container_image_uri=SERVING_CONTAINER_URI,
serving_container_args=SERVING_CONTAINER_ARGS,
serving_container_environment_variables=serving_env,
model_garden_source_model_name="publishers/google/models/keras-yolov8",
)
print("The model name is: ", upload_job_name)In [ ]:
# @title Deploy model
deploy_model_name = common_util.get_job_name_with_datetime("deploy_yolov8")
endpoints["yolov8"] = models["yolov8"].deploy(
deployed_model_display_name=deploy_model_name,
machine_type="n1-standard-4",
traffic_split={"0": 100},
accelerator_type="NVIDIA_TESLA_T4",
accelerator_count=1,
min_replica_count=1,
max_replica_count=1,
system_labels={"NOTEBOOK_NAME": "model_garden_keras_yolov8.ipynb"},
)
endpoint_id = endpoints["yolov8"].name
print("The endpoint id is: ", endpoint_id)In [ ]:
# @title Predict
test_filepath = "gs://cloud-ml-data/img/openimage/1302/4677521502_6f2767039c_o.jpg" # @param {type:"string"}
image_bytes = tf.io.read_file(test_filepath)
image_resized = tf.expand_dims(decode_image(image_bytes), axis=0)
instances = get_prediction_instances(test_filepath, new_width=640)
predictions, _ = predict_custom_trained_model(
project=PROJECT_ID, location=REGION, endpoint_id=endpoint_id, instances=instances
)
predictions_dict = {
"boxes": tf.expand_dims(predictions[0]["boxes"], axis=0),
"classes": tf.expand_dims(predictions[0]["classes"], axis=0),
"confidence": tf.expand_dims(predictions[0]["confidence"], axis=0),
"num_detections": predictions[0]["num_detections"],
}
label_map = get_label_map(os.path.join(model_dir, "label_map.yaml"))
visualization.plot_bounding_box_gallery(
image_resized,
value_range=(0, 255),
rows=1,
cols=1,
y_pred=predictions_dict,
scale=5,
font_scale=0.7,
bounding_box_format="xywh",
class_mapping=label_map,
)In [ ]:
# @title Delete the models and endpoints
# @markdown Delete the experiment models and endpoints to recycle the resources
# @markdown and avoid unnecessary continuous charges that may incur.
# Undeploy model and delete endpoint.
for endpoint in endpoints.values():
endpoint.delete(force=True)
# Delete models.
for model in models.values():
model.delete()
delete_bucket = False # @param {type:"boolean"}
if delete_bucket:
! gcloud storage rm --recursive $BUCKET_NAME