From be688297426633e63620430176d503c1dc033e01 Mon Sep 17 00:00:00 2001
From: Jaycee Li <102714969+jaycee-li@users.noreply.github.com>
Date: Thu, 26 Oct 2023 11:54:02 -0700
Subject: [PATCH] Refactor SDK 2.0 notebooks (#2367)
* Refactor SDK 2.0 notebooks
* fix linting
* Update CODEOWNERS
* change codeowner username
* install dependencies
* update dependencies
* update folder name
* Update CODEOWNERS
---
notebooks/official/CODEOWNERS | 1 +
.../remote_hyperparameter_tuning.ipynb} | 2 +-
.../remote_prediction.ipynb} | 8 +-
.../remote_training_lightning.ipynb | 672 ++++++++++++++++
.../remote_training_pytorch.ipynb | 728 +++++++++++++++++
.../remote_training_sklearn.ipynb} | 598 +-------------
...training_tensorflow_with_autologging.ipynb | 749 ++++++++++++++++++
7 files changed, 2167 insertions(+), 591 deletions(-)
rename notebooks/official/{training/sdk2_remote_hyperparameter_tuning.ipynb => vertex_ai_sdk/remote_hyperparameter_tuning.ipynb} (99%)
rename notebooks/official/{prediction/sdk2_remote_prediction.ipynb => vertex_ai_sdk/remote_prediction.ipynb} (99%)
create mode 100644 notebooks/official/vertex_ai_sdk/remote_training_lightning.ipynb
create mode 100644 notebooks/official/vertex_ai_sdk/remote_training_pytorch.ipynb
rename notebooks/official/{training/sdk2_remote_training.ipynb => vertex_ai_sdk/remote_training_sklearn.ipynb} (53%)
create mode 100644 notebooks/official/vertex_ai_sdk/remote_training_tensorflow_with_autologging.ipynb
diff --git a/notebooks/official/CODEOWNERS b/notebooks/official/CODEOWNERS
index 3056cc813..95b2ac332 100644
--- a/notebooks/official/CODEOWNERS
+++ b/notebooks/official/CODEOWNERS
@@ -15,6 +15,7 @@
/ml_metadata @andrewferlitsch
/model_monitoring @andrewferlitsch
/tensorboard @zbl94
+/vertex_ai_sdk @jaycee-li
/bigquery_ml/bqml-online-prediction.ipynb @polong-lin
/model_monitoring/model_monitoring.ipynb @andrewferlitsch
diff --git a/notebooks/official/training/sdk2_remote_hyperparameter_tuning.ipynb b/notebooks/official/vertex_ai_sdk/remote_hyperparameter_tuning.ipynb
similarity index 99%
rename from notebooks/official/training/sdk2_remote_hyperparameter_tuning.ipynb
rename to notebooks/official/vertex_ai_sdk/remote_hyperparameter_tuning.ipynb
index 30cc113fc..b65c7e125 100644
--- a/notebooks/official/training/sdk2_remote_hyperparameter_tuning.ipynb
+++ b/notebooks/official/vertex_ai_sdk/remote_hyperparameter_tuning.ipynb
@@ -1760,7 +1760,7 @@
"metadata": {
"accelerator": "GPU",
"colab": {
- "name": "sdk2_remote_hyperparameter_tuning.ipynb",
+ "name": "remote_hyperparameter_tuning.ipynb",
"toc_visible": true
},
"kernelspec": {
diff --git a/notebooks/official/prediction/sdk2_remote_prediction.ipynb b/notebooks/official/vertex_ai_sdk/remote_prediction.ipynb
similarity index 99%
rename from notebooks/official/prediction/sdk2_remote_prediction.ipynb
rename to notebooks/official/vertex_ai_sdk/remote_prediction.ipynb
index b99c1994f..51eb0c95b 100644
--- a/notebooks/official/prediction/sdk2_remote_prediction.ipynb
+++ b/notebooks/official/vertex_ai_sdk/remote_prediction.ipynb
@@ -34,18 +34,18 @@
"\n",
"
\n",
" \n",
- " \n",
+ " \n",
" Run in Colab\n",
" \n",
" | \n",
" \n",
- " \n",
+ " \n",
" \n",
" View on GitHub\n",
" \n",
" | \n",
" \n",
- " \n",
+ " \n",
" Open in Vertex AI Workbench\n",
" \n",
" |
"
@@ -1383,7 +1383,7 @@
"metadata": {
"accelerator": "GPU",
"colab": {
- "name": "sdk2_remote_prediction.ipynb",
+ "name": "remote_prediction.ipynb",
"toc_visible": true
},
"kernelspec": {
diff --git a/notebooks/official/vertex_ai_sdk/remote_training_lightning.ipynb b/notebooks/official/vertex_ai_sdk/remote_training_lightning.ipynb
new file mode 100644
index 000000000..79d3da4ba
--- /dev/null
+++ b/notebooks/official/vertex_ai_sdk/remote_training_lightning.ipynb
@@ -0,0 +1,672 @@
+{
+ "cells": [
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {
+ "cellView": "form",
+ "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": "fsv4jGuU89rX"
+ },
+ "source": [
+ "# Vertex AI SDK 2.0 Remote Training for lightning model\n",
+ "\n",
+ "\n",
+ " \n",
+ " \n",
+ " Run in Colab\n",
+ " \n",
+ " | \n",
+ " \n",
+ " \n",
+ " \n",
+ " View on GitHub\n",
+ " \n",
+ " | \n",
+ " \n",
+ " \n",
+ " Open in Vertex AI Workbench\n",
+ " \n",
+ " |
"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {
+ "id": "overview:automl"
+ },
+ "source": [
+ "## Overview\n",
+ "\n",
+ "This tutorial demonstrates how to use Vertex AI SDK 2.0 for remote model training of a local model training job for OSS ML frameworks."
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {
+ "id": "objective:pipelines,automl"
+ },
+ "source": [
+ "### Objective\n",
+ "\n",
+ "In this tutorial, you learn to use `Vertex AI SDK 2.0` to remotely training models of various ML frameworks as a local (on-prem) training job.\n",
+ "\n",
+ "This tutorial uses the following Google Cloud ML services:\n",
+ "\n",
+ "- `Vertex AI Training`\n",
+ "- `Vertex AI Remote Training`\n",
+ "\n",
+ "The steps performed include:\n",
+ "\n",
+ "- Download and split the dataset\n",
+ "- Perform transformations as a Vertex AI remote training.\n",
+ "- For scikit-learn, PyTorch, TensorFlow, PyTorch Lightning\n",
+ " - Train the model remotely.\n",
+ " - Uptrain the pretrained model remotely.\n",
+ " - Evaluate both the pretrained and uptrained model.\n",
+ "\n",
+ "**Local-to-remote training**\n",
+ "\n",
+ "```\n",
+ "import vertexai\n",
+ "from my_module import MyModelClass\n",
+ "\n",
+ "vertexai.init(project=\"my-project\", location=\"my-location\", staging_bucket=\"gs://my-bucket\")\n",
+ "\n",
+ "# Switch to remote mode\n",
+ "vertexai.preview.init(remote=True)\n",
+ "\n",
+ "# Wrap the model class with `vertexai.preview.remote`\n",
+ "MyModelClass = vertexai.preview.remote(MyModelClass)\n",
+ "\n",
+ "# Instantiate the class\n",
+ "model = MyModelClass(...)\n",
+ "\n",
+ "# Optional set training config\n",
+ "model.fit.vertex.remote_config.display_name = \"MyModelClass-remote-training\"\n",
+ "model.fit.vertex.remote_config.staging_bucket = \"gs://my-bucket\"\n",
+ "\n",
+ "# This `fit` call will be executed remotely\n",
+ "model.fit(...)\n",
+ "```\n",
+ "\n",
+ "*Remote training supported OSS ML frameworks*\n",
+ "1. scikit-learn\n",
+ "2. TensorFlow\n",
+ "3. PyTorch\n",
+ "4. Pytorch Lightning\n",
+ "5. Custom model\n",
+ "\n",
+ "\n",
+ "---\n",
+ "\n",
+ "**Uptraining**\n",
+ "```\n",
+ "...\n",
+ "model = MyModelClass(...)\n",
+ "model.fit(...)\n",
+ "\n",
+ "# Save the trained model to Model Registry\n",
+ "registered_model = vertexai.preview.register(model)\n",
+ "\n",
+ "# The model can be loaded to a new (or current) local runtime\n",
+ "loaded_model = vertexai.preview.from_pretrained(\"registered-model-resource-id\")\n",
+ "\n",
+ "# Loaded model can cuntinue perform local-to-remote training\n",
+ "loaded_model.fit(...)\n",
+ "\n",
+ "```\n",
+ "\n",
+ "*Remote training supported OSS ML frameworks*\n",
+ "1. scikit-learn\n",
+ "2. TensorFlow\n",
+ "3. Custom model\n",
+ "4. PyTorch\n",
+ "\n",
+ "\n",
+ "\n",
+ "---\n",
+ "\n",
+ "**GPU Training**\n",
+ "```\n",
+ "...\n",
+ "model = MyModelClass(...)\n",
+ "\n",
+ "# Set enable_cuda to True to enable GPU training.\n",
+ "model.fit.vertex.remote_config.enable_cuda = True\n",
+ "\n",
+ "# (Optional) Training image and compute resources will be automatically\n",
+ "# handled by Vertex, but you can also config by yourself.\n",
+ "model.fit.vertex.remote_config.container_uri = \"your-cuda-image\"\n",
+ "model.fit.vertex.remote_config.machine_type = \"a2-highgpu-8g\"\n",
+ "model.fit.vertex.remote_config.accelerator_type = \"NVIDIA_TESLA_A100\"\n",
+ "model.fit.vertex.remote_config.accelerator_count = 8\n",
+ "\n",
+ "# Model will be trained remotely using GPU\n",
+ "model.fit(...)\n",
+ "```\n",
+ "\n",
+ "*GPU remote training supported OSS ML frameworks*\n",
+ "1. TensorFlow\n",
+ "2. PyTorch"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {
+ "id": "aef4f59195ad"
+ },
+ "source": [
+ "### Dataset\n",
+ "\n",
+ "This tutorial uses the IRIS dataset, which predicts the iris species."
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {
+ "id": "costs"
+ },
+ "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\n",
+ "pricing](https://cloud.google.com/vertex-ai/pricing), [Cloud Storage\n",
+ "pricing](https://cloud.google.com/storage/pricing), and use the [Pricing\n",
+ "Calculator](https://cloud.google.com/products/calculator/)\n",
+ "to generate a cost estimate based on your projected usage."
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {
+ "id": "install_aip:mbsdk"
+ },
+ "source": [
+ "## Installation\n",
+ "\n",
+ "Install the following packages required to execute this notebook."
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {
+ "id": "Fpp5OjGI6BaR"
+ },
+ "outputs": [],
+ "source": [
+ "! pip3 install --upgrade --quiet google-cloud-aiplatform[preview]\n",
+ "! pip3 install --upgrade --quiet scikit-learn\n",
+ "! pip3 install --upgrade --quiet lightning"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {
+ "id": "restart"
+ },
+ "source": [
+ "### Colab only: Uncomment the following cell to restart the kernel"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {
+ "id": "D-ZBOjErv5mM"
+ },
+ "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": "before_you_begin:nogpu"
+ },
+ "source": [
+ "## Before you begin\n",
+ "\n",
+ "### 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": "set_project_id"
+ },
+ "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": "2dw8q9fdQEH5"
+ },
+ "outputs": [],
+ "source": [
+ "REGION = \"us-central1\""
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {
+ "id": "gcp_authenticate"
+ },
+ "source": [
+ "### Authenticate your Google Cloud account\n",
+ "\n",
+ "Depending on your Jupyter environment, you may have to manually authenticate. Follow the relevant instructions below.\n",
+ "\n",
+ "**1. Vertex AI Workbench**\n",
+ "* Do nothing as you are already authenticated.\n",
+ "\n",
+ "**2. Local JupyterLab instance, uncomment and run:**"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {
+ "id": "ce6043da7b33"
+ },
+ "outputs": [],
+ "source": [
+ "# ! gcloud auth login"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {
+ "id": "0367eac06a10"
+ },
+ "source": [
+ "**3. Colab, uncomment and run:**"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {
+ "id": "21ad4dbb4a61"
+ },
+ "outputs": [],
+ "source": [
+ "# from google.colab import auth\n",
+ "# auth.authenticate_user()"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {
+ "id": "c13224697bfb"
+ },
+ "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": "bucket:mbsdk"
+ },
+ "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": "bucket"
+ },
+ "outputs": [],
+ "source": [
+ "BUCKET_URI = f\"gs://your-bucket-name-{PROJECT_ID}-unique\" # @param {type:\"string\"}"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {
+ "id": "autoset_bucket"
+ },
+ "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": "91c46850b49b"
+ },
+ "outputs": [],
+ "source": [
+ "! gsutil mb -l {REGION} -p {PROJECT_ID} {BUCKET_URI}"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {
+ "id": "setup_vars"
+ },
+ "source": [
+ "### Import libraries and define constants"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {
+ "id": "LF2bwT6q-of1"
+ },
+ "outputs": [],
+ "source": [
+ "import vertexai\n",
+ "from sklearn.datasets import load_iris\n",
+ "from sklearn.model_selection import train_test_split\n",
+ "from sklearn.preprocessing import StandardScaler"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {
+ "id": "init_aip:mbsdk"
+ },
+ "source": [
+ "## Initialize Vertex AI SDK for Python\n",
+ "\n",
+ "Initialize the Vertex AI SDK for Python for your project and corresponding bucket."
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {
+ "id": "-p0YAOY64cA5"
+ },
+ "outputs": [],
+ "source": [
+ "REMOTE_JOB_NAME = \"remote-scalar\"\n",
+ "REMOTE_JOB_BUCKET = f\"{BUCKET_URI}/{REMOTE_JOB_NAME}\"\n",
+ "\n",
+ "vertexai.init(\n",
+ " project=PROJECT_ID,\n",
+ " location=REGION,\n",
+ " staging_bucket=REMOTE_JOB_BUCKET,\n",
+ ")"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {
+ "id": "Ud6sRD1WWFIr"
+ },
+ "source": [
+ "## Prepare the dataset\n",
+ "\n",
+ "Now load the Iris dataset and split the data into train, retrain and test sets."
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {
+ "id": "N_tbI76DWRWt"
+ },
+ "outputs": [],
+ "source": [
+ "dataset = load_iris()\n",
+ "\n",
+ "X, X_retrain, y, y_retrain = train_test_split(\n",
+ " dataset.data, dataset.target, test_size=0.60, random_state=42\n",
+ ")\n",
+ "X_train, X_test, y_train, y_test = train_test_split(\n",
+ " X, y, test_size=0.20, random_state=42\n",
+ ")\n",
+ "\n",
+ "\n",
+ "transformer = StandardScaler()\n",
+ "X_train = transformer.fit_transform(X_train)\n",
+ "X_test = transformer.transform(X_test)\n",
+ "X_retrain = transformer.transform(X_retrain)"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {
+ "id": "Z6amYysl-h_x"
+ },
+ "source": [
+ "## PyTorch Lightning"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {
+ "id": "kr6WLoCIwXbJ"
+ },
+ "source": [
+ "### Remote training\n",
+ "\n",
+ "First, train a PyTorch Ligtning model as a remote training job:\n",
+ "\n",
+ "- Reinitialize Vertex AI for remote training.\n",
+ "- Set Trainer for the remote training job.\n",
+ "- Invoke Trainer locally which will launch the remote training job."
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {
+ "id": "H2kWuzcZLdVE"
+ },
+ "outputs": [],
+ "source": [
+ "# Switch to local mode for testing\n",
+ "vertexai.preview.init(remote=True)\n",
+ "\n",
+ "import lightning.pytorch as pl\n",
+ "import torch\n",
+ "\n",
+ "# Wrap classes to enable Vertex remote execution\n",
+ "pl.Trainer = vertexai.preview.remote(pl.Trainer)\n",
+ "\n",
+ "\n",
+ "# prepare data loaders\n",
+ "train_loader = torch.utils.data.DataLoader(\n",
+ " torch.utils.data.TensorDataset(\n",
+ " torch.tensor(X_train).to(torch.float32),\n",
+ " torch.tensor(y_train),\n",
+ " ),\n",
+ " batch_size=10,\n",
+ " shuffle=True,\n",
+ ")\n",
+ "\n",
+ "val_loader = torch.utils.data.DataLoader(\n",
+ " torch.utils.data.TensorDataset(\n",
+ " torch.tensor(X_test).to(torch.float32),\n",
+ " torch.tensor(y_test),\n",
+ " ),\n",
+ " batch_size=10,\n",
+ ")\n",
+ "\n",
+ "\n",
+ "# define the model\n",
+ "class LitLogisticRegression(pl.LightningModule):\n",
+ " def __init__(self, input_size: int, output_size: int):\n",
+ " super().__init__()\n",
+ " self.linear = torch.nn.Linear(input_size, output_size)\n",
+ " self.softmax = torch.nn.Softmax(dim=1)\n",
+ "\n",
+ " def forward(self, x):\n",
+ " return self.softmax(self.linear(x))\n",
+ "\n",
+ " def training_step(self, batch, batch_idx):\n",
+ " x, y = batch\n",
+ " y_hat = self(x)\n",
+ " loss = torch.nn.functional.cross_entropy(y_hat, y)\n",
+ " return loss\n",
+ "\n",
+ " def validation_step(self, batch, batch_idx):\n",
+ " x, y = batch\n",
+ " y_hat = self(x)\n",
+ " loss = torch.nn.functional.cross_entropy(y_hat, y)\n",
+ " return loss\n",
+ "\n",
+ " def configure_optimizers(self):\n",
+ " return torch.optim.Adam(self.parameters(), lr=0.05)\n",
+ "\n",
+ " def predict(self, X):\n",
+ " X = torch.tensor(X).to(torch.float32)\n",
+ " with torch.no_grad():\n",
+ " pred = torch.argmax(self(X), dim=1)\n",
+ " return pred\n",
+ "\n",
+ "\n",
+ "model = LitLogisticRegression(4, 3)\n",
+ "\n",
+ "# set up the trainer and training config\n",
+ "trainer = pl.Trainer(max_epochs=100, accelerator=\"cpu\")\n",
+ "\n",
+ "# Train model on Vertex\n",
+ "trainer.fit(model, train_dataloaders=train_loader, val_dataloaders=val_loader)"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {
+ "id": "b9GglFxvoeuw"
+ },
+ "source": [
+ "### Local evaluation\n",
+ "\n",
+ "Next, evaluate the pretrained version of the model, and compare the results."
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {
+ "id": "tG9sGXYb-nuU"
+ },
+ "outputs": [],
+ "source": [
+ "from sklearn.metrics import accuracy_score\n",
+ "\n",
+ "# Switch to local mode for testing\n",
+ "vertexai.preview.init(remote=False)\n",
+ "\n",
+ "# Evaluate model's accuracy score\n",
+ "print(f\"Train accuracy: {accuracy_score(y_train, model.predict(X_train))}\")\n",
+ "print(f\"Test accuracy: {accuracy_score(y_test, model.predict(X_test))}\")"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {
+ "id": "ArvOcC4B6BaY"
+ },
+ "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": "yYnCN1gV6BaZ"
+ },
+ "outputs": [],
+ "source": [
+ "import os\n",
+ "\n",
+ "delete_bucket = False\n",
+ "\n",
+ "if delete_bucket or os.getenv(\"IS_TESTING\"):\n",
+ " ! gsutil rm -rf {BUCKET_URI}"
+ ]
+ }
+ ],
+ "metadata": {
+ "accelerator": "GPU",
+ "colab": {
+ "collapsed_sections": [
+ "-aNMacb-Ghxs",
+ "6fC-g1B7wGgZ"
+ ],
+ "name": "remote_training_lightning.ipynb",
+ "toc_visible": true
+ },
+ "kernelspec": {
+ "display_name": "Python 3",
+ "name": "python3"
+ }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 0
+}
diff --git a/notebooks/official/vertex_ai_sdk/remote_training_pytorch.ipynb b/notebooks/official/vertex_ai_sdk/remote_training_pytorch.ipynb
new file mode 100644
index 000000000..8d899320f
--- /dev/null
+++ b/notebooks/official/vertex_ai_sdk/remote_training_pytorch.ipynb
@@ -0,0 +1,728 @@
+{
+ "cells": [
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {
+ "cellView": "form",
+ "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": "fsv4jGuU89rX"
+ },
+ "source": [
+ "# Vertex AI SDK 2.0 Remote Training for pytorch model\n",
+ "\n",
+ "\n",
+ " \n",
+ " \n",
+ " Run in Colab\n",
+ " \n",
+ " | \n",
+ " \n",
+ " \n",
+ " \n",
+ " View on GitHub\n",
+ " \n",
+ " | \n",
+ " \n",
+ " \n",
+ " Open in Vertex AI Workbench\n",
+ " \n",
+ " |
"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {
+ "id": "overview:automl"
+ },
+ "source": [
+ "## Overview\n",
+ "\n",
+ "This tutorial demonstrates how to use Vertex AI SDK 2.0 for remote model training of a local model training job for OSS ML frameworks."
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {
+ "id": "objective:pipelines,automl"
+ },
+ "source": [
+ "### Objective\n",
+ "\n",
+ "In this tutorial, you learn to use `Vertex AI SDK 2.0` to remotely training models of various ML frameworks as a local (on-prem) training job.\n",
+ "\n",
+ "This tutorial uses the following Google Cloud ML services:\n",
+ "\n",
+ "- `Vertex AI Training`\n",
+ "- `Vertex AI Remote Training`\n",
+ "\n",
+ "The steps performed include:\n",
+ "\n",
+ "- Download and split the dataset\n",
+ "- Perform transformations as a Vertex AI remote training.\n",
+ "- For scikit-learn, PyTorch, TensorFlow, PyTorch Lightning\n",
+ " - Train the model remotely.\n",
+ " - Uptrain the pretrained model remotely.\n",
+ " - Evaluate both the pretrained and uptrained model.\n",
+ "\n",
+ "**Local-to-remote training**\n",
+ "\n",
+ "```\n",
+ "import vertexai\n",
+ "from my_module import MyModelClass\n",
+ "\n",
+ "vertexai.init(project=\"my-project\", location=\"my-location\", staging_bucket=\"gs://my-bucket\")\n",
+ "\n",
+ "# Switch to remote mode\n",
+ "vertexai.preview.init(remote=True)\n",
+ "\n",
+ "# Wrap the model class with `vertexai.preview.remote`\n",
+ "MyModelClass = vertexai.preview.remote(MyModelClass)\n",
+ "\n",
+ "# Instantiate the class\n",
+ "model = MyModelClass(...)\n",
+ "\n",
+ "# Optional set training config\n",
+ "model.fit.vertex.remote_config.display_name = \"MyModelClass-remote-training\"\n",
+ "model.fit.vertex.remote_config.staging_bucket = \"gs://my-bucket\"\n",
+ "\n",
+ "# This `fit` call will be executed remotely\n",
+ "model.fit(...)\n",
+ "```\n",
+ "\n",
+ "*Remote training supported OSS ML frameworks*\n",
+ "1. scikit-learn\n",
+ "2. TensorFlow\n",
+ "3. PyTorch\n",
+ "4. Pytorch Lightning\n",
+ "5. Custom model\n",
+ "\n",
+ "\n",
+ "---\n",
+ "\n",
+ "**Uptraining**\n",
+ "```\n",
+ "...\n",
+ "model = MyModelClass(...)\n",
+ "model.fit(...)\n",
+ "\n",
+ "# Save the trained model to Model Registry\n",
+ "registered_model = vertexai.preview.register(model)\n",
+ "\n",
+ "# The model can be loaded to a new (or current) local runtime\n",
+ "loaded_model = vertexai.preview.from_pretrained(\"registered-model-resource-id\")\n",
+ "\n",
+ "# Loaded model can cuntinue perform local-to-remote training\n",
+ "loaded_model.fit(...)\n",
+ "\n",
+ "```\n",
+ "\n",
+ "*Remote training supported OSS ML frameworks*\n",
+ "1. scikit-learn\n",
+ "2. TensorFlow\n",
+ "3. Custom model\n",
+ "4. PyTorch\n",
+ "\n",
+ "\n",
+ "\n",
+ "---\n",
+ "\n",
+ "**GPU Training**\n",
+ "```\n",
+ "...\n",
+ "model = MyModelClass(...)\n",
+ "\n",
+ "# Set enable_cuda to True to enable GPU training.\n",
+ "model.fit.vertex.remote_config.enable_cuda = True\n",
+ "\n",
+ "# (Optional) Training image and compute resources will be automatically\n",
+ "# handled by Vertex, but you can also config by yourself.\n",
+ "model.fit.vertex.remote_config.container_uri = \"your-cuda-image\"\n",
+ "model.fit.vertex.remote_config.machine_type = \"a2-highgpu-8g\"\n",
+ "model.fit.vertex.remote_config.accelerator_type = \"NVIDIA_TESLA_A100\"\n",
+ "model.fit.vertex.remote_config.accelerator_count = 8\n",
+ "\n",
+ "# Model will be trained remotely using GPU\n",
+ "model.fit(...)\n",
+ "```\n",
+ "\n",
+ "*GPU remote training supported OSS ML frameworks*\n",
+ "1. TensorFlow\n",
+ "2. PyTorch"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {
+ "id": "aef4f59195ad"
+ },
+ "source": [
+ "### Dataset\n",
+ "\n",
+ "This tutorial uses the IRIS dataset, which predicts the iris species."
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {
+ "id": "costs"
+ },
+ "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\n",
+ "pricing](https://cloud.google.com/vertex-ai/pricing), [Cloud Storage\n",
+ "pricing](https://cloud.google.com/storage/pricing), and use the [Pricing\n",
+ "Calculator](https://cloud.google.com/products/calculator/)\n",
+ "to generate a cost estimate based on your projected usage."
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {
+ "id": "install_aip:mbsdk"
+ },
+ "source": [
+ "## Installation\n",
+ "\n",
+ "Install the following packages required to execute this notebook."
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {
+ "id": "Fpp5OjGI6BaR"
+ },
+ "outputs": [],
+ "source": [
+ "! pip3 install --upgrade --quiet google-cloud-aiplatform[preview]\n",
+ "! pip3 install --upgrade --quiet scikit-learn\n",
+ "! pip3 install --upgrade --quiet torch"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {
+ "id": "restart"
+ },
+ "source": [
+ "### Colab only: Uncomment the following cell to restart the kernel"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {
+ "id": "D-ZBOjErv5mM"
+ },
+ "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": "before_you_begin:nogpu"
+ },
+ "source": [
+ "## Before you begin\n",
+ "\n",
+ "### 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": "set_project_id"
+ },
+ "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": "2dw8q9fdQEH5"
+ },
+ "outputs": [],
+ "source": [
+ "REGION = \"us-central1\""
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {
+ "id": "gcp_authenticate"
+ },
+ "source": [
+ "### Authenticate your Google Cloud account\n",
+ "\n",
+ "Depending on your Jupyter environment, you may have to manually authenticate. Follow the relevant instructions below.\n",
+ "\n",
+ "**1. Vertex AI Workbench**\n",
+ "* Do nothing as you are already authenticated.\n",
+ "\n",
+ "**2. Local JupyterLab instance, uncomment and run:**"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {
+ "id": "ce6043da7b33"
+ },
+ "outputs": [],
+ "source": [
+ "# ! gcloud auth login"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {
+ "id": "0367eac06a10"
+ },
+ "source": [
+ "**3. Colab, uncomment and run:**"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {
+ "id": "21ad4dbb4a61"
+ },
+ "outputs": [],
+ "source": [
+ "# from google.colab import auth\n",
+ "# auth.authenticate_user()"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {
+ "id": "c13224697bfb"
+ },
+ "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": "bucket:mbsdk"
+ },
+ "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": "bucket"
+ },
+ "outputs": [],
+ "source": [
+ "BUCKET_URI = f\"gs://your-bucket-name-{PROJECT_ID}-unique\" # @param {type:\"string\"}"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {
+ "id": "autoset_bucket"
+ },
+ "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": "91c46850b49b"
+ },
+ "outputs": [],
+ "source": [
+ "! gsutil mb -l {REGION} -p {PROJECT_ID} {BUCKET_URI}"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {
+ "id": "setup_vars"
+ },
+ "source": [
+ "### Import libraries and define constants"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {
+ "id": "LF2bwT6q-of1"
+ },
+ "outputs": [],
+ "source": [
+ "import vertexai\n",
+ "from sklearn.datasets import load_iris\n",
+ "from sklearn.model_selection import train_test_split\n",
+ "from sklearn.preprocessing import StandardScaler"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {
+ "id": "init_aip:mbsdk"
+ },
+ "source": [
+ "## Initialize Vertex AI SDK for Python\n",
+ "\n",
+ "Initialize the Vertex AI SDK for Python for your project and corresponding bucket."
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {
+ "id": "-p0YAOY64cA5"
+ },
+ "outputs": [],
+ "source": [
+ "REMOTE_JOB_NAME = \"remote-scalar\"\n",
+ "REMOTE_JOB_BUCKET = f\"{BUCKET_URI}/{REMOTE_JOB_NAME}\"\n",
+ "\n",
+ "vertexai.init(\n",
+ " project=PROJECT_ID,\n",
+ " location=REGION,\n",
+ " staging_bucket=REMOTE_JOB_BUCKET,\n",
+ ")"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {
+ "id": "Ud6sRD1WWFIr"
+ },
+ "source": [
+ "## Prepare the dataset\n",
+ "\n",
+ "Now load the Iris dataset and split the data into train, retrain and test sets."
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {
+ "id": "N_tbI76DWRWt"
+ },
+ "outputs": [],
+ "source": [
+ "dataset = load_iris()\n",
+ "\n",
+ "X, X_retrain, y, y_retrain = train_test_split(\n",
+ " dataset.data, dataset.target, test_size=0.60, random_state=42\n",
+ ")\n",
+ "X_train, X_test, y_train, y_test = train_test_split(\n",
+ " X, y, test_size=0.20, random_state=42\n",
+ ")\n",
+ "\n",
+ "\n",
+ "transformer = StandardScaler()\n",
+ "X_train = transformer.fit_transform(X_train)\n",
+ "X_test = transformer.transform(X_test)\n",
+ "X_retrain = transformer.transform(X_retrain)"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {
+ "id": "kdg47y-wjoxT"
+ },
+ "source": [
+ "## PyTorch"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {
+ "id": "eTSejqbdv2RE"
+ },
+ "source": [
+ "### Remote training with CPU\n",
+ "\n",
+ "First, train a PyTorch model as a remote training job:\n",
+ "\n",
+ "- Reinitialize Vertex AI for remote training.\n",
+ "- Set TorchLogisticRegression for the remote training job.\n",
+ "- Invoke TorchLogisticRegression locally which will launch the remote training job."
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {
+ "id": "wtk_KqDX-5WI"
+ },
+ "outputs": [],
+ "source": [
+ "# Switch to remote mode for training\n",
+ "vertexai.preview.init(remote=True)\n",
+ "\n",
+ "import torch\n",
+ "from vertexai.preview import VertexModel\n",
+ "\n",
+ "\n",
+ "# define the custom model\n",
+ "class TorchLogisticRegression(VertexModel, torch.nn.Module):\n",
+ " def __init__(self, input_size: int, output_size: int):\n",
+ " torch.nn.Module.__init__(self)\n",
+ " VertexModel.__init__(self)\n",
+ " self.linear = torch.nn.Linear(input_size, output_size)\n",
+ " self.softmax = torch.nn.Softmax(dim=1)\n",
+ "\n",
+ " def forward(self, x):\n",
+ " return self.softmax(self.linear(x))\n",
+ "\n",
+ " @vertexai.preview.developer.mark.train()\n",
+ " def train(self, X, y, num_epochs, lr):\n",
+ " X, y = torch.tensor(X).to(torch.float32), torch.tensor(y)\n",
+ " dataloader = torch.utils.data.DataLoader(\n",
+ " torch.utils.data.TensorDataset(X, y),\n",
+ " batch_size=10,\n",
+ " shuffle=True,\n",
+ " generator=torch.Generator(device=X.device),\n",
+ " )\n",
+ "\n",
+ " criterion = torch.nn.CrossEntropyLoss()\n",
+ " optimizer = torch.optim.SGD(self.parameters(), lr=lr)\n",
+ "\n",
+ " for t in range(num_epochs):\n",
+ " for batch, (X, y) in enumerate(dataloader):\n",
+ " optimizer.zero_grad()\n",
+ " pred = self(X)\n",
+ " loss = criterion(pred, y)\n",
+ " loss.backward()\n",
+ " optimizer.step()\n",
+ "\n",
+ " @vertexai.preview.developer.mark.predict()\n",
+ " def predict(self, X):\n",
+ " X = torch.tensor(X).to(torch.float32)\n",
+ " with torch.no_grad():\n",
+ " pred = torch.argmax(self(X), dim=1)\n",
+ " return pred\n",
+ "\n",
+ "\n",
+ "# Instantiate model\n",
+ "model = TorchLogisticRegression(4, 3)\n",
+ "\n",
+ "# Train model on Vertex\n",
+ "model.train(X_train, y_train, num_epochs=100, lr=0.05)"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {
+ "id": "qQ53E9NK-lYB"
+ },
+ "source": [
+ "### Uptrain the pretrained model with GPU\n",
+ "\n",
+ "Next, get the registered model from the Vertex AI Model Registry. Then request the pretrained version of the model."
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {
+ "id": "1dV4pWtbRdkq"
+ },
+ "outputs": [],
+ "source": [
+ "registered_model = vertexai.preview.register(model)\n",
+ "\n",
+ "pulled_model = vertexai.preview.from_pretrained(\n",
+ " model_name=registered_model.resource_name\n",
+ ")"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {
+ "id": "rYuE40b36BaW"
+ },
+ "source": [
+ "Now train the model remotely via Vertex AI Training."
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {
+ "id": "pbjDOXzk9AkG"
+ },
+ "outputs": [],
+ "source": [
+ "pulled_model.train(X_retrain, y_retrain, num_epochs=100, lr=0.05)"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {
+ "id": "l35WfW-dkERU"
+ },
+ "source": [
+ "### Local evaluation\n",
+ "\n",
+ "Next, evaluate the pretrained and uptrained versions of the model, and compare the results."
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {
+ "id": "GFo8usi_kB-F"
+ },
+ "outputs": [],
+ "source": [
+ "from sklearn.metrics import accuracy_score\n",
+ "\n",
+ "# Switch to local mode for testing\n",
+ "vertexai.preview.init(remote=False)\n",
+ "\n",
+ "# Evaluate model's accuracy score\n",
+ "print(f\"Train accuracy: {accuracy_score(y_train, model.predict(X_train))}\")\n",
+ "print(f\"Test accuracy: {accuracy_score(y_test, model.predict(X_test))}\")\n",
+ "\n",
+ "# Evaluate uptrained model's accuracy score\n",
+ "print(f\"Train accuracy: {accuracy_score(y_train, pulled_model.predict(X_train))}\")\n",
+ "print(f\"Test accuracy: {accuracy_score(y_test, pulled_model.predict(X_test))}\")"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {
+ "id": "zZ6DLHxo6BaX"
+ },
+ "source": [
+ "#### Delete the registered model\n",
+ "\n",
+ "You can delete the registered model in the Vertex AI Model Registry with the delete() method."
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {
+ "id": "esqEbYAA6BaX"
+ },
+ "outputs": [],
+ "source": [
+ "registered_model.delete()"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {
+ "id": "ArvOcC4B6BaY"
+ },
+ "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": "yYnCN1gV6BaZ"
+ },
+ "outputs": [],
+ "source": [
+ "import os\n",
+ "\n",
+ "delete_bucket = False\n",
+ "\n",
+ "if delete_bucket or os.getenv(\"IS_TESTING\"):\n",
+ " ! gsutil rm -rf {BUCKET_URI}"
+ ]
+ }
+ ],
+ "metadata": {
+ "accelerator": "GPU",
+ "colab": {
+ "collapsed_sections": [
+ "-aNMacb-Ghxs",
+ "6fC-g1B7wGgZ"
+ ],
+ "name": "remote_training_pytorch.ipynb",
+ "toc_visible": true
+ },
+ "kernelspec": {
+ "display_name": "Python 3",
+ "name": "python3"
+ }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 0
+}
diff --git a/notebooks/official/training/sdk2_remote_training.ipynb b/notebooks/official/vertex_ai_sdk/remote_training_sklearn.ipynb
similarity index 53%
rename from notebooks/official/training/sdk2_remote_training.ipynb
rename to notebooks/official/vertex_ai_sdk/remote_training_sklearn.ipynb
index ab9cd26e0..f8d582cda 100644
--- a/notebooks/official/training/sdk2_remote_training.ipynb
+++ b/notebooks/official/vertex_ai_sdk/remote_training_sklearn.ipynb
@@ -30,22 +30,22 @@
"id": "fsv4jGuU89rX"
},
"source": [
- "# Vertex AI SDK 2.0 Vertex AI Remote Training for OSS ML frameworks\n",
+ "# Vertex AI SDK 2.0 Remote Training for sklearn model\n",
"\n",
"\n",
" \n",
- " \n",
+ " \n",
" Run in Colab\n",
" \n",
" | \n",
" \n",
- " \n",
+ " \n",
" \n",
" View on GitHub\n",
" \n",
" | \n",
" \n",
- " \n",
+ " \n",
" Open in Vertex AI Workbench\n",
" \n",
" |
"
@@ -222,9 +222,8 @@
},
"outputs": [],
"source": [
- "! pip3 install --upgrade --quiet google-cloud-aiplatform[preview,autologging]\n",
- "! pip3 install --upgrade --quiet lightning\n",
- "! pip3 install --upgrade --quiet tensorflow==2.12"
+ "! pip3 install --upgrade --quiet google-cloud-aiplatform[preview]\n",
+ "! pip3 install --upgrade --quiet scikit-learn"
]
},
{
@@ -446,10 +445,13 @@
},
"outputs": [],
"source": [
+ "REMOTE_JOB_NAME = \"remote-scalar\"\n",
+ "REMOTE_JOB_BUCKET = f\"{BUCKET_URI}/{REMOTE_JOB_NAME}\"\n",
+ "\n",
"vertexai.init(\n",
" project=PROJECT_ID,\n",
" location=REGION,\n",
- " staging_bucket=BUCKET_URI,\n",
+ " staging_bucket=REMOTE_JOB_BUCKET,\n",
")"
]
},
@@ -531,9 +533,6 @@
},
"outputs": [],
"source": [
- "REMOTE_JOB_NAME = \"remote-scalar\"\n",
- "REMOTE_JOB_BUCKET = f\"{BUCKET_URI}/{REMOTE_JOB_NAME}\"\n",
- "\n",
"# Wrap classes to enable Vertex remote execution\n",
"StandardScaler = vertexai.preview.remote(StandardScaler)\n",
"\n",
@@ -541,11 +540,6 @@
"# Instantiate transformer\n",
"transformer = StandardScaler()\n",
"\n",
- "# Set training config\n",
- "transformer.fit_transform.vertex.remote_config.display_name = (\n",
- " f\"{REMOTE_JOB_NAME}-fit-transformer\"\n",
- ")\n",
- "transformer.fit_transform.vertex.remote_config.staging_bucket = REMOTE_JOB_BUCKET\n",
"\n",
"# Execute transformer on Vertex\n",
"X_train = transformer.fit_transform(X_train)"
@@ -569,11 +563,6 @@
"outputs": [],
"source": [
"# Transform test dataset before calculate test score\n",
- "transformer.transform.vertex.remote_config.display_name = (\n",
- " REMOTE_JOB_NAME + \"-transformer\"\n",
- ")\n",
- "transformer.transform.vertex.remote_config.staging_bucket = REMOTE_JOB_BUCKET\n",
- "\n",
"X_test = transformer.transform(X_test)"
]
},
@@ -636,15 +625,13 @@
"vertexai.preview.init(remote=True)\n",
"\n",
"# Wrap classes to enable Vertex remote execution\n",
- "# Don't need this step after import hook is implemented\n",
"LogisticRegression = vertexai.preview.remote(LogisticRegression)\n",
"\n",
"# Instantiate model, warm_start=True for uptraining\n",
"model = LogisticRegression(warm_start=True)\n",
"\n",
- "# Set training config\n",
+ "# (Optional) Customize the display name of the remote job\n",
"model.fit.vertex.remote_config.display_name = REMOTE_JOB_NAME + \"-sklearn-model\"\n",
- "model.fit.vertex.remote_config.staging_bucket = REMOTE_JOB_BUCKET\n",
"\n",
"# Train model on Vertex\n",
"model = model.fit(X_train, y_train)"
@@ -749,567 +736,6 @@
"registered_model.delete()"
]
},
- {
- "cell_type": "markdown",
- "metadata": {
- "id": "kdg47y-wjoxT"
- },
- "source": [
- "## PyTorch"
- ]
- },
- {
- "cell_type": "markdown",
- "metadata": {
- "id": "eTSejqbdv2RE"
- },
- "source": [
- "### Remote training with CPU\n",
- "\n",
- "First, train a PyTorch model as a remote training job:\n",
- "\n",
- "- Reinitialize Vertex AI for remote training.\n",
- "- Set TorchLogisticRegression for the remote training job.\n",
- "- Invoke TorchLogisticRegression locally which will launch the remote training job."
- ]
- },
- {
- "cell_type": "code",
- "execution_count": null,
- "metadata": {
- "id": "wtk_KqDX-5WI"
- },
- "outputs": [],
- "source": [
- "# Switch to remote mode for training\n",
- "vertexai.preview.init(remote=True)\n",
- "\n",
- "import torch\n",
- "from vertexai.preview import VertexModel\n",
- "\n",
- "\n",
- "# define the custom model\n",
- "class TorchLogisticRegression(VertexModel, torch.nn.Module):\n",
- " def __init__(self, input_size: int, output_size: int):\n",
- " torch.nn.Module.__init__(self)\n",
- " VertexModel.__init__(self)\n",
- " self.linear = torch.nn.Linear(input_size, output_size)\n",
- " self.softmax = torch.nn.Softmax(dim=1)\n",
- "\n",
- " def forward(self, x):\n",
- " return self.softmax(self.linear(x))\n",
- "\n",
- " @vertexai.preview.developer.mark.train()\n",
- " def train(self, X, y, num_epochs, lr):\n",
- " X, y = torch.tensor(X).to(torch.float32), torch.tensor(y)\n",
- " dataloader = torch.utils.data.DataLoader(\n",
- " torch.utils.data.TensorDataset(X, y),\n",
- " batch_size=10,\n",
- " shuffle=True,\n",
- " generator=torch.Generator(device=X.device),\n",
- " )\n",
- "\n",
- " criterion = torch.nn.CrossEntropyLoss()\n",
- " optimizer = torch.optim.SGD(self.parameters(), lr=lr)\n",
- "\n",
- " for t in range(num_epochs):\n",
- " for batch, (X, y) in enumerate(dataloader):\n",
- " optimizer.zero_grad()\n",
- " pred = self(X)\n",
- " loss = criterion(pred, y)\n",
- " loss.backward()\n",
- " optimizer.step()\n",
- "\n",
- " @vertexai.preview.developer.mark.predict()\n",
- " def predict(self, X):\n",
- " X = torch.tensor(X).to(torch.float32)\n",
- " with torch.no_grad():\n",
- " pred = torch.argmax(self(X), dim=1)\n",
- " return pred\n",
- "\n",
- "\n",
- "# Instantiate model\n",
- "model = TorchLogisticRegression(4, 3)\n",
- "\n",
- "# Set training config\n",
- "model.train.vertex.remote_config.display_name = (\n",
- " REMOTE_JOB_NAME + \"-pytorch-custom-model\"\n",
- ")\n",
- "model.train.vertex.remote_config.staging_bucket = REMOTE_JOB_BUCKET\n",
- "\n",
- "# Train model on Vertex\n",
- "model.train(X_train, y_train, num_epochs=100, lr=0.05)"
- ]
- },
- {
- "cell_type": "markdown",
- "metadata": {
- "id": "qQ53E9NK-lYB"
- },
- "source": [
- "### Uptrain the pretrained model with GPU\n",
- "\n",
- "Next, get the registered model from the Vertex AI Model Registry. Then request the pretrained version of the model."
- ]
- },
- {
- "cell_type": "code",
- "execution_count": null,
- "metadata": {
- "id": "1dV4pWtbRdkq"
- },
- "outputs": [],
- "source": [
- "registered_model = vertexai.preview.register(model)\n",
- "\n",
- "pulled_model = vertexai.preview.from_pretrained(\n",
- " model_name=registered_model.resource_name\n",
- ")"
- ]
- },
- {
- "cell_type": "markdown",
- "metadata": {
- "id": "rYuE40b36BaW"
- },
- "source": [
- "Now train the model remotely via Vertex AI Training."
- ]
- },
- {
- "cell_type": "code",
- "execution_count": null,
- "metadata": {
- "id": "pbjDOXzk9AkG"
- },
- "outputs": [],
- "source": [
- "# Get Python version\n",
- "py_v = ! python3 --version\n",
- "py_v = py_v[0].split(\".\")[1]\n",
- "\n",
- "# only supported by Pythn 3.10+\n",
- "if int(py_v) >= 10:\n",
- " pulled_model.train.vertex.remote_config.enable_cuda = True\n",
- " pulled_model.train.vertex.remote_config.display_name = (\n",
- " REMOTE_JOB_NAME + \"-pytorch-custom-model-gpu\"\n",
- " )\n",
- " pulled_model.train(X_retrain, y_retrain, num_epochs=100, lr=0.05)"
- ]
- },
- {
- "cell_type": "markdown",
- "metadata": {
- "id": "l35WfW-dkERU"
- },
- "source": [
- "### Local evaluation\n",
- "\n",
- "Next, evaluate the pretrained and uptrained versions of the model, and compare the results."
- ]
- },
- {
- "cell_type": "code",
- "execution_count": null,
- "metadata": {
- "id": "GFo8usi_kB-F"
- },
- "outputs": [],
- "source": [
- "from sklearn.metrics import accuracy_score\n",
- "\n",
- "# Switch to local mode for testing\n",
- "vertexai.preview.init(remote=False)\n",
- "\n",
- "# Evaluate model's accuracy score\n",
- "print(f\"Train accuracy: {accuracy_score(y_train, model.predict(X_train))}\")\n",
- "print(f\"Test accuracy: {accuracy_score(y_test, model.predict(X_test))}\")\n",
- "\n",
- "# only supported by Pythn 3.10+\n",
- "if py_v >= \"3.10\":\n",
- " # Evaluate uptrained model's accuracy score\n",
- " print(f\"Train accuracy: {accuracy_score(y_train, pulled_model.predict(X_train))}\")\n",
- " print(f\"Test accuracy: {accuracy_score(y_test, pulled_model.predict(X_test))}\")"
- ]
- },
- {
- "cell_type": "markdown",
- "metadata": {
- "id": "zZ6DLHxo6BaX"
- },
- "source": [
- "#### Delete the registered model\n",
- "\n",
- "You can delete the registered model in the Vertex AI Model Registry with the delete() method."
- ]
- },
- {
- "cell_type": "code",
- "execution_count": null,
- "metadata": {
- "id": "esqEbYAA6BaX"
- },
- "outputs": [],
- "source": [
- "registered_model.delete()"
- ]
- },
- {
- "cell_type": "markdown",
- "metadata": {
- "id": "6fC-g1B7wGgZ"
- },
- "source": [
- "## TensorFlow"
- ]
- },
- {
- "cell_type": "markdown",
- "metadata": {
- "id": "ZRbU6gMUzaxq"
- },
- "source": [
- "### Remote training with GPU\n",
- "\n",
- "First, train a TensorFlow model as a remote training job:\n",
- "\n",
- "- Reinitialize Vertex AI for remote training.\n",
- "- Set Sequential for the remote training job.\n",
- "- Invoke Sequential locally which will launch the remote training job."
- ]
- },
- {
- "cell_type": "code",
- "execution_count": null,
- "metadata": {
- "id": "hgV3626pGoH2"
- },
- "outputs": [],
- "source": [
- "# Switch to remote mode for training\n",
- "vertexai.preview.init(remote=True)\n",
- "\n",
- "from tensorflow import keras\n",
- "\n",
- "# Wrap classes to enable Vertex remote execution\n",
- "keras.Sequential = vertexai.preview.remote(keras.Sequential)\n",
- "\n",
- "# Instantiate model\n",
- "model = keras.Sequential(\n",
- " [keras.layers.Dense(5, input_shape=(4,)), keras.layers.Softmax()]\n",
- ")\n",
- "\n",
- "# Specify optimizer and loss function\n",
- "model.compile(optimizer=\"adam\", loss=\"mean_squared_error\")\n",
- "\n",
- "# Set training config\n",
- "model.fit.vertex.remote_config.enable_cuda = True\n",
- "model.fit.vertex.remote_config.display_name = REMOTE_JOB_NAME + \"-keras-model-gpu\"\n",
- "model.fit.vertex.remote_config.staging_bucket = REMOTE_JOB_BUCKET\n",
- "\n",
- "# Manually set compute resources this time\n",
- "model.fit.vertex.remote_config.machine_type = \"n1-highmem-4\"\n",
- "model.fit.vertex.remote_config.accelerator_type = \"NVIDIA_TESLA_K80\"\n",
- "model.fit.vertex.remote_config.accelerator_count = 4\n",
- "\n",
- "\n",
- "# Train model on Vertex\n",
- "model.fit(X_train, y_train, epochs=10, batch_size=32)"
- ]
- },
- {
- "cell_type": "markdown",
- "metadata": {
- "id": "7JAwOWXX-uTV"
- },
- "source": [
- "### Uptrain the pretrained model with autologging feature\n",
- "\n",
- "Next, get the registered model from the Vertex AI Model Registry. Then request the pretrained version of the model."
- ]
- },
- {
- "cell_type": "code",
- "execution_count": null,
- "metadata": {
- "id": "oD2Bw3QzFdVV"
- },
- "outputs": [],
- "source": [
- "registered_model = vertexai.preview.register(model)\n",
- "\n",
- "pulled_model = vertexai.preview.from_pretrained(\n",
- " model_name=registered_model.resource_name\n",
- ")"
- ]
- },
- {
- "cell_type": "markdown",
- "metadata": {
- "id": "XZaN8KVu6BaX"
- },
- "source": [
- "Now train the model remotely via Vertex AI Training."
- ]
- },
- {
- "cell_type": "code",
- "execution_count": null,
- "metadata": {
- "id": "oBsY7MRBjeM8"
- },
- "outputs": [],
- "source": [
- "# Config experiment and turn on autologging\n",
- "vertexai.init(\n",
- " project=PROJECT_ID,\n",
- " location=REGION,\n",
- " staging_bucket=BUCKET_URI,\n",
- " experiment=\"test-remote-training-autologging\",\n",
- ")\n",
- "vertexai.preview.init(remote=True, autolog=True)\n",
- "\n",
- "# service account is required since autolog is True\n",
- "pulled_model.fit.vertex.remote_config.service_account = \"GCE\"\n",
- "\n",
- "# Set GPU configs to None\n",
- "pulled_model.fit.vertex.remote_config.enable_cuda = False\n",
- "pulled_model.fit.vertex.remote_config.machine_type = None\n",
- "pulled_model.fit.vertex.remote_config.accelerator_type = None\n",
- "pulled_model.fit.vertex.remote_config.accelerator_count = None\n",
- "\n",
- "pulled_model.fit.vertex.remote_config.display_name = (\n",
- " REMOTE_JOB_NAME + \"-keras-model-autologging\"\n",
- ")\n",
- "\n",
- "pulled_model.fit(X_retrain, y_retrain, epochs=10, batch_size=32)"
- ]
- },
- {
- "cell_type": "markdown",
- "metadata": {
- "id": "yO9fryXI6BaX"
- },
- "source": [
- "### Get experiments results\n",
- "\n",
- "Finally, get the Vertex AI Experiments results from the remote training job."
- ]
- },
- {
- "cell_type": "code",
- "execution_count": null,
- "metadata": {
- "id": "1ANPcw0Q-Fx7"
- },
- "outputs": [],
- "source": [
- "# View logged metrics & params\n",
- "vertexai.preview.get_experiment_df()\n",
- "\n",
- "# Turn off the autologging\n",
- "vertexai.preview.init(autolog=False)"
- ]
- },
- {
- "cell_type": "markdown",
- "metadata": {
- "id": "gHrWH_r0e6Kw"
- },
- "source": [
- "### Local evaluation\n",
- "\n",
- "Next, evaluate the pretrained and uptrained versions of the model, and compare the results."
- ]
- },
- {
- "cell_type": "code",
- "execution_count": null,
- "metadata": {
- "id": "gMyrRUe9a2ci"
- },
- "outputs": [],
- "source": [
- "# Switch to local mode for testing\n",
- "vertexai.preview.init(remote=False)\n",
- "\n",
- "# Evaluate model's mean square errors\n",
- "print(f\"Train loss: {model.evaluate(X_train, y_train)}\")\n",
- "print(f\"Test loss: {model.evaluate(X_test, y_test)}\")\n",
- "\n",
- "# Evaluate uptrained model's mean square errors\n",
- "print(f\"Train loss: {pulled_model.evaluate(X_train, y_train)}\")\n",
- "print(f\"Test loss: {pulled_model.evaluate(X_test, y_test)}\")"
- ]
- },
- {
- "cell_type": "markdown",
- "metadata": {
- "id": "u3RU6l0X6BaY"
- },
- "source": [
- "#### Delete the registered model\n",
- "\n",
- "You can delete the registered model in the Vertex AI Model Registry with the delete() method."
- ]
- },
- {
- "cell_type": "code",
- "execution_count": null,
- "metadata": {
- "id": "nlyFEi8M6BaY"
- },
- "outputs": [],
- "source": [
- "registered_model.delete()"
- ]
- },
- {
- "cell_type": "markdown",
- "metadata": {
- "id": "Z6amYysl-h_x"
- },
- "source": [
- "## PyTorch Lightning"
- ]
- },
- {
- "cell_type": "markdown",
- "metadata": {
- "id": "kr6WLoCIwXbJ"
- },
- "source": [
- "### Remote training\n",
- "\n",
- "First, train a PyTorch Ligtning model as a remote training job:\n",
- "\n",
- "- Reinitialize Vertex AI for remote training.\n",
- "- Set Trainer for the remote training job.\n",
- "- Invoke Trainer locally which will launch the remote training job."
- ]
- },
- {
- "cell_type": "code",
- "execution_count": null,
- "metadata": {
- "id": "H2kWuzcZLdVE"
- },
- "outputs": [],
- "source": [
- "# Switch to local mode for testing\n",
- "vertexai.preview.init(remote=True)\n",
- "\n",
- "import lightning.pytorch as pl\n",
- "import torch\n",
- "\n",
- "# Wrap classes to enable Vertex remote execution\n",
- "# Don't need this step after import hook is implemented\n",
- "pl.Trainer = vertexai.preview.remote(pl.Trainer)\n",
- "\n",
- "\n",
- "# prepare data loaders\n",
- "train_loader = torch.utils.data.DataLoader(\n",
- " torch.utils.data.TensorDataset(\n",
- " torch.tensor(X_train).to(torch.float32),\n",
- " torch.tensor(y_train),\n",
- " ),\n",
- " batch_size=10,\n",
- " shuffle=True,\n",
- ")\n",
- "\n",
- "val_loader = torch.utils.data.DataLoader(\n",
- " torch.utils.data.TensorDataset(\n",
- " torch.tensor(X_test).to(torch.float32),\n",
- " torch.tensor(y_test),\n",
- " ),\n",
- " batch_size=10,\n",
- ")\n",
- "\n",
- "\n",
- "# define the model\n",
- "class LitLogisticRegression(pl.LightningModule):\n",
- " def __init__(self, input_size: int, output_size: int):\n",
- " super().__init__()\n",
- " self.linear = torch.nn.Linear(input_size, output_size)\n",
- " self.softmax = torch.nn.Softmax(dim=1)\n",
- "\n",
- " def forward(self, x):\n",
- " return self.softmax(self.linear(x))\n",
- "\n",
- " def training_step(self, batch, batch_idx):\n",
- " x, y = batch\n",
- " y_hat = self(x)\n",
- " loss = torch.nn.functional.cross_entropy(y_hat, y)\n",
- " return loss\n",
- "\n",
- " def validation_step(self, batch, batch_idx):\n",
- " x, y = batch\n",
- " y_hat = self(x)\n",
- " loss = torch.nn.functional.cross_entropy(y_hat, y)\n",
- " return loss\n",
- "\n",
- " def configure_optimizers(self):\n",
- " return torch.optim.Adam(self.parameters(), lr=0.05)\n",
- "\n",
- " def predict(self, X):\n",
- " X = torch.tensor(X).to(torch.float32)\n",
- " with torch.no_grad():\n",
- " pred = torch.argmax(self(X), dim=1)\n",
- " return pred\n",
- "\n",
- "\n",
- "model = LitLogisticRegression(4, 3)\n",
- "\n",
- "# set up the trainer and training config\n",
- "trainer = pl.Trainer(max_epochs=100, accelerator=\"cpu\")\n",
- "\n",
- "trainer.fit.vertex.remote_config.display_name = REMOTE_JOB_NAME + \"-lightning-model\"\n",
- "trainer.fit.vertex.remote_config.staging_bucket = REMOTE_JOB_BUCKET\n",
- "\n",
- "# Train model on Vertex\n",
- "trainer.fit(model, train_dataloaders=train_loader, val_dataloaders=val_loader)"
- ]
- },
- {
- "cell_type": "markdown",
- "metadata": {
- "id": "b9GglFxvoeuw"
- },
- "source": [
- "### Local evaluation\n",
- "\n",
- "Next, evaluate the pretrained version of the model, and compare the results."
- ]
- },
- {
- "cell_type": "code",
- "execution_count": null,
- "metadata": {
- "id": "tG9sGXYb-nuU"
- },
- "outputs": [],
- "source": [
- "from sklearn.metrics import accuracy_score\n",
- "\n",
- "# Switch to local mode for testing\n",
- "vertexai.preview.init(remote=False)\n",
- "\n",
- "# Evaluate model's accuracy score\n",
- "print(f\"Train accuracy: {accuracy_score(y_train, model.predict(X_train))}\")\n",
- "print(f\"Test accuracy: {accuracy_score(y_test, model.predict(X_test))}\")"
- ]
- },
- {
- "cell_type": "markdown",
- "metadata": {
- "id": "FduZbn2s6BaY"
- },
- "source": [
- "#### Delete the registered model\n",
- "\n",
- "You can delete the registered model in the Vertex AI Model Registry with the delete() method."
- ]
- },
{
"cell_type": "markdown",
"metadata": {
@@ -1348,7 +774,7 @@
"-aNMacb-Ghxs",
"6fC-g1B7wGgZ"
],
- "name": "sdk2_remote_training.ipynb",
+ "name": "remote_training_sklearn.ipynb",
"toc_visible": true
},
"kernelspec": {
diff --git a/notebooks/official/vertex_ai_sdk/remote_training_tensorflow_with_autologging.ipynb b/notebooks/official/vertex_ai_sdk/remote_training_tensorflow_with_autologging.ipynb
new file mode 100644
index 000000000..1f039a218
--- /dev/null
+++ b/notebooks/official/vertex_ai_sdk/remote_training_tensorflow_with_autologging.ipynb
@@ -0,0 +1,749 @@
+{
+ "cells": [
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {
+ "cellView": "form",
+ "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": "fsv4jGuU89rX"
+ },
+ "source": [
+ "# Vertex AI SDK 2.0 Remote Training for keras model\n",
+ "\n",
+ "\n",
+ " \n",
+ " \n",
+ " Run in Colab\n",
+ " \n",
+ " | \n",
+ " \n",
+ " \n",
+ " \n",
+ " View on GitHub\n",
+ " \n",
+ " | \n",
+ " \n",
+ " \n",
+ " Open in Vertex AI Workbench\n",
+ " \n",
+ " |
"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {
+ "id": "overview:automl"
+ },
+ "source": [
+ "## Overview\n",
+ "\n",
+ "This tutorial demonstrates how to use Vertex AI SDK 2.0 for remote model training of a local model training job for OSS ML frameworks."
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {
+ "id": "objective:pipelines,automl"
+ },
+ "source": [
+ "### Objective\n",
+ "\n",
+ "In this tutorial, you learn to use `Vertex AI SDK 2.0` to remotely training models of various ML frameworks as a local (on-prem) training job.\n",
+ "\n",
+ "This tutorial uses the following Google Cloud ML services:\n",
+ "\n",
+ "- `Vertex AI Training`\n",
+ "- `Vertex AI Remote Training`\n",
+ "\n",
+ "The steps performed include:\n",
+ "\n",
+ "- Download and split the dataset\n",
+ "- Perform transformations as a Vertex AI remote training.\n",
+ "- For scikit-learn, PyTorch, TensorFlow, PyTorch Lightning\n",
+ " - Train the model remotely.\n",
+ " - Uptrain the pretrained model remotely.\n",
+ " - Evaluate both the pretrained and uptrained model.\n",
+ "\n",
+ "**Local-to-remote training**\n",
+ "\n",
+ "```\n",
+ "import vertexai\n",
+ "from my_module import MyModelClass\n",
+ "\n",
+ "vertexai.init(project=\"my-project\", location=\"my-location\", staging_bucket=\"gs://my-bucket\")\n",
+ "\n",
+ "# Switch to remote mode\n",
+ "vertexai.preview.init(remote=True)\n",
+ "\n",
+ "# Wrap the model class with `vertexai.preview.remote`\n",
+ "MyModelClass = vertexai.preview.remote(MyModelClass)\n",
+ "\n",
+ "# Instantiate the class\n",
+ "model = MyModelClass(...)\n",
+ "\n",
+ "# Optional set training config\n",
+ "model.fit.vertex.remote_config.display_name = \"MyModelClass-remote-training\"\n",
+ "model.fit.vertex.remote_config.staging_bucket = \"gs://my-bucket\"\n",
+ "\n",
+ "# This `fit` call will be executed remotely\n",
+ "model.fit(...)\n",
+ "```\n",
+ "\n",
+ "*Remote training supported OSS ML frameworks*\n",
+ "1. scikit-learn\n",
+ "2. TensorFlow\n",
+ "3. PyTorch\n",
+ "4. Pytorch Lightning\n",
+ "5. Custom model\n",
+ "\n",
+ "\n",
+ "---\n",
+ "\n",
+ "**Uptraining**\n",
+ "```\n",
+ "...\n",
+ "model = MyModelClass(...)\n",
+ "model.fit(...)\n",
+ "\n",
+ "# Save the trained model to Model Registry\n",
+ "registered_model = vertexai.preview.register(model)\n",
+ "\n",
+ "# The model can be loaded to a new (or current) local runtime\n",
+ "loaded_model = vertexai.preview.from_pretrained(\"registered-model-resource-id\")\n",
+ "\n",
+ "# Loaded model can cuntinue perform local-to-remote training\n",
+ "loaded_model.fit(...)\n",
+ "\n",
+ "```\n",
+ "\n",
+ "*Remote training supported OSS ML frameworks*\n",
+ "1. scikit-learn\n",
+ "2. TensorFlow\n",
+ "3. Custom model\n",
+ "4. PyTorch\n",
+ "\n",
+ "\n",
+ "\n",
+ "---\n",
+ "\n",
+ "**GPU Training**\n",
+ "```\n",
+ "...\n",
+ "model = MyModelClass(...)\n",
+ "\n",
+ "# Set enable_cuda to True to enable GPU training.\n",
+ "model.fit.vertex.remote_config.enable_cuda = True\n",
+ "\n",
+ "# (Optional) Training image and compute resources will be automatically\n",
+ "# handled by Vertex, but you can also config by yourself.\n",
+ "model.fit.vertex.remote_config.container_uri = \"your-cuda-image\"\n",
+ "model.fit.vertex.remote_config.machine_type = \"a2-highgpu-8g\"\n",
+ "model.fit.vertex.remote_config.accelerator_type = \"NVIDIA_TESLA_A100\"\n",
+ "model.fit.vertex.remote_config.accelerator_count = 8\n",
+ "\n",
+ "# Model will be trained remotely using GPU\n",
+ "model.fit(...)\n",
+ "```\n",
+ "\n",
+ "*GPU remote training supported OSS ML frameworks*\n",
+ "1. TensorFlow\n",
+ "2. PyTorch"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {
+ "id": "aef4f59195ad"
+ },
+ "source": [
+ "### Dataset\n",
+ "\n",
+ "This tutorial uses the IRIS dataset, which predicts the iris species."
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {
+ "id": "costs"
+ },
+ "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\n",
+ "pricing](https://cloud.google.com/vertex-ai/pricing), [Cloud Storage\n",
+ "pricing](https://cloud.google.com/storage/pricing), and use the [Pricing\n",
+ "Calculator](https://cloud.google.com/products/calculator/)\n",
+ "to generate a cost estimate based on your projected usage."
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {
+ "id": "install_aip:mbsdk"
+ },
+ "source": [
+ "## Installation\n",
+ "\n",
+ "Install the following packages required to execute this notebook."
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {
+ "id": "Fpp5OjGI6BaR"
+ },
+ "outputs": [],
+ "source": [
+ "! pip3 install --upgrade --quiet google-cloud-aiplatform[preview,autologging]\n",
+ "! pip3 install --upgrade --quiet scikit-learn\n",
+ "! pip3 install --upgrade --quiet tensorflow==2.12"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {
+ "id": "restart"
+ },
+ "source": [
+ "### Colab only: Uncomment the following cell to restart the kernel"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {
+ "id": "D-ZBOjErv5mM"
+ },
+ "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": "before_you_begin:nogpu"
+ },
+ "source": [
+ "## Before you begin\n",
+ "\n",
+ "### 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": "set_project_id"
+ },
+ "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": "2dw8q9fdQEH5"
+ },
+ "outputs": [],
+ "source": [
+ "REGION = \"us-central1\""
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {
+ "id": "gcp_authenticate"
+ },
+ "source": [
+ "### Authenticate your Google Cloud account\n",
+ "\n",
+ "Depending on your Jupyter environment, you may have to manually authenticate. Follow the relevant instructions below.\n",
+ "\n",
+ "**1. Vertex AI Workbench**\n",
+ "* Do nothing as you are already authenticated.\n",
+ "\n",
+ "**2. Local JupyterLab instance, uncomment and run:**"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {
+ "id": "ce6043da7b33"
+ },
+ "outputs": [],
+ "source": [
+ "# ! gcloud auth login"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {
+ "id": "0367eac06a10"
+ },
+ "source": [
+ "**3. Colab, uncomment and run:**"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {
+ "id": "21ad4dbb4a61"
+ },
+ "outputs": [],
+ "source": [
+ "# from google.colab import auth\n",
+ "# auth.authenticate_user()"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {
+ "id": "c13224697bfb"
+ },
+ "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": "bucket:mbsdk"
+ },
+ "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": "bucket"
+ },
+ "outputs": [],
+ "source": [
+ "BUCKET_URI = f\"gs://your-bucket-name-{PROJECT_ID}-unique\" # @param {type:\"string\"}"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {
+ "id": "autoset_bucket"
+ },
+ "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": "91c46850b49b"
+ },
+ "outputs": [],
+ "source": [
+ "! gsutil mb -l {REGION} -p {PROJECT_ID} {BUCKET_URI}"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {
+ "id": "setup_vars"
+ },
+ "source": [
+ "### Import libraries and define constants"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {
+ "id": "LF2bwT6q-of1"
+ },
+ "outputs": [],
+ "source": [
+ "import vertexai\n",
+ "from sklearn.datasets import load_iris\n",
+ "from sklearn.model_selection import train_test_split\n",
+ "from sklearn.preprocessing import StandardScaler"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {
+ "id": "init_aip:mbsdk"
+ },
+ "source": [
+ "## Initialize Vertex AI SDK for Python\n",
+ "\n",
+ "Initialize the Vertex AI SDK for Python for your project and corresponding bucket."
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {
+ "id": "-p0YAOY64cA5"
+ },
+ "outputs": [],
+ "source": [
+ "REMOTE_JOB_NAME = \"remote-scalar\"\n",
+ "REMOTE_JOB_BUCKET = f\"{BUCKET_URI}/{REMOTE_JOB_NAME}\"\n",
+ "\n",
+ "vertexai.init(\n",
+ " project=PROJECT_ID,\n",
+ " location=REGION,\n",
+ " staging_bucket=REMOTE_JOB_BUCKET,\n",
+ ")"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {
+ "id": "Ud6sRD1WWFIr"
+ },
+ "source": [
+ "## Prepare the dataset\n",
+ "\n",
+ "Now load the Iris dataset turn it into tf dataset."
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {
+ "id": "N_tbI76DWRWt"
+ },
+ "outputs": [],
+ "source": [
+ "import tensorflow as tf\n",
+ "\n",
+ "dataset = load_iris()\n",
+ "\n",
+ "X, X_retrain, y, y_retrain = train_test_split(\n",
+ " dataset.data, dataset.target, test_size=0.60, random_state=42\n",
+ ")\n",
+ "X_train, X_test, y_train, y_test = train_test_split(\n",
+ " X, y, test_size=0.20, random_state=42\n",
+ ")\n",
+ "\n",
+ "\n",
+ "transformer = StandardScaler()\n",
+ "X_train = transformer.fit_transform(X_train)\n",
+ "X_test = transformer.transform(X_test)\n",
+ "X_retrain = transformer.transform(X_retrain)\n",
+ "\n",
+ "\n",
+ "tf_train_dataset = tf.data.Dataset.from_tensor_slices((X_train, y_train))\n",
+ "tf_train_dataset = tf_train_dataset.shuffle(buffer_size=64).batch(32)\n",
+ "\n",
+ "tf_retrain_dataset = tf.data.Dataset.from_tensor_slices((X_retrain, y_retrain))\n",
+ "tf_retrain_dataset = tf_retrain_dataset.shuffle(buffer_size=64).batch(32)\n",
+ "\n",
+ "tf_test_dataset = tf.data.Dataset.from_tensor_slices((X_test, y_test))\n",
+ "tf_test_dataset = tf_test_dataset.shuffle(buffer_size=64).batch(32)"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {
+ "id": "6fC-g1B7wGgZ"
+ },
+ "source": [
+ "## TensorFlow"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {
+ "id": "ZRbU6gMUzaxq"
+ },
+ "source": [
+ "### Remote training with GPU\n",
+ "\n",
+ "First, train a TensorFlow model as a remote training job:\n",
+ "\n",
+ "- Reinitialize Vertex AI for remote training.\n",
+ "- Set Sequential for the remote training job.\n",
+ "- Invoke Sequential locally which will launch the remote training job."
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {
+ "id": "hgV3626pGoH2"
+ },
+ "outputs": [],
+ "source": [
+ "# Switch to remote mode for training\n",
+ "vertexai.preview.init(remote=True)\n",
+ "\n",
+ "from tensorflow import keras\n",
+ "\n",
+ "# Wrap classes to enable Vertex remote execution\n",
+ "keras.Sequential = vertexai.preview.remote(keras.Sequential)\n",
+ "\n",
+ "# Instantiate model\n",
+ "model = keras.Sequential(\n",
+ " [keras.layers.Dense(5, input_shape=(4,)), keras.layers.Softmax()]\n",
+ ")\n",
+ "\n",
+ "# Specify optimizer and loss function\n",
+ "model.compile(optimizer=\"adam\", loss=\"mean_squared_error\")\n",
+ "\n",
+ "# Enable GPU training in remote_config\n",
+ "model.fit.vertex.remote_config.enable_cuda = True\n",
+ "\n",
+ "# Train model on Vertex\n",
+ "model.fit(tf_train_dataset, epochs=10)"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {
+ "id": "7JAwOWXX-uTV"
+ },
+ "source": [
+ "### Uptrain the pretrained model with autologging feature\n",
+ "\n",
+ "Next, get the registered model from the Vertex AI Model Registry. Then request the pretrained version of the model."
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {
+ "id": "oD2Bw3QzFdVV"
+ },
+ "outputs": [],
+ "source": [
+ "registered_model = vertexai.preview.register(model)\n",
+ "\n",
+ "pulled_model = vertexai.preview.from_pretrained(\n",
+ " model_name=registered_model.resource_name\n",
+ ")"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {
+ "id": "XZaN8KVu6BaX"
+ },
+ "source": [
+ "Now train the model remotely via Vertex AI Training."
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {
+ "id": "oBsY7MRBjeM8"
+ },
+ "outputs": [],
+ "source": [
+ "# Config experiment and turn on autologging\n",
+ "vertexai.init(\n",
+ " project=PROJECT_ID,\n",
+ " location=REGION,\n",
+ " staging_bucket=REMOTE_JOB_BUCKET,\n",
+ " experiment=\"test-remote-training-autologging\",\n",
+ ")\n",
+ "vertexai.preview.init(remote=True, autolog=True)\n",
+ "\n",
+ "# service account is required since autolog is True\n",
+ "pulled_model.fit.vertex.remote_config.service_account = \"GCE\"\n",
+ "\n",
+ "# Turn off GPU training\n",
+ "pulled_model.fit.vertex.remote_config.enable_cuda = False\n",
+ "\n",
+ "# Train model on Vertex\n",
+ "pulled_model.fit(tf_retrain_dataset, epochs=10)"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {
+ "id": "yO9fryXI6BaX"
+ },
+ "source": [
+ "### Get experiments results\n",
+ "\n",
+ "Finally, get the Vertex AI Experiments results from the remote training job."
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {
+ "id": "1ANPcw0Q-Fx7"
+ },
+ "outputs": [],
+ "source": [
+ "# View logged metrics & params\n",
+ "vertexai.preview.get_experiment_df()\n",
+ "\n",
+ "# Turn off the autologging\n",
+ "vertexai.preview.init(autolog=False)"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {
+ "id": "gHrWH_r0e6Kw"
+ },
+ "source": [
+ "### Local evaluation\n",
+ "\n",
+ "Next, evaluate the pretrained and uptrained versions of the model, and compare the results."
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {
+ "id": "gMyrRUe9a2ci"
+ },
+ "outputs": [],
+ "source": [
+ "# Switch to local mode for testing\n",
+ "vertexai.preview.init(remote=False)\n",
+ "\n",
+ "# Evaluate model's mean square errors\n",
+ "print(f\"Train loss: {model.evaluate(tf_train_dataset)}\")\n",
+ "print(f\"Test loss: {model.evaluate(tf_test_dataset)}\")\n",
+ "\n",
+ "# Evaluate uptrained model's mean square errors\n",
+ "print(f\"Train loss: {pulled_model.evaluate(tf_retrain_dataset)}\")\n",
+ "print(f\"Test loss: {pulled_model.evaluate(tf_test_dataset)}\")"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {
+ "id": "u3RU6l0X6BaY"
+ },
+ "source": [
+ "#### Delete the registered model\n",
+ "\n",
+ "You can delete the registered model in the Vertex AI Model Registry with the delete() method."
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {
+ "id": "nlyFEi8M6BaY"
+ },
+ "outputs": [],
+ "source": [
+ "registered_model.delete()"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {
+ "id": "ArvOcC4B6BaY"
+ },
+ "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": "yYnCN1gV6BaZ"
+ },
+ "outputs": [],
+ "source": [
+ "import os\n",
+ "\n",
+ "delete_bucket = False\n",
+ "\n",
+ "if delete_bucket or os.getenv(\"IS_TESTING\"):\n",
+ " ! gsutil rm -rf {BUCKET_URI}"
+ ]
+ }
+ ],
+ "metadata": {
+ "accelerator": "GPU",
+ "colab": {
+ "collapsed_sections": [
+ "-aNMacb-Ghxs",
+ "6fC-g1B7wGgZ"
+ ],
+ "name": "remote_training_tensorflow_with_autologging.ipynb",
+ "toc_visible": true
+ },
+ "kernelspec": {
+ "display_name": "Python 3",
+ "name": "python3"
+ }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 0
+}