From cb72a56e55e6db12f21efcd34bcc1e6ff7f7d7ff Mon Sep 17 00:00:00 2001 From: gericdong Date: Mon, 17 Apr 2023 10:56:43 -0400 Subject: [PATCH] feat: PyTorch train and deploy E2E with pre-built containers (#1728) * feat: PyTorch train and deploy E2E with pre-built containers * address review comments * address review comment 2 * suppress gsutil warning messags * Use unique names * Corrected the project id template --- ...ploy_models_with_prebuilt_containers.ipynb | 1174 +++++++++++++++++ 1 file changed, 1174 insertions(+) create mode 100644 notebooks/official/prediction/pytorch_train_deploy_models_with_prebuilt_containers.ipynb diff --git a/notebooks/official/prediction/pytorch_train_deploy_models_with_prebuilt_containers.ipynb b/notebooks/official/prediction/pytorch_train_deploy_models_with_prebuilt_containers.ipynb new file mode 100644 index 000000000..15cdfc0ae --- /dev/null +++ b/notebooks/official/prediction/pytorch_train_deploy_models_with_prebuilt_containers.ipynb @@ -0,0 +1,1174 @@ +{ + "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": [ + "# Train and deploy PyTorch models with prebuilt containers on Vertex AI\n", + "\n", + "\n", + "\n", + "\n", + " \n", + " \n", + " \n", + "
\n", + " \n", + " \"Colab Run in Colab\n", + " \n", + " \n", + " \n", + " \"GitHub\n", + " View on GitHub\n", + " \n", + " \n", + " \n", + " \"Vertex\n", + " Open in Vertex AI Workbench\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 custom training and prediction. These containers, which are organized by machine learning (ML) framework and framework version, include common dependencies that you might want to use in training code and serving predictions. Using pre-built containers is generally simpler than creating your own custom containers.\n", + "\n", + "This tutorial demonstrates how to train and deploy a PyTorch image model with prebuilt containers on Vertex AI.\n", + "\n", + "Learn more about [Prebuilt containers for custom training](https://cloud.google.com/vertex-ai/docs/training/pre-built-containers)\n", + "and [Prebuilt containers for prediction and explanation](https://cloud.google.com/vertex-ai/docs/predictions/pre-built-containers).\n" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "d975e698c9a4" + }, + "source": [ + "### Objective\n", + "\n", + "In this tutorial, you learn how to build, train and deploy a PyTorch image classification model using prebuilt containers for custom training and prediction.\n", + "\n", + "This tutorial uses the following Google Cloud ML services and resources:\n", + "\n", + "- `Vertex AI Training` service\n", + "- `Vertex AI Prediction` service\n", + "- `Vertex AI Model Registry`\n", + "- `Vertex AI Model` resources\n", + "- `Vertex AI Endpoint` resources\n", + "\n", + "\n", + "The steps performed include:\n", + "\n", + "- Package training application into a Python source distribution\n", + "- Configure and run training job in a prebuilt container\n", + "- Package model artifacts in a model archive file\n", + "- Upload model for deployment\n", + "- Deploy model using a prebuilt container for prediction\n", + "- Make online predictions\n" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "08d289fa873f" + }, + "source": [ + "### Dataset/Model\n", + "\n", + "In this tutorial, you use the [MNIST](https://pytorch.org/vision/main/generated/torchvision.datasets.MNIST.html) handwritten digit recognition dataset from PyTorch torchvision dataset. You train a simple Convolutional Neural Network on the MNIST dataset to identify handwritten digits." + ] + }, + { + "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.\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "2b4ef9b72d43" + }, + "outputs": [], + "source": [ + "! pip3 install --upgrade --quiet google-cloud-aiplatform \\\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**, try the following:\n", + "* Run `gcloud config list`.\n", + "* Run `gcloud projects list`.\n", + "* 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": "gkaCh23Pg2h8" + }, + "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": "f6b2ccc891ed" + }, + "source": [ + "**4. Service account or other**\n", + "* See how to grant Cloud Storage permissions to your service account at https://cloud.google.com/storage/docs/gsutil/commands/iam#ch-examples." + ] + }, + { + "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 = f\"gs://your-bucket-name-{PROJECT_ID}-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 io\n", + "import os\n", + "\n", + "from google.cloud import aiplatform\n", + "from matplotlib import pyplot as plt\n", + "from torch.utils.data import DataLoader\n", + "from torchvision import datasets, transforms" + ] + }, + { + "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": "4KyAr2MHg2h9" + }, + "outputs": [], + "source": [ + "aiplatform.init(project=PROJECT_ID, location=REGION, staging_bucket=BUCKET_URI)" + ] + }, + { + "attachments": {}, + "cell_type": "markdown", + "metadata": { + "id": "99DcIE4H54HA" + }, + "source": [ + "## Train a PyTorch model with a prebuilt container\n", + "\n", + "In this tutorial, you train a custom model using a prebuilt container for PyTorch models." + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "UOSGkpmJ6r7v" + }, + "source": [ + "### Package a training application\n", + "\n", + "#### Package layout\n", + "\n", + "Before you start the training, let's take a look at how a Python package is assembled for a custom training job. When extracted, the package contains the following:\n", + "\n", + "- PKG-INFO\n", + "- README.md\n", + "- setup.cfg\n", + "- setup.py\n", + "- trainer\n", + " - \\_\\_init\\_\\_.py\n", + " - task.py\n", + "\n", + "The files `setup.cfg` and `setup.py` are the instructions for installing the package into the operating environment of the Docker image." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "KIWLwZe3q1uu" + }, + "outputs": [], + "source": [ + "APP_NAME = \"pytorch-101-trainer\"\n", + "PYTHON_PACKAGE_APPLICATION_DIR = \"python_package\"\n", + "\n", + "source_package_file_name = (\n", + " f\"{PYTHON_PACKAGE_APPLICATION_DIR}/dist/{APP_NAME}-0.1.tar.gz\"\n", + ")\n", + "python_package_gcs_uri = f\"{BUCKET_URI}/pytorch/training/{APP_NAME}-0.1.tar.gz\"" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "KRy-N9edq2OM" + }, + "outputs": [], + "source": [ + "! mkdir {PYTHON_PACKAGE_APPLICATION_DIR}\n", + "! mkdir {PYTHON_PACKAGE_APPLICATION_DIR}/trainer\n", + "\n", + "! touch {PYTHON_PACKAGE_APPLICATION_DIR}/README.md\n", + "! touch {PYTHON_PACKAGE_APPLICATION_DIR}/trainer/__init__.py" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "1fi56Stvq2SG" + }, + "outputs": [], + "source": [ + "%%writefile ./{PYTHON_PACKAGE_APPLICATION_DIR}/setup.py\n", + "\n", + "import os\n", + "from setuptools import find_packages\n", + "from setuptools import setup\n", + "import setuptools\n", + "\n", + "from distutils.command.build import build as _build\n", + "import subprocess\n", + "\n", + "\n", + "REQUIRED_PACKAGES = [\n", + "]\n", + "\n", + "setup(\n", + " name='pytorch-101-trainer',\n", + " version='0.1',\n", + " install_requires=REQUIRED_PACKAGES,\n", + " packages=find_packages(),\n", + " include_package_data=True,\n", + " description='Vertex AI | Training | PyTorch | Image Classification | Python Package'\n", + ")" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "IVz8V3b57R0M" + }, + "source": [ + "#### Prepare the training script\n", + "\n", + "The file `trainer/task.py` is the Python script for executing the custom training job." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "cVgqlhM9q2Vo" + }, + "outputs": [], + "source": [ + "%%writefile ./{PYTHON_PACKAGE_APPLICATION_DIR}/trainer/task.py\n", + "\n", + "import os\n", + "import argparse\n", + "\n", + "import torch\n", + "from torch import nn\n", + "from torch.utils.data import DataLoader\n", + "from torchvision import datasets\n", + "from torchvision.transforms import ToTensor\n", + "\n", + "from google.cloud import storage\n", + "\n", + "\n", + "def load_data(batch_size):\n", + " # Download training data from open datasets\n", + " training_data = datasets.MNIST(\n", + " root=\"data\",\n", + " train=True,\n", + " download=True,\n", + " transform=ToTensor(),\n", + " )\n", + "\n", + " # Download test data from open datasets\n", + " test_data = datasets.MNIST(\n", + " root=\"data\",\n", + " train=False,\n", + " download=True,\n", + " transform=ToTensor(),\n", + " )\n", + "\n", + " # Create data loaders\n", + " train_dataloader = DataLoader(training_data, batch_size=batch_size)\n", + " test_dataloader = DataLoader(test_data, batch_size=batch_size)\n", + "\n", + " return train_dataloader, test_dataloader\n", + "\n", + "def create_model(device):\n", + " # Define model\n", + " class NeuralNetwork(nn.Module):\n", + " def __init__(self):\n", + " super().__init__()\n", + " self.flatten = nn.Flatten()\n", + " self.linear_relu_stack = nn.Sequential(\n", + " nn.Linear(28*28, 512),\n", + " nn.ReLU(),\n", + " nn.Linear(512, 512),\n", + " nn.ReLU(),\n", + " nn.Linear(512, 10)\n", + " )\n", + "\n", + " def forward(self, x):\n", + " x = self.flatten(x)\n", + " logits = self.linear_relu_stack(x)\n", + " return logits\n", + "\n", + " model = NeuralNetwork().to(device)\n", + "\n", + " return model\n", + "\n", + "def train(dataloader, model, loss_fn, optimizer, device):\n", + " size = len(dataloader.dataset)\n", + " model.train()\n", + " for batch, (X, y) in enumerate(dataloader):\n", + " X, y = X.to(device), y.to(device)\n", + "\n", + " pred = model(X)\n", + " loss = loss_fn(pred, y)\n", + "\n", + " optimizer.zero_grad()\n", + " loss.backward()\n", + " optimizer.step()\n", + "\n", + " if batch % 100 == 0:\n", + " loss, current = loss.item(), batch * len(X)\n", + " print(f\"loss: {loss:>7f} [{current:>5d}/{size:>5d}]\")\n", + "\n", + "def test(dataloader, model, loss_fn, device):\n", + " size = len(dataloader.dataset)\n", + " num_batches = len(dataloader)\n", + " model.eval()\n", + " test_loss, correct = 0, 0\n", + " with torch.no_grad():\n", + " for X, y in dataloader:\n", + " X, y = X.to(device), y.to(device)\n", + " pred = model(X)\n", + " test_loss += loss_fn(pred, y).item()\n", + " correct += (pred.argmax(1) == y).type(torch.float).sum().item()\n", + " test_loss /= num_batches\n", + " correct /= size\n", + " print(f\"Test Error: \\n Accuracy: {(100*correct):>0.1f}%, Avg loss: {test_loss:>8f} \\n\")\n", + "\n", + "\n", + "def main():\n", + " parser = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter)\n", + " parser.add_argument(\"--epochs\", type=int, help=\"Number of training epochs.\", default=2)\n", + " parser.add_argument(\"--batch_size\", type=int, help=\"Training batch size for one process.\", default=32)\n", + " parser.add_argument(\"--model_dir\", type=str, help=\"Directory for saving models.\", default=os.environ['AIP_MODEL_DIR'] if 'AIP_MODEL_DIR' in os.environ else \"\")\n", + " argv = parser.parse_args()\n", + "\n", + " train_dataloader, test_dataloader = load_data(argv.batch_size)\n", + "\n", + " # Get cpu or gpu device for training\n", + " device = \"cuda\" if torch.cuda.is_available() else \"mps\" if torch.backends.mps.is_available() else \"cpu\"\n", + " print(f\"Using {device} device\")\n", + "\n", + " model = create_model(device)\n", + "\n", + " # Define a loss function and an optimizer.\n", + " loss_fn = nn.CrossEntropyLoss()\n", + " optimizer = torch.optim.SGD(model.parameters(), lr=1e-3)\n", + "\n", + " for t in range(argv.epochs):\n", + " print(f\"Epoch {t+1}\\n-------------------------------\")\n", + " train(train_dataloader, model, loss_fn, optimizer, device)\n", + " test(test_dataloader, model, loss_fn, device)\n", + " print(\"Done!\")\n", + "\n", + " # Export the model to TorchScript\n", + " model_filename = \"pytorch-mnist.pt\"\n", + " local_path = os.path.join(\"/tmp\", model_filename)\n", + " model_scripted = torch.jit.script(model)\n", + " model_scripted.save(local_path)\n", + "\n", + " if (os.path.exists(local_path)):\n", + " # Upload the trained model to Cloud storage\n", + " storage_path = os.path.join(argv.model_dir, model_filename)\n", + " blob = storage.blob.Blob.from_string(storage_path, client=storage.Client())\n", + " blob.upload_from_filename(local_path)\n", + " print(f\"Saved model files in {argv.model_dir}\")\n", + "\n", + "if __name__ == \"__main__\":\n", + " main()" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "5WRTlkEw7ugt" + }, + "source": [ + "#### Create a Python source distribution\n", + "\n", + "You create a Python source distribution with your training application and upload the source distribution to your Cloud Storage bucket." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "WUS6ap2W1JeI" + }, + "outputs": [], + "source": [ + "! cd {PYTHON_PACKAGE_APPLICATION_DIR} && python3 setup.py sdist --formats=gztar\n", + "\n", + "! gsutil cp {source_package_file_name} {python_package_gcs_uri}\n", + "\n", + "print(f\"Python source distribution package location: {python_package_gcs_uri}\")" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "VOyihOMh8Q3B" + }, + "source": [ + "### Configure custom training job\n", + "\n", + "Configure a [custom job](https://cloud.google.com/vertex-ai/docs/training/create-custom-job) with the [prebuilt container](https://cloud.google.com/vertex-ai/docs/training/pre-built-containers) image for training code packaged as Python source distribution." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "x6ehn8V_1Prx" + }, + "outputs": [], + "source": [ + "JOB_DISPLAY_NAME = \"pytorch-custom-job-unique\"\n", + "python_module_name = \"trainer.task\"\n", + "PRE_BUILT_TRAINING_CONTAINER_IMAGE_URI = (\n", + " \"us-docker.pkg.dev/vertex-ai/training/pytorch-gpu.1-13:latest\"\n", + ")\n", + "\n", + "job = aiplatform.CustomPythonPackageTrainingJob(\n", + " display_name=JOB_DISPLAY_NAME,\n", + " python_package_gcs_uri=python_package_gcs_uri,\n", + " python_module_name=python_module_name,\n", + " container_uri=PRE_BUILT_TRAINING_CONTAINER_IMAGE_URI,\n", + ")" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "Un0gwmrv80Ya" + }, + "source": [ + "### Run custom training job\n", + "\n", + "Next, you run the custom job to start the training job by invoking the method `run`.\n", + "\n", + "**NOTE:** When using Vertex AI SDK for Python for submitting a training job, it creates a [training pipeline](https://cloud.google.com/vertex-ai/docs/training/create-training-pipeline) which launches the custom job on `Vertex AI Training` service.\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "VXdr0et51h-H" + }, + "outputs": [], + "source": [ + "MACHINE_TYPE = \"n1-standard-4\"\n", + "ACCELERATOR_TYPE = \"NVIDIA_TESLA_V100\"\n", + "ACCELERATOR_COUNT = 1\n", + "\n", + "EPOCHS = 2\n", + "BATCH_SIZE = 32\n", + "MODEL_DIR = f\"{BUCKET_URI}/{JOB_DISPLAY_NAME}\"\n", + "\n", + "training_args = [\n", + " \"--epochs\",\n", + " str(EPOCHS),\n", + " \"--batch_size\",\n", + " str(BATCH_SIZE),\n", + " \"--model_dir\",\n", + " MODEL_DIR,\n", + "]" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "Vl4cDy9r1nnX" + }, + "outputs": [], + "source": [ + "model = job.run(\n", + " machine_type=MACHINE_TYPE,\n", + " accelerator_type=ACCELERATOR_TYPE,\n", + " accelerator_count=ACCELERATOR_COUNT,\n", + " base_output_dir=MODEL_DIR,\n", + " replica_count=1,\n", + " args=training_args,\n", + " sync=True,\n", + ")" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "abjl3sgM-sOt" + }, + "source": [ + "#### Check model artifacts\n", + "\n", + "When the custom training job has completed, you check the model artifacts in the Cloud Storage location.\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "ch36B6ju16z9" + }, + "outputs": [], + "source": [ + "print(f\"Model artifacts are available at {MODEL_DIR}\")\n", + "! gsutil ls -l {MODEL_DIR}" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "ATeuG5I12BO9" + }, + "source": [ + "## Deploy the PyTorch model to a prebuilt container for prediction\n", + "\n", + "\n", + "You create a local directory and copy the model artifacts from the Cloud Storage to this local directory." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "awSZOd1i2EtZ" + }, + "outputs": [], + "source": [ + "# Create a local directory for model artifacts\n", + "model_path = \"model-unique\"\n", + "\n", + "if not os.path.exists(model_path):\n", + " ! mkdir {model_path}" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "9raKhhpt3epG" + }, + "outputs": [], + "source": [ + "model_name = \"pytorch-mnist.pt\"\n", + "model_file = f\"{model_path}/{model_name}\"\n", + "\n", + "! gsutil cp {MODEL_DIR}/{model_name} {model_file}" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "U6W-18UUCBMn" + }, + "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. 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. This custom handler extends class ImageClassifier from defaul handler image_classifier.py, takes an image and returns the number in that image." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "nW67E74lCDwA" + }, + "outputs": [], + "source": [ + "hander_file = f\"{model_path}/custom_handler.py\"" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "YYUOBdvq3Z8G" + }, + "outputs": [], + "source": [ + "%%writefile {hander_file}\n", + "\n", + "from torchvision import transforms\n", + "from ts.torch_handler.image_classifier import ImageClassifier\n", + "from torch.profiler import ProfilerActivity\n", + "\n", + "\n", + "class MNISTDigitClassifier(ImageClassifier):\n", + " \"\"\"\n", + " MNISTDigitClassifier handler class. This handler extends class ImageClassifier from image_classifier.py, a\n", + " default handler. This handler takes an image and returns the number in that image.\n", + " \"\"\"\n", + "\n", + " # Normalize MNIST images\n", + " image_processing = transforms.Compose([\n", + " transforms.ToTensor(),\n", + " transforms.Normalize((0.1307,), (0.3081,))\n", + " ])\n", + "\n", + " def __init__(self):\n", + " super(MNISTDigitClassifier, self).__init__()\n", + " self.profiler_args = {\n", + " \"activities\" : [ProfilerActivity.CPU],\n", + " \"record_shapes\": True,\n", + " }\n", + "\n", + "\n", + " def postprocess(self, data):\n", + " \"\"\"The post process of MNIST converts the predicted output response to a label.\n", + " Args:\n", + " data (list): The predicted output from the Inference with probabilities is passed\n", + " to the post-process function\n", + " Returns:\n", + " list : A list of dictionaries with predictions\n", + " \"\"\"\n", + " return data.argmax(1).tolist()" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "JY-yZ9xzCSSe" + }, + "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": "gMuQRwOB2oa3" + }, + "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": "t1mhQeX43jMP" + }, + "outputs": [], + "source": [ + "! torch-model-archiver -f \\\n", + " --model-name model \\\n", + " --version 1.0 \\\n", + " --serialized-file $model_file \\\n", + " --handler $hander_file \\\n", + " --export-path $model_path" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "nyx95vSPDJ3C" + }, + "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": "i9ltOguF322J" + }, + "outputs": [], + "source": [ + "MODEL_URI = f\"{BUCKET_URI}/model\"\n", + "\n", + "! gsutil -m rm -r -f $MODEL_URI\n", + "! gsutil -m cp -r $model_path $MODEL_URI\n", + "! gsutil ls $MODEL_URI" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "Yi3I6JotDw3k" + }, + "source": [ + "### Upload model for deployment\n", + "\n", + "Next, you upload the model artifacts to `Vertex AI Model Registry`, which will create a `Vertex AI Model` resource for your model. This tutorial uses the PyTorch v1.11 container, but for your own use case, you can choose from the list of [PyTorch prebuilt containers](https://cloud.google.com/vertex-ai/docs/predictions/pre-built-containers#pytorch)." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "9kyOITop38H7" + }, + "outputs": [], + "source": [ + "DEPLOY_IMAGE_URI = \"us-docker.pkg.dev/vertex-ai/prediction/pytorch-cpu.1-11:latest\"\n", + "\n", + "uploaded_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": "ko3IfZgbEG1J" + }, + "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": "MSukFOf04IMv" + }, + "outputs": [], + "source": [ + "DEPLOY_COMPUTE = \"n1-standard-4\"\n", + "\n", + "endpoint = uploaded_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": "AzHtXzv0E2sg" + }, + "source": [ + "## Make online predictions\n", + "\n", + "You use the MNIST dataset for the input for online predictions." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "kdRjvCOX4hVn" + }, + "outputs": [], + "source": [ + "# Download test data from PyTorch torchvision dataset\n", + "test_data = datasets.MNIST(\n", + " root=\"data\",\n", + " train=False,\n", + " download=True,\n", + " transform=transforms.ToTensor(),\n", + ")\n", + "\n", + "# Create data loaders\n", + "test_dataloader = DataLoader(test_data, batch_size=64)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "fDWLHAuPWo_7" + }, + "outputs": [], + "source": [ + "# Take one image as example for prediction\n", + "image, _ = test_data[0]\n", + "pil_image = transforms.ToPILImage()(image)\n", + "plt.imshow(pil_image)\n", + "plt.show()" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "BOV_iS-jFcTK" + }, + "source": [ + "### Get online predictions\n", + "\n", + "You send a `predict` request with encoded input image data to the `endpoint` and get prediction." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "kPqDTh184fM0" + }, + "outputs": [], + "source": [ + "buffered_image = io.BytesIO()\n", + "pil_image.save(buffered_image, format=\"JPEG\")\n", + "\n", + "data = {\"data\": base64.b64encode(buffered_image.getvalue()).decode(\"utf-8\")}\n", + "prediction = endpoint.predict(instances=[data])\n", + "\n", + "print(prediction.predictions)" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "-_pV0fq5FxgE" + }, + "source": [ + "### Make batch predictions (optional)\n", + "\n", + "Learn more about [making batch predictions](https://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/official/prediction/pytorch_image_classification_with_prebuilt_serving_containers.ipynb) from your PyTorch model." + ] + }, + { + "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": [ + "# Delete the deployment endpoint\n", + "endpoint.undeploy_all()\n", + "endpoint.delete()\n", + "\n", + "# Delete the model from Model Registry\n", + "uploaded_model.delete()\n", + "\n", + "# Delete Cloud Storage objects that were created\n", + "delete_bucket = False\n", + "if delete_bucket or os.getenv(\"IS_TESTING\"):\n", + " ! gsutil -m rm -r $BUCKET_URI" + ] + } + ], + "metadata": { + "colab": { + "name": "pytorch_train_deploy_models_with_prebuilt_containers.ipynb", + "toc_visible": true + }, + "kernelspec": { + "display_name": "Python 3", + "name": "python3" + } + }, + "nbformat": 4, + "nbformat_minor": 0 +}