model_garden/notebooks/official/prediction/sdk2_remote_prediction.ipynb

In [ ]:

# Copyright 2023 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

Vertex AI SDK 2.0 Vertex AI Remote Prediction

Run in Colab

View on GitHub

Open in Vertex AI Workbench

Overview

This tutorial demonstrates how to use Vertex AI SDK 2.0 for remote prediction of a local model prediction request.

Objective

In this tutorial, you learn to use Vertex AI SDK 2.0 to do a remote prediction from a local prediction of models trained in various OSS ML frameworks.

This tutorial uses the following Google Cloud ML services:

• Vertex AI Training
• Vertex AI Remote Training
• Vertex AI Remote Prediction

The steps performed include:

• Download and split the dataset
• Perform transformations as a Vertex AI remote training.
• For scikit-learn, PyTorch, TensorFlow, PyTorch Lightning
  • Train the model remotely.
  • Uptrain the pretrained model remotely.
  • Evaluate both the pretrained and uptrained model.
  • Make a prediction remotely

Local-to-remote prediction

from google.cloud.aiplatform.private_preview import vertex_ai
from my_module import MyModelClass

MyModelClass = vertex_ai.remote(MyModelClass)

vertex_ai.init(remote=True, project="my-project", location="my-location")

model = MyModelClass(...)

# optional
model.PredictionConfig.display_name = "MyModelClass-remote-prediction"
model.PredictionConfig.staging_bucket = "my-bucket"

model.predict(...)

Remote prediction support OSS ML frameworks

1. scikit-learn
2. TensorFlow
3. Custom model
4. PyTorch
5. PyTorch Lightning

Dataset

This tutorial uses the IRIS dataset, which predicts the iris species.

Costs

This tutorial uses billable components of Google Cloud:

• Vertex AI
• Cloud Storage

Learn about Vertex AI pricing, Cloud Storage pricing, and use the Pricing Calculator to generate a cost estimate based on your projected usage.

Installation

Install the following packages required to execute this notebook.

In [ ]:

! pip3 install --quiet google-cloud-aiplatform[preview,autologging]
! pip3 install --upgrade --quiet lightning
! pip3 install --upgrade --quiet tensorflow==2.12

Colab only: Uncomment the following cell to restart the kernel

In [ ]:

# Automatically restart kernel after installs so that your environment can access the new packages
# import IPython

# app = IPython.Application.instance()
# app.kernel.do_shutdown(True)

Before you begin

Set your project ID

If you don't know your project ID, try the following:

• Run gcloud config list.
• Run gcloud projects list.
• See the support page: Locate the project ID

In [ ]:

PROJECT_ID = "[your-project-id]"  # @param {type:"string"}

# Set the project id
! gcloud config set project {PROJECT_ID}

Region

You can also change the REGION variable used by Vertex AI. Learn more about Vertex AI regions.

In [ ]:

REGION = "us-central1"

Authenticate your Google Cloud account

Depending on your Jupyter environment, you may have to manually authenticate. Follow the relevant instructions below.

1. Vertex AI Workbench

• Do nothing as you are already authenticated.

2. Local JupyterLab instance, uncomment and run:

In [ ]:

# ! gcloud auth login

3. Colab, uncomment and run:

In [ ]:

# from google.colab import auth
# auth.authenticate_user()

4. Service account or other

• See how to grant Cloud Storage permissions to your service account at https://cloud.google.com/storage/docs/gsutil/commands/iam#ch-examples.

Create a Cloud Storage bucket

Create a storage bucket to store intermediate artifacts such as datasets.

In [ ]:

BUCKET_URI = f"gs://your-bucket-name-{PROJECT_ID}-unique"  # @param {type:"string"}

Only if your bucket doesn't already exist: Run the following cell to create your Cloud Storage bucket.

In [ ]:

! gsutil mb -l {REGION} -p {PROJECT_ID} {BUCKET_URI}

Import libraries and define constants

In [ ]:

import vertexai.preview
from sklearn.datasets import load_iris
from sklearn.linear_model import LogisticRegression
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import StandardScaler

Initialize Vertex AI SDK for Python

Initialize the Vertex AI SDK for Python for your project and corresponding bucket.

In [ ]:

vertexai.init(
    project=PROJECT_ID,
    location=REGION,
    staging_bucket=BUCKET_URI,
)

Prepare the dataset

Now load the Iris dataset and split the data into train, retrain and test sets.

In [ ]:

dataset = load_iris()

X, X_retrain, y, y_retrain = train_test_split(
    dataset.data, dataset.target, test_size=0.60, random_state=42
)
X_train, X_test, y_train, y_test = train_test_split(
    X, y, test_size=0.20, random_state=42
)

print("Data size: ", len(dataset.target))
print("X_train size: ", len(X_train))
print("X_retrain size: ", len(X_retrain))
print("X_test size: ", len(X_test))

Feature transformation

Next, you do feature transformations on the data using the Vertex AI remote training service.

First, you re-initialize Vertex AI to enable remote training.

In [ ]:

# Switch to remote mode for training
vertexai.preview.init(remote=True)

Execute remote job for fit_transform() on training data

Next, indicate that the StandardScalar class is to be executed remotely. Then set up the data transform and call the fit_transform() method is executed remotely.

In [ ]:

REMOTE_JOB_NAME = "remote-scalar"
REMOTE_JOB_BUCKET = f"{BUCKET_URI}/{REMOTE_JOB_NAME}"

# Wrap classes to enable Vertex remote execution
# Don't need this step after import hook is implemented
StandardScaler = vertexai.preview.remote(StandardScaler)


# Instantiate transformer
transformer = StandardScaler()

# Set training config
transformer.fit_transform.vertex.remote_config.display_name = (
    f"{REMOTE_JOB_NAME}-fit-transformer"
)
transformer.fit_transform.vertex.remote_config.staging_bucket = REMOTE_JOB_BUCKET

# Execute transformer on Vertex
X_train = transformer.fit_transform(X_train)

Remote transform on test data

In [ ]:

# Transform test dataset before calculate test score
transformer.transform.vertex.remote_config.display_name = (
    REMOTE_JOB_NAME + "-transformer"
)
transformer.transform.vertex.remote_config.staging_bucket = REMOTE_JOB_BUCKET

X_test = transformer.transform(X_test)

Local transform on retrain data

In [ ]:

# Switch to local transformation
vertexai.preview.init(remote=False)

X_retrain = transformer.transform(X_retrain)

scikit-learn

Remote training

First, train the scikit-learn model as a remote training job:

• Reinitialize Vertex AI for remote training.
• Set LogisticRegression for the remote training job.
• Invoke LogisticRegression locally which will launch the remote training job.

In [ ]:

# Switch to remote mode for training
vertexai.preview.init(remote=True)

# Wrap classes to enable Vertex remote execution
# Don't need this step after import hook is implemented
LogisticRegression = vertexai.preview.remote(LogisticRegression)

# Instantiate model, warm_start=True for uptraining
model = LogisticRegression(warm_start=True)

# Set training config
model.fit.vertex.remote_config.display_name = REMOTE_JOB_NAME + "-sklearn-model"
model.fit.vertex.remote_config.staging_bucket = REMOTE_JOB_BUCKET

# Train model on Vertex
model = model.fit(X_train, y_train)

Uptrain the pretrained model

Next, get the registered model from the Vertex AI Model Registry. Then request the pretrained version of the model.

In [ ]:

registered_model = vertexai.preview.register(model)

pulled_model = vertexai.preview.from_pretrained(
    model_name=registered_model.resource_name
)

Now train the model remotely via Vertex AI Training.

In [ ]:

pulled_model.fit(X_retrain, y_retrain)

Local evaluation

Next, evaluate the pretrained and uptrained versions of the model, and compare the results.

In [ ]:

# Switch to local mode for testing
vertexai.preview.init(remote=False)

# Evaluate model's accuracy score
print(f"Train accuracy: {model.score(X_train, y_train)}")
print(f"Test accuracy: {model.score(X_test, y_test)}")

# Evaluate uptrained model's accuracy score
print(f"Train accuracy: {pulled_model.score(X_train, y_train)}")
print(f"Test accuracy: {pulled_model.score(X_test, y_test)}")

Remote prediction

Finally, you do a remote prediction as a local prediction.

In [ ]:

# Remote mode for prediction
vertexai.preview.init(remote=True)

predictions = model.predict(X_test)
print(predictions)

Delete the registered model

You can delete the registered model in the Vertex AI Model Registry with the delete() method.

In [ ]:

registered_model.delete()

PyTorch

Remote training with CPU

First, train a PyTorch model as a remote training job:

• Reinitialize Vertex AI for remote training.
• Set TorchLogisticRegression for the remote training job.
• Invoke TorchLogisticRegression locally which will launch the remote training job.

In [ ]:

# Switch to remote mode for training
vertexai.preview.init(remote=True)

import torch
from vertexai.preview import VertexModel


# define the custom model
class TorchLogisticRegression(VertexModel, torch.nn.Module):
    def __init__(self, input_size: int, output_size: int):
        torch.nn.Module.__init__(self)
        VertexModel.__init__(self)
        self.linear = torch.nn.Linear(input_size, output_size)
        self.softmax = torch.nn.Softmax(dim=1)

    def forward(self, x):
        return self.softmax(self.linear(x))

    @vertexai.preview.developer.mark.train()
    def train(self, X, y, num_epochs, lr):
        X, y = torch.tensor(X).to(torch.float32), torch.tensor(y)
        dataloader = torch.utils.data.DataLoader(
            torch.utils.data.TensorDataset(X, y),
            batch_size=10,
            shuffle=True,
            generator=torch.Generator(device=X.device),
        )

        criterion = torch.nn.CrossEntropyLoss()
        optimizer = torch.optim.SGD(self.parameters(), lr=lr)

        for t in range(num_epochs):
            for batch, (X, y) in enumerate(dataloader):
                optimizer.zero_grad()
                pred = self(X)
                loss = criterion(pred, y)
                loss.backward()
                optimizer.step()

    @vertexai.preview.developer.mark.predict()
    def predict(self, X):
        X = torch.tensor(X).to(torch.float32)
        with torch.no_grad():
            pred = torch.argmax(self(X), dim=1)
        return pred


# Instantiate model
model = TorchLogisticRegression(4, 3)

# Set training config
model.train.vertex.remote_config.display_name = (
    REMOTE_JOB_NAME + "-pytorch-custom-model"
)
model.train.vertex.remote_config.staging_bucket = REMOTE_JOB_BUCKET

# Train model on Vertex
# Currently update trained model in place hasn't been implemented, so need to get the return value
model.train(X_train, y_train, num_epochs=100, lr=0.05)

Uptrain the pretrained model with GPU

Next, get the registered model from the Vertex AI Model Registry. Then request the pretrained version of the model.

In [ ]:

registered_model = vertexai.preview.register(model)

pulled_model = vertexai.preview.from_pretrained(
    model_name=registered_model.resource_name
)

Now train the model remotely via Vertex AI Training.

In [ ]:

# Get Python version
py_v = ! python3 --version
py_v = py_v[0].split(".")[1]

# only supported by Pythn 3.10+
if int(py_v) >= 10:
    pulled_model.train.vertex.remote_config.enable_cuda = True
    pulled_model.train.vertex.remote_config.display_name = (
        REMOTE_JOB_NAME + "-pytorch-custom-model-gpu"
    )
    pulled_model.train(X_retrain, y_retrain, num_epochs=100, lr=0.05)

Local evaluation

Next, evaluate the pretrained and uptrained versions of the model, and compare the results.

In [ ]:

from sklearn.metrics import accuracy_score

# Switch to local mode for testing
vertexai.preview.init(remote=False)

# Evaluate model's accuracy score
print(f"Train accuracy: {accuracy_score(y_train, model.predict(X_train))}")
print(f"Test accuracy: {accuracy_score(y_test, model.predict(X_test))}")

# Evaluate uptrained model's accuracy score
# only supported by Pythn 3.10+
if py_v >= "3.10":
    print(f"Train accuracy: {accuracy_score(y_train, pulled_model.predict(X_train))}")
    print(f"Test accuracy: {accuracy_score(y_test, pulled_model.predict(X_test))}")

Remote prediction

Finally, you do a remote prediction as a local prediction.

In [ ]:

# Remote mode for prediction
vertexai.preview.init(remote=True)

predictions = model.predict(X_test)
print(predictions)

Delete the registered model

You can delete the registered model in the Vertex AI Model Registry with the delete() method.

In [ ]:

registered_model.delete()

TensorFlow

Remote training with GPU

First, train a TensorFlow model as a remote training job:

• Reinitialize Vertex AI for remote training.
• Set Sequential for the remote training job.
• Invoke Sequential locally which will launch the remote training job.

In [ ]:

# Switch to remote mode for training
vertexai.preview.init(remote=True)

from tensorflow import keras

# Wrap classes to enable Vertex remote execution
# Don't need this step after import hook is implemented
keras.Sequential = vertexai.preview.remote(keras.Sequential)

# Instantiate model
model = keras.Sequential(
    [keras.layers.Dense(5, input_shape=(4,)), keras.layers.Softmax()]
)

# Specify optimizer and loss function
model.compile(optimizer="adam", loss="mean_squared_error")

# Set training config
model.fit.vertex.remote_config.enable_cuda = True
model.fit.vertex.remote_config.display_name = REMOTE_JOB_NAME + "-keras-model-gpu"
model.fit.vertex.remote_config.staging_bucket = REMOTE_JOB_BUCKET

# TODO: Remove
# Manually set compute resources this time
# model.fit.vertex.remote_config.machine_type = "n1-highmem-4"
# model.fit.vertex.remote_config.accelerator_type = "NVIDIA_TESLA_K80"
# model.fit.vertex.remote_config.accelerator_count = 4


# Train model on Vertex
# Currently update trained model in place hasn't been implemented, so need to get the return value
model.fit(X_train, y_train, epochs=10, batch_size=32)

Uptrain the pretrained model with autologging feature

Next, get the registered model from the Vertex AI Model Registry. Then request the pretrained version of the model.

In [ ]:

registered_model = vertexai.preview.register(model)

pulled_model = vertexai.preview.from_pretrained(
    model_name=registered_model.resource_name
)

Now train the model remotely via Vertex AI Training.

In [ ]:

# Config experiment and turn on autologging
# Config experiment and turn on autologging
vertexai.init(
    project=PROJECT_ID,
    location=REGION,
    staging_bucket=BUCKET_URI,
    experiment="test-remote-training-autologging",
)
vertexai.preview.init(remote=True, autolog=True)

# service account is required since autolog is True
pulled_model.fit.vertex.remote_config.service_account = "GCE"

# Set GPU configs to None
pulled_model.fit.vertex.remote_config.enable_cuda = False
pulled_model.fit.vertex.remote_config.machine_type = None
pulled_model.fit.vertex.remote_config.accelerator_type = None
pulled_model.fit.vertex.remote_config.accelerator_count = None

pulled_model.fit.vertex.remote_config.display_name = (
    REMOTE_JOB_NAME + "-keras-model-autologging"
)

pulled_model.fit(X_retrain, y_retrain, epochs=10, batch_size=32)

# TODO InvalidArgument: 400 User-specified resource ID must match the regular expression '[a-z0-9][a-z0-9-]{0,127}'

Get experiments results

Finally, get the Vertex AI Experiments results from the remote training job.

In [ ]:

# View logged metrics & params
vertexai.preview.get_experiment_df()

# Turn off the autologging
vertexai.preview.init(autolog=False)

Local evaluation

Next, evaluate the pretrained and uptrained versions of the model, and compare the results.

In [ ]:

# Switch to local mode for testing
vertexai.preview.init(remote=False)

# Evaluate model's mean square errors
print(f"Train loss: {model.evaluate(X_train, y_train)}")
print(f"Test loss: {model.evaluate(X_test, y_test)}")

# Evaluate uptrained model's mean square errors
print(f"Train loss: {pulled_model.evaluate(X_train, y_train)}")
print(f"Test loss: {pulled_model.evaluate(X_test, y_test)}")

Remote prediction

Finally, you do a remote prediction as a local prediction.

In [ ]:

# Remote mode for prediction
vertexai.preview.init(remote=True)

predictions = model.predict(X_test)
print(predictions)

Delete the registered model

You can delete the registered model in the Vertex AI Model Registry with the delete() method.

In [ ]:

registered_model.delete()

PyTorch Lightning

Remote training

First, train a PyTorch Ligtning model as a remote training job:

• Reinitialize Vertex AI for remote training.
• Set Trainer.pl for the remote training job.
• Invoke Trainer.pl locally which will launch the remote training job.

In [ ]:

# Switch to local mode for testing
vertexai.preview.init(remote=True, autolog=True)

import lightning.pytorch as pl
import torch

# Wrap classes to enable Vertex remote execution
# Don't need this step after import hook is implemented
pl.Trainer = vertexai.preview.remote(pl.Trainer)


# prepare data loader
train_loader = torch.utils.data.DataLoader(
    torch.utils.data.TensorDataset(
        torch.tensor(X_train).to(torch.float32),
        torch.tensor(y_train),
    ),
    batch_size=10,
    shuffle=True,
)


# define the model
class LitLogisticRegression(pl.LightningModule):
    def __init__(self, input_size: int, output_size: int):
        super().__init__()
        self.linear = torch.nn.Linear(input_size, output_size)
        self.softmax = torch.nn.Softmax(dim=1)

    def forward(self, x):
        return self.softmax(self.linear(x))

    def training_step(self, batch, batch_idx):
        x, y = batch
        y_hat = self(x)
        loss = torch.nn.functional.cross_entropy(y_hat, y)
        return loss

    def configure_optimizers(self):
        return torch.optim.Adam(self.parameters(), lr=0.05)

    def predict(self, X):
        X = torch.tensor(X).to(torch.float32)
        with torch.no_grad():
            pred = torch.argmax(self(X), dim=1)
        return pred


model = LitLogisticRegression(4, 3)

# set up the trainer and training config
trainer = pl.Trainer(max_epochs=100, accelerator="cpu")

trainer.fit.vertex.remote_config.display_name = REMOTE_JOB_NAME + "-lightning-model"
trainer.fit.vertex.remote_config.staging_bucket = REMOTE_JOB_BUCKET
trainer.fit.vertex.remote_config.service_account = "GCE"

# Train model on Vertex
trainer.fit(model, train_dataloaders=train_loader)

Local evaluation

Next, evaluate the pretrained version of the model, and compare the results.

In [ ]:

from sklearn.metrics import accuracy_score

# Switch to local mode for testing
vertexai.preview.init(remote=False)

# Evaluate model's accuracy score
print(f"Train accuracy: {accuracy_score(y_train, model.predict(X_train))}")
print(f"Test accuracy: {accuracy_score(y_test, model.predict(X_test))}")

Remote prediction

Finally, you do a remote prediction as a local prediction.

In [ ]:

# Remote mode for prediction
vertexai.preview.init(remote=True)

predictions = model.predict(X_test)
print(predictions)

Cleaning up

To clean up all Google Cloud resources used in this project, you can delete the Google Cloud project you used for the tutorial.

Otherwise, you can delete the individual resources you created in this tutorial.

In [ ]:

import os

from google.cloud import aiplatform

delete_bucket = False

if delete_bucket or os.getenv("IS_TESTING"):
    ! gsutil rm -rf {BUCKET_URI}

try:
    experiment_name = vertexai.preview.get_experiment_df()["experiment_name"]
    aiplatform.Experiment(experiment_name[0]).delete()
except Exception as e:
    print(e)