diff --git a/notebooks/official/prediction/pytorch_image_classification_with_prebuilt_serving_containers.ipynb b/notebooks/official/prediction/pytorch_image_classification_with_prebuilt_serving_containers.ipynb
new file mode 100644
index 000000000..ded1e0e13
--- /dev/null
+++ b/notebooks/official/prediction/pytorch_image_classification_with_prebuilt_serving_containers.ipynb
@@ -0,0 +1,883 @@
+{
+ "cells": [
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {
+ "id": "ur8xi4C7S06n"
+ },
+ "outputs": [],
+ "source": [
+ "# Copyright 2023 Google LLC\n",
+ "#\n",
+ "# Licensed under the Apache License, Version 2.0 (the \"License\");\n",
+ "# you may not use this file except in compliance with the License.\n",
+ "# You may obtain a copy of the License at\n",
+ "#\n",
+ "# https://www.apache.org/licenses/LICENSE-2.0\n",
+ "#\n",
+ "# Unless required by applicable law or agreed to in writing, software\n",
+ "# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
+ "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
+ "# See the License for the specific language governing permissions and\n",
+ "# limitations under the License."
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {
+ "id": "JAPoU8Sm5E6e"
+ },
+ "source": [
+ "# Serving PyTorch image models with prebuilt containers on Vertex AI\n",
+ "\n",
+ "\n",
+ "
\n",
+ "\n",
+ " \n",
+ " \n",
+ " Run in Colab\n",
+ " \n",
+ " | \n",
+ " \n",
+ " \n",
+ " \n",
+ " View on GitHub\n",
+ " \n",
+ " | \n",
+ " \n",
+ " \n",
+ " \n",
+ " Open in Vertex AI Workbench\n",
+ " \n",
+ " | \n",
+ "
"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {
+ "id": "tvgnzT1CKxrO"
+ },
+ "source": [
+ "## Overview\n",
+ "\n",
+ "Vertex AI provides Docker container images that you run as prebuilt containers for serving predictions and explanations from trained model artifacts. Using prebuilt containers, you can deploy models for inference and prediction quickly and reliably at any scale.\n",
+ "\n",
+ "This tutorial demonstrates how to deploy a PyTorch image model using a prebuilt TorchServe container for serving online and batch predictions with minimal configuration.\n",
+ "\n",
+ "Learn more about [Pre-built containers for prediction](https://cloud.google.com/vertex-ai/docs/predictions/pre-built-containers)."
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {
+ "id": "d975e698c9a4"
+ },
+ "source": [
+ "### Objective\n",
+ "\n",
+ "In this tutorial, you learn how to package and deploy a PyTorch image classification model using a prebuilt TorchServe container for serving online and batch predictions.\n",
+ "\n",
+ "This tutorial uses the following Google Cloud ML services and resources:\n",
+ "\n",
+ "- `Vertex AI Model Registry`\n",
+ "- `Vertex AI Model` resources\n",
+ "- `Vertex AI Endpoint` resources\n",
+ "\n",
+ "The steps performed include:\n",
+ "\n",
+ "- Download a pretrained image model from PyTorch\n",
+ "- Create a custom model handler\n",
+ "- Package model artifacts in a model archive file\n",
+ "- Upload model for deployment\n",
+ "- Deploy model for prediction\n",
+ "- Make online predictions\n",
+ "- Make batch predictions"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {
+ "id": "08d289fa873f"
+ },
+ "source": [
+ "### Dataset\n",
+ "\n",
+ "In this example, you use the TensorFlow [flowers](http://download.tensorflow.org/example_images/flower_photos.tgz) dataset for inference."
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {
+ "id": "aed92deeb4a0"
+ },
+ "source": [
+ "### 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),\n",
+ "and [Cloud Storage pricing](https://cloud.google.com/storage/pricing),\n",
+ "and use the [Pricing Calculator](https://cloud.google.com/products/calculator/)\n",
+ "to generate a cost estimate based on your projected usage."
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {
+ "id": "i7EUnXsZhAGF"
+ },
+ "source": [
+ "## Installation\n",
+ "\n",
+ "Install the following packages required to execute this notebook."
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {
+ "id": "2b4ef9b72d43"
+ },
+ "outputs": [],
+ "source": [
+ "! pip3 install --upgrade --quiet google-cloud-aiplatform \\\n",
+ " tensorflow \\\n",
+ " torch \\\n",
+ " torchvision \\\n",
+ " torch-model-archiver"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {
+ "id": "58707a750154"
+ },
+ "source": [
+ "### Colab only: Uncomment the following cell to restart the kernel."
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {
+ "id": "f200f10a1da3"
+ },
+ "outputs": [],
+ "source": [
+ "# Automatically restart kernel after installs so that your environment can access the new packages\n",
+ "# import IPython\n",
+ "\n",
+ "# app = IPython.Application.instance()\n",
+ "# app.kernel.do_shutdown(True)"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {
+ "id": "BF1j6f9HApxa"
+ },
+ "source": [
+ "## Before you begin\n",
+ "\n",
+ "### 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",
+ "2. [Make sure that billing is enabled for your project](https://cloud.google.com/billing/docs/how-to/modify-project).\n",
+ "\n",
+ "3. [Enable the Vertex AI API](https://console.cloud.google.com/flows/enableapi?apiid=aiplatform.googleapis.com).\n",
+ "\n",
+ "4. If you are running this notebook locally, you need to install the [Cloud SDK](https://cloud.google.com/sdk)."
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {
+ "id": "WReHDGG5g0XY"
+ },
+ "source": [
+ "#### Set your project ID\n",
+ "\n",
+ "**If you don't know your project ID**, see the support page: [Locate the project ID](https://support.google.com/googleapi/answer/7014113)"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {
+ "id": "oM1iC_MfAts1"
+ },
+ "outputs": [],
+ "source": [
+ "PROJECT_ID = \"[your-project-id]\" # @param {type:\"string\"}\n",
+ "\n",
+ "# Set the project id\n",
+ "! gcloud config set project {PROJECT_ID}"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {
+ "id": "region"
+ },
+ "source": [
+ "#### Region\n",
+ "\n",
+ "You can also change the `REGION` variable used by Vertex AI. Learn more about [Vertex AI regions](https://cloud.google.com/vertex-ai/docs/general/locations)."
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {
+ "id": "tTy1gX11kCJY"
+ },
+ "outputs": [],
+ "source": [
+ "REGION = \"us-central1\" # @param {type: \"string\"}"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {
+ "id": "sBCra4QMA2wR"
+ },
+ "source": [
+ "### Authenticate your Google Cloud account\n",
+ "\n",
+ "Depending on your Jupyter environment, you may have to manually authenticate. Follow the relevant instructions below."
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {
+ "id": "74ccc9e52986"
+ },
+ "source": [
+ "**1. Vertex AI Workbench**\n",
+ "* Do nothing as you are already authenticated."
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {
+ "id": "de775a3773ba"
+ },
+ "source": [
+ "**2. Local JupyterLab instance, uncomment and run:**"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {
+ "id": "254614fa0c46"
+ },
+ "outputs": [],
+ "source": [
+ "# ! gcloud auth login"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {
+ "id": "ef21552ccea8"
+ },
+ "source": [
+ "**3. Colab, uncomment and run:**"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {
+ "id": "603adbbf0532"
+ },
+ "outputs": [],
+ "source": [
+ "# from google.colab import auth\n",
+ "# auth.authenticate_user()"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {
+ "id": "zgPO1eR3CYjk"
+ },
+ "source": [
+ "### Create a Cloud Storage bucket\n",
+ "\n",
+ "Create a storage bucket to store intermediate artifacts such as datasets."
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {
+ "id": "MzGDU7TWdts_"
+ },
+ "outputs": [],
+ "source": [
+ "BUCKET_URI = \"gs://your-bucket-name-unique\" # @param {type:\"string\"}"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {
+ "id": "-EcIXiGsCePi"
+ },
+ "source": [
+ "**Only if your bucket doesn't already exist**: Run the following cell to create your Cloud Storage bucket."
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {
+ "id": "NIq7R4HZCfIc"
+ },
+ "outputs": [],
+ "source": [
+ "! gsutil mb -l $REGION -p $PROJECT_ID $BUCKET_URI"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {
+ "id": "960505627ddf"
+ },
+ "source": [
+ "### Import libraries"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {
+ "id": "PyQmSRbKA8r-"
+ },
+ "outputs": [],
+ "source": [
+ "import base64\n",
+ "import json\n",
+ "import os\n",
+ "import pathlib\n",
+ "import urllib.request\n",
+ "\n",
+ "import tensorflow as tf\n",
+ "import torch\n",
+ "from google.cloud import aiplatform\n",
+ "from PIL import Image\n",
+ "from torchvision import models"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {
+ "id": "init_aip:mbsdk,all"
+ },
+ "source": [
+ "### Initialize Vertex AI SDK for Python\n",
+ "\n",
+ "Initialize the Vertex AI SDK for Python for your project."
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {
+ "id": "BnBAXs5XkCJZ"
+ },
+ "outputs": [],
+ "source": [
+ "aiplatform.init(project=PROJECT_ID, location=REGION, staging_bucket=BUCKET_URI)"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {
+ "id": "_McUaTTABIqu"
+ },
+ "source": [
+ "## Download a pre-trained image model\n",
+ "\n",
+ "For demonstration purposes, this tutorial uses a pretrained image model [resnet18](https://pytorch.org/vision/master/models/generated/torchvision.models.resnet18.html) from the PyTorch TorchVision."
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {
+ "id": "9Ss_3SGpqOWy"
+ },
+ "outputs": [],
+ "source": [
+ "# Create a local directory for model artifacts\n",
+ "model_path = \"model\"\n",
+ "\n",
+ "!rm -r $model_path\n",
+ "!mkdir $model_path\n",
+ "\n",
+ "model_name = \"resnet-18-custom-handler\"\n",
+ "model_file = f\"{model_path}/{model_name}.pt\""
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {
+ "id": "z62_eZEcEBKt"
+ },
+ "outputs": [],
+ "source": [
+ "# Use scripted mode to save the PyTorch model locally\n",
+ "model = models.resnet18(pretrained=True)\n",
+ "script_module = torch.jit.script(model)\n",
+ "script_module.save(model_file)"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {
+ "id": "FZpKK4FgDHeu"
+ },
+ "source": [
+ "## Create a custom model handler\n",
+ "\n",
+ "A custom model handler is a Python script that you package with the model when you use the model archiver. The script typically defines how to pre-process input data, invoke the model and post-process the output.\n",
+ "\n",
+ "TorchServe has [default handlers](https://pytorch.org/serve/default_handlers.html) for `image_classifier`, `image_segmenter`, `object_detector` and `text_classifier`. In this tutorial, you create a custom handler extending the default [`image_classifier`](https://github.com/pytorch/serve/blob/master/ts/torch_handler/image_classifier.py) handler.\n"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {
+ "id": "0jW_oUKNs0vQ"
+ },
+ "outputs": [],
+ "source": [
+ "hander_file = f\"{model_path}/custom_handler.py\""
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {
+ "id": "toIO7mTus453"
+ },
+ "outputs": [],
+ "source": [
+ "%%writefile {hander_file}\n",
+ "\n",
+ "import torch\n",
+ "import torch.nn.functional as F\n",
+ "from ts.torch_handler.image_classifier import ImageClassifier\n",
+ "from ts.utils.util import map_class_to_label\n",
+ "\n",
+ "\n",
+ "class CustomImageClassifier(ImageClassifier):\n",
+ "\n",
+ " # Only return the top 3 predictions\n",
+ " topk = 3\n",
+ "\n",
+ " def postprocess(self, data):\n",
+ " ps = F.softmax(data, dim=1)\n",
+ " probs, classes = torch.topk(ps, self.topk, dim=1)\n",
+ " probs = probs.tolist()\n",
+ " classes = classes.tolist()\n",
+ " return map_class_to_label(probs, self.mapping, classes)"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {
+ "id": "sD9ttWmqhPME"
+ },
+ "source": [
+ "## Download an index_to_name.json file\n",
+ "\n",
+ "PyTorch `image_classifier`, `text_classifier` and `object_detector` can all automatically map from numeric classes (0,1,2...) to friendly strings. To do this, simply include `index_to_name.json` that contains a mapping of class number to friendly name in your model archive file."
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {
+ "id": "m19wquufqOaL"
+ },
+ "outputs": [],
+ "source": [
+ "index_to_name_file = f\"{model_path}/index_to_name.json\"\n",
+ "\n",
+ "urllib.request.urlretrieve(\n",
+ " \"https://github.com/pytorch/serve/raw/master/examples/image_classifier/index_to_name.json\",\n",
+ " index_to_name_file,\n",
+ ")"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {
+ "id": "MJe1luzPjPCJ"
+ },
+ "source": [
+ "## Package the model artifacts in a model archive file\n",
+ "\n",
+ "You package all the model artifacts in a model archive file using the [`Torch model archiver`](https://github.com/pytorch/serve/tree/master/model-archiver).\n",
+ "\n",
+ "Note that the prebuilt PyTorch serving containers require the model archive file named as `model.mar` so you need to set the model-name as `model` in the `torch-model-archiver` command."
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {
+ "id": "mtP3s0MLqOL8"
+ },
+ "outputs": [],
+ "source": [
+ "# Add torch-model-archiver to the PATH\n",
+ "os.environ[\"PATH\"] = f'{os.environ.get(\"PATH\")}:~/.local/bin'"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {
+ "id": "pN1BqokBtHAS"
+ },
+ "outputs": [],
+ "source": [
+ "!torch-model-archiver -f \\\n",
+ " --model-name model \\\n",
+ " --version 1.0 \\\n",
+ " --serialized-file $model_file \\\n",
+ " --handler $hander_file \\\n",
+ " --extra-files $index_to_name_file \\\n",
+ " --export-path $model_path"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {
+ "id": "8qXVlPVTlRgP"
+ },
+ "source": [
+ "## Copy the model artifacts to Cloud Storage\n",
+ "\n",
+ "Next, use `gsutil` to copy the model artifacts to your Cloud Storage bucket."
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {
+ "id": "hsJ6YVaitHVW"
+ },
+ "outputs": [],
+ "source": [
+ "MODEL_URI = f\"{BUCKET_URI}/{model_name}\"\n",
+ "\n",
+ "!gsutil rm -r $MODEL_URI\n",
+ "!gsutil cp -r $model_path $MODEL_URI\n",
+ "!gsutil ls -al $MODEL_URI"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {
+ "id": "B9Ep0eANl7te"
+ },
+ "source": [
+ "## Upload model for deployment\n",
+ "\n",
+ "Next, you upload the model to `Vertex AI Model Registry`, which will create a `Vertex AI Model` resource for your model. You select a PyTorch prebuilt container image from a list of available [PyTorch prebuilt containers](https://cloud.google.com/vertex-ai/docs/predictions/pre-built-containers#pytorch), and specify the image URI in the parameter `serving_container_image_uri`."
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {
+ "id": "igJgzA6btPL7"
+ },
+ "outputs": [],
+ "source": [
+ "DEPLOY_IMAGE_URI = \"us-docker.pkg.dev/vertex-ai/prediction/pytorch-cpu.1-11:latest\"\n",
+ "\n",
+ "deployed_model = aiplatform.Model.upload(\n",
+ " display_name=model_name,\n",
+ " serving_container_image_uri=DEPLOY_IMAGE_URI,\n",
+ " artifact_uri=MODEL_URI,\n",
+ ")"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {
+ "id": "Tys97v6XpLmF"
+ },
+ "source": [
+ "## Deploy model for prediction\n",
+ "\n",
+ "Next, deploy your model for online prediction. You set the variable `DEPLOY_COMPUTE` to configure the machine type for the [compute resources](https://cloud.google.com/vertex-ai/docs/predictions/configure-compute) you will use for prediction."
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {
+ "id": "ka_zUL6-tPxW"
+ },
+ "outputs": [],
+ "source": [
+ "DEPLOY_COMPUTE = \"n1-standard-4\"\n",
+ "\n",
+ "endpoint = deployed_model.deploy(\n",
+ " deployed_model_display_name=model_name,\n",
+ " machine_type=DEPLOY_COMPUTE,\n",
+ " accelerator_type=None,\n",
+ " accelerator_count=0,\n",
+ ")"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {
+ "id": "Dp2oUReOpx7X"
+ },
+ "source": [
+ "## Make online predictions\n",
+ "\n",
+ "### Download an image dataset\n",
+ "In this example, you use the TensorFlow flowers dataset for the input data for predictions."
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {
+ "id": "nwKmNS-UqOP3"
+ },
+ "outputs": [],
+ "source": [
+ "data_dir = tf.keras.utils.get_file(\n",
+ " \"flower_photos\",\n",
+ " origin=\"https://storage.googleapis.com/download.tensorflow.org/example_images/flower_photos.tgz\",\n",
+ " untar=True,\n",
+ ")\n",
+ "\n",
+ "data_dir = pathlib.Path(data_dir)\n",
+ "images_files = list(data_dir.glob(\"daisy/*\"))"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {
+ "id": "0xRXkbxZqDkc"
+ },
+ "source": [
+ "### Get online predictions\n",
+ "\n",
+ "You send an `predict` request with encoded input image data to the `endpoint` and get predictions."
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {
+ "id": "lzooJLwUtXiU"
+ },
+ "outputs": [],
+ "source": [
+ "with open(images_files[0], \"rb\") as f:\n",
+ " data = {\"data\": base64.b64encode(f.read()).decode(\"utf-8\")}\n",
+ "\n",
+ "response = endpoint.predict(instances=[data])"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {
+ "id": "w3qgFiMyuUb3"
+ },
+ "outputs": [],
+ "source": [
+ "prediction = response.predictions[0]\n",
+ "prediction = dict(sorted(prediction.items(), key=lambda item: item[1], reverse=True))\n",
+ "\n",
+ "print(prediction)\n",
+ "image = Image.open(images_files[0])\n",
+ "image"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {
+ "id": "6o93-TkuqXS-"
+ },
+ "source": [
+ "## Make batch predictions\n",
+ "\n",
+ "### Create the batch input file\n",
+ "\n",
+ "You create a batch input file in JSONL format and store the input file in your Cloud Storage bucket."
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {
+ "id": "ChsjFzwW0rBj"
+ },
+ "outputs": [],
+ "source": [
+ "TEST_IMAGE_SIZE = 2\n",
+ "test_image_list = []\n",
+ "for i in range(TEST_IMAGE_SIZE):\n",
+ " test_image_list.append(str(images_files[i]))\n",
+ "\n",
+ "gcs_input_uri = f\"{BUCKET_URI}/test_images.json\"\n",
+ "\n",
+ "with tf.io.gfile.GFile(gcs_input_uri, \"w\") as f:\n",
+ " for test_image in test_image_list:\n",
+ " with open(test_image, \"rb\") as image_f:\n",
+ " data = {\"data\": base64.b64encode(image_f.read()).decode(\"utf-8\")}\n",
+ " f.write(json.dumps(data) + \"\\n\")"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {
+ "id": "_Ap-UO4tsJiO"
+ },
+ "source": [
+ "### Submit a batch prediction job\n",
+ "\n",
+ "You make a batch prediction by submitting a batch prediction job."
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {
+ "id": "HhH6KyIk0t3n"
+ },
+ "outputs": [],
+ "source": [
+ "JOB_DISPLAY_NAME = f\"{model_name}_batch_predict_job_unique\"\n",
+ "\n",
+ "batch_predict_job = deployed_model.batch_predict(\n",
+ " job_display_name=JOB_DISPLAY_NAME,\n",
+ " gcs_source=gcs_input_uri,\n",
+ " gcs_destination_prefix=BUCKET_URI,\n",
+ " instances_format=\"jsonl\",\n",
+ " model_parameters=None,\n",
+ " machine_type=DEPLOY_COMPUTE,\n",
+ ")"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {
+ "id": "2gVBOadysOSY"
+ },
+ "source": [
+ "### Get batch predictions\n",
+ "\n",
+ "Next, you get the predictions from the completed batch prediction job. The results are written to the Cloud Storage output bucket you specified in the batch prediction request. You call the method `iter_outputs()` to get a list of each Cloud Storage file generated with the results."
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {
+ "id": "wXcrGXrf0yVb"
+ },
+ "outputs": [],
+ "source": [
+ "bp_iter_outputs = batch_predict_job.iter_outputs()\n",
+ "\n",
+ "prediction_files = list()\n",
+ "for blob in bp_iter_outputs:\n",
+ " if blob.name.split(\"/\")[-1].startswith(\"prediction.results\"):\n",
+ " prediction_files.append(blob.name)"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {
+ "id": "4f7nXK1Y0vmd"
+ },
+ "outputs": [],
+ "source": [
+ "prediction_file = prediction_files[0]\n",
+ "\n",
+ "results = []\n",
+ "gfile_name = f\"{BUCKET_URI}/{prediction_file}\"\n",
+ "with tf.io.gfile.GFile(name=gfile_name, mode=\"r\") as gfile:\n",
+ " for line in gfile.readlines():\n",
+ " results.append(json.loads(line))\n",
+ "\n",
+ "# Take one result as an example and print out the prediction.\n",
+ "prediction = results[0][\"prediction\"]\n",
+ "prediction = dict(sorted(prediction.items(), key=lambda item: item[1], reverse=True))\n",
+ "print(prediction)"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {
+ "id": "TpV-iwP9qw9c"
+ },
+ "source": [
+ "## Cleaning up\n",
+ "\n",
+ "To clean up all Google Cloud resources used in this project, you can [delete the Google Cloud\n",
+ "project](https://cloud.google.com/resource-manager/docs/creating-managing-projects#shutting_down_projects) you used for the tutorial.\n",
+ "\n",
+ "Otherwise, you can delete the individual resources you created in this tutorial."
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {
+ "id": "sx_vKniMq9ZX"
+ },
+ "outputs": [],
+ "source": [
+ "endpoint.undeploy_all()\n",
+ "endpoint.delete()\n",
+ "\n",
+ "deployed_model.delete()\n",
+ "batch_predict_job.delete()\n",
+ "\n",
+ "delete_bucket = False\n",
+ "if delete_bucket:\n",
+ " ! gsutil -m rm -r $BUCKET_URI"
+ ]
+ }
+ ],
+ "metadata": {
+ "colab": {
+ "name": "pytorch_image_classification_with_prebuilt_serving_containers.ipynb",
+ "toc_visible": true
+ },
+ "kernelspec": {
+ "display_name": "Python 3",
+ "name": "python3"
+ }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 0
+}