model_garden/notebooks/official/training/hyperparameter_tuning_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.

Vertex AI Hyperparameter Tuning for XGBoost

Open in Colab

Open in Colab Enterprise

Open in Workbench

View on GitHub

Overview

This tutorial demonstrates how to Vertex AI hyperparameter tuning with XGBoost.

Learn more about Vertex AI hyperparameter tuning.

Objective

In this tutorial, you learn how to use the Vertex AI hyperparameter tuning service for training an XGBoost model.

This tutorial uses the following Vertex AI services:

• Vertex AI training
• Vertex AI hyperparameter tuning (uses Vertex AI Vizier)

The steps performed include:

• Train using a Python training application package.
• Report accuracy during hyperparameter tuning.
• Save the model artifacts to Cloud Storage using GCSFuse.
• List the best model.

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

In [ ]:

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

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

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 [ ]:

! gcloud storage buckets create --location={LOCATION} --project={PROJECT_ID} {BUCKET_URI}

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.

In [ ]:

from google.cloud import aiplatform

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

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 Telsa T4 GPUs allocated to each VM, you would specify:

(aiplatform.gapic.AcceleratorType.NVIDIA_TESLA_T4, 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 version 2.3 for GPU support are known to fail while loading the custom model in this tutorial. The issue is fixed in TF versions 2.3 and above. 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

Set the variable TRAIN_COMPUTE to configure the compute resources for VMs that you use for training. Learn more about the machine types supported for training.

In [ ]:

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

Python training application package

In this example, you use Vertex AI hyperparameter tuning service with a training job that executes a Python training application package.

Learn more about hyperparameter tuning in Vertex AI.

Take a look at how a Python package can be structured for running a custom training job in Vertex AI. The package contains the following directory structure:

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

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

The file trainer/task.py is the Python script that is executed when you run the custom training job.

Create a folder structure as Python package

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 a training script

Next, you create task.py script for your training job. Some noteable steps include:

1. Handling command-line arguments:
   • model-dir: The location to save the trained model. If no value is passed, the location to save the model is obtained from the environment variable AIP_MODEL_DIR that defaults to the staging bucket location.
   • 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
2. Data preprocessing (get_data()):
   • Download the dataset and split into training and test.
3. Training (train_model()):
   • Trains the model
4. Evaluation (evaluate_model()):
   • Evaluates the model.
   • If hyperparameter tuning, reports the metric for accuracy.
5. Saving model artifacts:
   • Saves the model artifacts and evaluation metrics where the Cloud Storage location specified by model-dir.

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

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. Hence, the need to divert it to stdout.
    subprocess.check_call(['gcloud', 'storage', 'cp', args.dataset_data_url, 'data.csv'], stderr=sys.stdout)
    subprocess.check_call(['gcloud', 'storage', '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
    model = xgb.train({}, 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 = [round(value) for value in pred]
    # evaluate predictions
    accuracy = accuracy_score(test_labels, predictions)
    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 Cloud Storage bucket

Compress the whole training folder as a tar ball and then store it in a Cloud Storage bucket.

In [ ]:

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

Define machine configuration

Define the machine configuration for your custom hyperparameter tuning job. This tells Vertex AI what type of machine instance to provision when the job gets started.

You can specify the following parameters:

• machine_type: The type of GCP instance to provision -- e.g., n1-standard-8.
• accelerator_type: The type, if any, of hardware accelerator. In this tutorial if you previously set the variable TRAIN_GPU, you're using a GPU. Otherwise you're using a CPU.
• accelerator_count: Number of accelerators.

In [ ]:

if TRAIN_GPU:
    machine_spec = {
        "machine_type": TRAIN_COMPUTE,
        "accelerator_type": TRAIN_GPU,
        "accelerator_count": TRAIN_NGPU,
    }
else:
    machine_spec = {"machine_type": TRAIN_COMPUTE, "accelerator_count": 0}

Define disk configuration (optional)

Optionally, define the disk configuration for your custom hyperparameter tuning job. This tells Vertex AI what type and size of disk to provision in each machine instance for the hyperparameter tuning.

You can specify the following parameters:

• boot_disk_type: Either SSD or Standard. SSD is faster, and Standard is less expensive. Defaults to SSD.
• boot_disk_size_gb: Size of disk in GB.

In [ ]:

DISK_TYPE = "pd-ssd"  # [ pd-ssd, pd-standard]
DISK_SIZE = 100  # GB

disk_spec = {"boot_disk_type": DISK_TYPE, "boot_disk_size_gb": DISK_SIZE}

Set worker pool specs

Specify the following worker pool specs 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: Specifies where to store the model artifacts in the Cloud Storage bucket.
  • --dataset-data-url: The location of the training data to download.
  • --dataset-labels-url: The location of the training labels to download.
  • --boost-rounds: Sets the tunable hyperparameter num_boost_round while training the XGBoost model.

In [ ]:

# Set path to save model
MODEL_DIR = "{}/aiplatform-custom-job".format(BUCKET_URI)
# Set the source path to the dataset
DATASET_DIR = "gs://cloud-samples-data/ai-platform/iris"

# Set the command-line arguments
CMDARGS = [
    "--dataset-data-url=" + DATASET_DIR + "/iris_data.csv",
    "--dataset-labels-url=" + DATASET_DIR + "/iris_target.csv",
]

# Set the worker pool specs
worker_pool_spec = [
    {
        "replica_count": 1,
        "machine_spec": machine_spec,
        "disk_spec": disk_spec,
        "python_package_spec": {
            "executor_image_uri": TRAIN_IMAGE,
            "package_uris": [BUCKET_URI + "/trainer_iris.tar.gz"],
            "python_module": "trainer.task",
            "args": CMDARGS,
        },
    }
]

Create a custom training job

Use the CustomJob class to create a custom training job with the following parameters:

• display_name: A human readable name for the custom job.
• worker_pool_specs: The specification for the corresponding VM instances.
• base_output_dir: The Cloud Storage location for storing the model artifacts.

In [ ]:

job = aiplatform.CustomJob(
    display_name="iris",
    worker_pool_specs=worker_pool_spec,
    base_output_dir=MODEL_DIR,
)

Create a hyperparameter tuning job

Use the HyperparameterTuningJob class to create a hyperparameter tuning job with the following parameters:

• display_name: A human readable name for the custom job.
• custom_job: The CustomJob object created for training. The worker pool spec from this custom job applies to the CustomJobs created in all the trials.
• metrics_spec: The metrics to optimize. The dictionary key is the metric_id, which is reported by your training job, and the dictionary value is the optimization goal of the metric('minimize' or 'maximize').
• parameter_spec: The parameters to optimize. The dictionary key is the metric_id, which is passed into your training job as a command line key word argument, and the dictionary value is the parameter specification of the metric.
• search_algorithm: The search algorithm to use. Takes grid, random and None as values. Hyperparameter tuning for custom training uses Vertex AI Vizier for training jobs.
• max_trial_count: The maximum number of trials to perform.

In [ ]:

from google.cloud.aiplatform import hyperparameter_tuning as hpt

hpt_job = aiplatform.HyperparameterTuningJob(
    display_name="iris",
    custom_job=job,
    metric_spec={
        "accuracy": "maximize",
    },
    parameter_spec={
        "boost-rounds": hpt.IntegerParameterSpec(min=10, max=100, scale="linear"),
    },
    search_algorithm=None,
    max_trial_count=6,
    parallel_trial_count=1,
)

Run the hyperparameter tuning job

Use the run() method to execute the hyperparameter tuning job.

In [ ]:

hpt_job.run()

Display the hyperparameter tuning job trial results

Once the hyperparameter tuning job successfully finishes, you can access the results from each trial using the trials property.

In [ ]:

print(hpt_job.trials)

Fetch the best trial

Identify the best trial and print the details.

In [ ]:

# Initialize a tuple to identify the best configuration
best = (None, None, None, 0.0)
# Iterate through the trails and update the best configuration
for trial in hpt_job.trials:
    # Keep track of the best outcome
    if float(trial.final_measurement.metrics[0].value) > best[3]:
        try:
            best = (
                trial.id,
                float(trial.parameters[0].value),
                float(trial.parameters[1].value),
                float(trial.final_measurement.metrics[0].value),
            )
        except:
            best = (
                trial.id,
                float(trial.parameters[0].value),
                None,
                float(trial.final_measurement.metrics[0].value),
            )

# print details of the best configuration
print(best)

List the best model

The model artifacts for the best model are saved at:

MODEL_DIR/<best_trial_id>/model

In [ ]:

# Fetch the best model
BEST_MODEL_DIR = MODEL_DIR + "/" + best[0] + "/model"

! gcloud storage ls {BEST_MODEL_DIR}

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:

• HyperparameterTuning Job
• Cloud Storage bucket
• Locally generated files

In [ ]:

# Delete the hyperparameter tuning job
hpt_job.delete()

# Delete the Cloud Storage bucket
delete_bucket = False  # Set True to delete the bucket

if delete_bucket:
    ! gcloud storage rm --recursive $BUCKET_URI

# Delete the locally generated files
! rm -rf custom/
! rm custom.tar.gz