model_garden/notebooks/official/custom/get_started_vertex_training_xgboost.ipynb

In [ ]:

# Copyright 2022 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.

Get started with Vertex AI Training for XGBoost

Run in Colab

Open in Colab Enterprise

Open in Vertex AI Workbench

View on GitHub

Overview

This tutorial demonstrates how to use Vertex AI Training for XGBoost models.

Learn more about Custom training.

Objective

In this tutorial, you learn how to use Vertex AI Training for training a XGBoost custom model.

This tutorial uses the following Google Cloud ML services:

• Vertex AI Training
• Vertex AI model resource

The steps performed include:

• Training using a Python package.
• Report accuracy when hyperparameter tuning.
• Save the model artifacts to Cloud Storage using GCSFuse.
• Create a Vertex AI model resource.

Dataset

The dataset used for this tutorial is the Iris dataset from TensorFlow Datasets. This dataset does not require any feature engineering. The version of the dataset in this tutorial is stored in a public Cloud Storage bucket. The trained model predicts the type of Iris flower species from a class of three species: setosa, virginica, or versicolor.

Costs

This tutorial uses billable components of Google Cloud:

• Vertex AI
• Cloud Storage

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

Get started

Install Vertex AI SDK for Python and other required packages

In [ ]:

! pip3 install --upgrade google-cloud-aiplatform  --quiet

Restart runtime (Colab only)

To use the newly installed packages, you must restart the runtime on Google Colab.

In [ ]:

import sys

if "google.colab" in sys.modules:

    import IPython

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

⚠️ The kernel is going to restart. Wait until it's finished before continuing to the next step. ⚠️

Authenticate your notebook environment (Colab only)

Authenticate your environment on Google Colab.

In [ ]:

import sys

if "google.colab" in sys.modules:

    from google.colab import auth

    auth.authenticate_user()

Set Google Cloud project information and initialize Vertex AI SDK for Python

To get started using Vertex AI, you must have an existing Google Cloud project and enable the Vertex AI API. Learn more about setting up a project and a development environment.

In [ ]:

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

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"}

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

In [ ]:

! gsutil mb -l $LOCATION $BUCKET_URI

Initialize Vertex AI SDK for Python

In [ ]:

import google.cloud.aiplatform as aiplatform

aiplatform.init(project=PROJECT_ID, location=LOCATION, staging_bucket=BUCKET_URI)

Set up variables

Next, set up some variables used throughout the tutorial.

Set hardware accelerators

You can set hardware accelerators for training and prediction.

Set the variables TRAIN_GPU/TRAIN_NGPU and DEPLOY_GPU/DEPLOY_NGPU to use a container image supporting a GPU and the number of GPUs allocated to the virtual machine (VM) instance. For example, to use a GPU container image with 4 NVIDIA Tesla K80 GPUs allocated to each VM, you would specify:

(aiplatform.gapic..AcceleratorType.NVIDIA_TESLA_K80, 4)

Otherwise specify (None, None) to use a container image to run on a CPU.

Learn more about hardware accelerator support for your region.

Note: TF releases before 2.3 for GPU support is expected to fail to load the custom model in this tutorial. It's a known issue and fixed in TF 2.3. This is caused by static graph ops that are generated in the serving function. If you encounter this issue on your own custom models, use a container image for TF 2.3 with GPU support.

In [ ]:

TRAIN_GPU, TRAIN_NGPU = (None, None)

DEPLOY_GPU, DEPLOY_NGPU = (None, None)

Set pre-built containers

Set the pre-built Docker container image for training and prediction.

For the latest list, see Pre-built containers for training.

For the latest list, see Pre-built containers for prediction.

In [ ]:

TRAIN_VERSION = "xgboost-cpu.1-1"
DEPLOY_VERSION = "xgboost-cpu.1-1"

TRAIN_IMAGE = "{}-docker.pkg.dev/vertex-ai/training/{}:latest".format(
    LOCATION.split("-")[0], TRAIN_VERSION
)
DEPLOY_IMAGE = "{}-docker.pkg.dev/vertex-ai/prediction/{}:latest".format(
    LOCATION.split("-")[0], DEPLOY_VERSION
)

Set machine type

Next, set the machine type to use for training.

• Set the variable TRAIN_COMPUTE to configure the compute resources for the VMs used for training.
• machine type
  • n1-standard: 3.75GB of memory per vCPU.
  • n1-highmem: 6.5GB of memory per vCPU
  • n1-highcpu: 0.9 GB of memory per vCPU
• vCPUs: number of [2, 4, 8, 16, 32, 64, 96 ]

Note: The following is not supported for training:

• standard: 2 vCPUs
• highcpu: 2, 4 and 8 vCPUs

Note: You may also use n2 and e2 machine types for training and deployment, but they do not support GPUs.

In [ ]:

TRAIN_COMPUTE = "n1-standard-4"
print("Train machine type", TRAIN_COMPUTE)

Introduction to XGBoost training

Once you have trained a XGBoost model, save it at a Cloud Storage location, so it can subsequently be uploaded to a Vertex AI model resource. The XGBoost package does not have support to save the model to a Cloud Storage location. Instead, do the following steps to save to a Cloud Storage location.

1. Save the in-memory model to the local filesystem (e.g., model.bst).
2. Use gsutil to copy the local copy to the specified Cloud Storage location.

Note: You can do hyperparameter tuning with a XGBoost model.

Examine the training package

Package layout

Before you start the training, look at how a Python package is assembled for a custom training job. When unarchived, the package contains the following directory/file layout.

• PKG-INFO
• README.md
• setup.cfg
• setup.py
• trainer
  • __init__.py
  • task.py

The files setup.cfg and setup.py are the instructions for installing the package into the operating environment of the Docker image.

The file trainer/task.py is the Python script for executing the custom training job. Note, when trainer/task.py is referred to in the worker pool specification, the directory slash is replaced with a dot and the file suffix (.py) suffix is dropped: (trainer.task).

Package Assembly

In [ ]:

# Make folder for Python training script
! rm -rf custom
! mkdir custom

# Add package information
! touch custom/README.md

setup_cfg = "[egg_info]\n\ntag_build =\n\ntag_date = 0"
! echo "$setup_cfg" > custom/setup.cfg

setup_py = "import setuptools\n\nsetuptools.setup(\n\n    install_requires=[\n\n        'cloudml-hypertune',\n\n    ],\n\n    packages=setuptools.find_packages())"
! echo "$setup_py" > custom/setup.py

pkg_info = "Metadata-Version: 1.0\n\nName: Iris tabular classification\n\nVersion: 0.0.0\n\nSummary: Demostration training script\n\nHome-page: www.google.com\n\nAuthor: Google\n\nAuthor-email: aferlitsch@google.com\n\nLicense: Public\n\nDescription: Demo\n\nPlatform: Vertex"
! echo "$pkg_info" > custom/PKG-INFO

# Make the training subfolder
! mkdir custom/trainer
! touch custom/trainer/__init__.py

Create the task script for the Python training package

Next, you create the task.py script for driving the training package. Some noteable steps include:

• Command-line arguments:
  • model-dir: The location to save the trained model. When using Vertex AI custom training, the location is specified in the environment variable: AIP_MODEL_DIR,
  • dataset_data_url: The location of the training data to download.
  • dataset_labels_url: The location of the training labels to download.
  • boost-rounds: Tunable hyperparameter
• Data preprocessing (get_data()):
  • Download the dataset and split into training and test.
• Training (train_model()):
  • Trains the model
• Evaluation (evaluate_model()):
  • Evaluates the model.
  • If hyperparameter tuning, reports the metric for accuracy.
• Model artifact saving
  • Saves the model artifacts and evaluation metrics where the Cloud Storage location specified by model-dir.

Note: The training script uses GCSFuse which mounts the Cloud Storage bucket as a network filesystem, allowing the script to perform filesystem operations (e.g., read and write) within a Cloud Storage bucket.

In [ ]:

%%writefile custom/trainer/task.py
import datetime
import os
import subprocess
import sys
import pandas as pd
import xgboost as xgb
import hypertune
import argparse
import logging
import numpy as np

from sklearn.model_selection import train_test_split
from sklearn.metrics import accuracy_score

parser = argparse.ArgumentParser()
parser.add_argument('--model-dir', dest='model_dir',
                    default=os.getenv('AIP_MODEL_DIR'), type=str, help='Model dir.')
parser.add_argument("--dataset-data-url", dest="dataset_data_url",
                    type=str, help="Download url for the training data.")
parser.add_argument("--dataset-labels-url", dest="dataset_labels_url",
                    type=str, help="Download url for the training data labels.")
parser.add_argument("--boost-rounds", dest="boost_rounds",
                    default=20, type=int, help="Number of boosted rounds") 
args = parser.parse_args()

logging.getLogger().setLevel(logging.INFO)

def get_data():
    logging.info("Downloading training data and labelsfrom: {}, {}".format(args.dataset_data_url, args.dataset_labels_url))
    # gsutil outputs everything to stderr so we need to divert it to stdout.
    subprocess.check_call(['gsutil', 'cp', args.dataset_data_url, 'data.csv'], stderr=sys.stdout)
    # gsutil outputs everything to stderr so we need to divert it to stdout.
    subprocess.check_call(['gsutil', 'cp', args.dataset_labels_url, 'labels.csv'], stderr=sys.stdout)


    # Load data into pandas, then use `.values` to get NumPy arrays
    data = pd.read_csv('data.csv').values
    labels = pd.read_csv('labels.csv').values

    # Convert one-column 2D array into 1D array for use with XGBoost
    labels = labels.reshape((labels.size,))

    train_data, test_data, train_labels, test_labels = train_test_split(data, labels, test_size=0.2, random_state=7)

    # Load data into DMatrix object
    dtrain = xgb.DMatrix(train_data, label=train_labels)
    return dtrain, test_data, test_labels

def train_model(dtrain):
    logging.info("Start training ...")
    # Train XGBoost model
    params = {
        'objective': 'multi:softprob',
        'num_class': 3
    }
    model = xgb.train(params, dtrain, num_boost_round=args.boost_rounds)
    logging.info("Training completed")
    return model

def evaluate_model(model, test_data, test_labels):
    dtest = xgb.DMatrix(test_data)
    pred = model.predict(dtest)
    predictions = [np.around(value) for value in pred]
    # evaluate predictions
    try:
        accuracy = accuracy_score(test_labels, predictions)
    except:
        accuracy = 0.0
    logging.info(f"Evaluation completed with model accuracy: {accuracy}")

    # report metric for hyperparameter tuning
    hpt = hypertune.HyperTune()
    hpt.report_hyperparameter_tuning_metric(
        hyperparameter_metric_tag='accuracy',
        metric_value=accuracy
    )
    return accuracy


dtrain, test_data, test_labels = get_data()
model = train_model(dtrain)
accuracy = evaluate_model(model, test_data, test_labels)

# GCSFuse conversion
gs_prefix = 'gs://'
gcsfuse_prefix = '/gcs/'
if args.model_dir.startswith(gs_prefix):
    args.model_dir = args.model_dir.replace(gs_prefix, gcsfuse_prefix)
    dirpath = os.path.split(args.model_dir)[0]
    if not os.path.isdir(dirpath):
        os.makedirs(dirpath)

# Export the classifier to a file
gcs_model_path = os.path.join(args.model_dir, 'model.bst')
logging.info("Saving model artifacts to {}". format(gcs_model_path))
model.save_model(gcs_model_path)

logging.info("Saving metrics to {}/metrics.json". format(args.model_dir))
gcs_metrics_path = os.path.join(args.model_dir, 'metrics.json')
with open(gcs_metrics_path, "w") as f:
    f.write(f"{'accuracy: {accuracy}'}")

Store training script on your Cloud Storage bucket

Next, you package the training folder into a compressed tar ball, and then store it in your Cloud Storage bucket.

In [ ]:

! rm -f custom.tar custom.tar.gz
! tar cvf custom.tar custom
! gzip custom.tar
! gsutil cp custom.tar.gz $BUCKET_URI/trainer_iris.tar.gz

Create and run custom training job

To train a custom model, you perform two steps: 1) create a custom training job, and 2) run the job.

Create custom training job

A custom training job is created with the CustomTrainingJob class, with the following parameters:

• display_name: The human readable name for the custom training job.

• container_uri: The training container image.

• python_package_gcs_uri: The location of the Python training package as a tarball.

• python_module_name: The relative path to the training script in the Python package.

• model_serving_container_uri: The container image for deploying the model.

Note: There is no requirements parameter. You specify any requirements in the setup.py script in your Python package.

In [ ]:

DISPLAY_NAME = "iris"

job = aiplatform.CustomPythonPackageTrainingJob(
    display_name=DISPLAY_NAME,
    python_package_gcs_uri=f"{BUCKET_URI}/trainer_iris.tar.gz",
    python_module_name="trainer.task",
    container_uri=TRAIN_IMAGE,
    model_serving_container_image_uri=DEPLOY_IMAGE,
    project=PROJECT_ID,
)

Prepare your command-line arguments

Now define the command-line arguments for your custom training container:

• args: The command-line arguments to pass to the executable that is set as the entry point into the container.
  • --model-dir : For our demonstrations, we use this command-line argument to specify where to store the model artifacts.
    • direct: You pass the Cloud Storage location as a command line argument to your training script (set variable DIRECT = True), or
    • indirect: The service passes the Cloud Storage location as the environment variable AIP_MODEL_DIR to your training script (set variable DIRECT = False). In this case, you tell the service the model artifact location in the job specification.
  • --dataset-data-url: The location of the training data to download.
  • --dataset-labels-url: The location of the training labels to download.
  • --boost-rounds: Tunable hyperparameter.

In [ ]:

MODEL_DIR = "{}/{}".format(BUCKET_URI, "model")
DATASET_DIR = "gs://cloud-samples-data/ai-platform/iris"

ROUNDS = 20

DIRECT = False
if DIRECT:
    CMDARGS = [
        "--dataset-data-url=" + DATASET_DIR + "/iris_data.csv",
        "--dataset-labels-url=" + DATASET_DIR + "/iris_target.csv",
        "--boost-rounds=" + str(ROUNDS),
        "--model_dir=" + MODEL_DIR,
    ]
else:
    CMDARGS = [
        "--dataset-data-url=" + DATASET_DIR + "/iris_data.csv",
        "--dataset-labels-url=" + DATASET_DIR + "/iris_target.csv",
        "--boost-rounds=" + str(ROUNDS),
    ]

Run the custom training job

Next, you run the custom job to start the training job by invoking the method run, with the following parameters:

• model_display_name: The human readable name for the Model resource.
• args: The command-line arguments to pass to the training script.
• replica_count: The number of compute instances for training (replica_count = 1 is single node training).
• machine_type: The machine type for the compute instances.
• accelerator_type: The hardware accelerator type.
• accelerator_count: The number of accelerators to attach to a worker replica.
• base_output_dir: The Cloud Storage location to write the model artifacts to.
• sync: Whether to block until completion of the job.

In [ ]:

if TRAIN_GPU:
    model = job.run(
        model_display_name="iris",
        args=CMDARGS,
        replica_count=1,
        machine_type=TRAIN_COMPUTE,
        accelerator_type=TRAIN_GPU.name,
        accelerator_count=TRAIN_NGPU,
        base_output_dir=MODEL_DIR,
        sync=False,
    )
else:
    model = job.run(
        model_display_name="iris",
        args=CMDARGS,
        replica_count=1,
        machine_type=TRAIN_COMPUTE,
        base_output_dir=MODEL_DIR,
        sync=False,
    )

model_path_to_deploy = MODEL_DIR

List a custom training job

In [ ]:

_job = job.list(filter=f"display_name={DISPLAY_NAME}")
print(_job)

Wait for completion of custom training job

Next, wait for the custom training job to complete. Alternatively, one can set the parameter sync to True in the run() method to block until the custom training job is completed.

In [ ]:

model.wait()

Delete a custom training job

After a training job is completed, you can delete the training job with the method delete(). Prior to completion, a training job can be canceled with the method cancel().

In [ ]:

job.delete()

Clean 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:

• Custom Job (Custom Training job is removed in the previous step)
• Model
• Cloud Storage Bucket

In [ ]:

delete_bucket = False

model.delete()

if delete_bucket:
    ! gsutil rm -r $BUCKET_URI