model_garden/notebooks/official/migration/sdk-hyperparameter-tuning.ipynb

In [ ]:

# Copyright 2021 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 Migration: Hyperparameter Tuning

Run in Colab

View on GitHub

Open in Vertex AI Workbench

Overview

This tutorial demonstrates how to use the Vertex AI SDK for Python to tune hyperparameters in a custom tabular classification TensorFlow model.

Learn more about Migrate to Vertex AI and Custom training.

Objective

In this tutorial, you learn to use Vertex AI Hyperparameter to create and tune a custom trained model.

You learn how to create and tune a custom-trained model from a Python script in a Docker container using the Vertex AI SDK for Python.

This tutorial uses the following Google Cloud ML services:

• Vertex AI Training
• Vertex AI Hyperparameter Tuning

The steps performed include:

• Create a Vertex AI hyperparameter tuning job for training a TensorFlow model.

Dataset

The dataset used for this tutorial is the Boston Housing Prices dataset. The version of the dataset you use in this tutorial is built into TensorFlow. The trained model predicts the median price of a house in units of 1K USD.

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.

Installation

Install the packages required for executing this notebook.

In [ ]:

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

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"  # @param {type: "string"}

UUID

If you are in a live tutorial session, you might be using a shared test account or project. To avoid name collisions between users on resources created, you create a uuid for each instance session, and append it onto the name of resources you create in this tutorial.

In [ ]:

import random
import string


# Generate a uuid of a specifed length(default=8)
def generate_uuid(length: int = 8) -> str:
    return "".join(random.choices(string.ascii_lowercase + string.digits, k=length))


UUID = generate_uuid()

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} {BUCKET_URI}

Set up variables

Next, set up some variables used throughout the tutorial.

Import libraries and define constants

In [ ]:

import google.cloud.aiplatform as aip

Initialize Vertex SDK for Python

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

In [ ]:

aip.init(project=PROJECT_ID, 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 K80 GPUs allocated to each VM, you would specify:

(aip.gapic.AcceleratorType.NVIDIA_TESLA_K80, 4)

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

Learn more here hardware accelerator support for your region

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

TF = "2-5"

if TRAIN_GPU:
    TRAIN_VERSION = "tf-gpu.{}".format(TF)
else:
    TRAIN_VERSION = "tf-cpu.{}".format(TF)
if DEPLOY_GPU:
    DEPLOY_VERSION = "tf2-gpu.{}".format(TF)
else:
    DEPLOY_VERSION = "tf2-cpu.{}".format(TF)

TRAIN_IMAGE = "us-docker.pkg.dev/vertex-ai/training/{}:latest".format(TRAIN_VERSION)
DEPLOY_IMAGE = "us-docker.pkg.dev/vertex-ai/prediction/{}:latest".format(DEPLOY_VERSION)

print("Training:", TRAIN_IMAGE, TRAIN_GPU, TRAIN_NGPU)
print("Deployment:", DEPLOY_IMAGE, DEPLOY_GPU, DEPLOY_NGPU)

Set machine type

Next, set the machine type to use for training and prediction.

• Set the variables TRAIN_COMPUTE and DEPLOY_COMPUTE to configure the compute resources for the VMs you use for for training and prediction.
• 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 [ ]:

MACHINE_TYPE = "n1-standard-4"

TRAIN_COMPUTE = MACHINE_TYPE
print("Train machine type", TRAIN_COMPUTE)

DEPLOY_COMPUTE = MACHINE_TYPE
print("Deploy machine type", DEPLOY_COMPUTE)

Examine the training package

Package layout

Before you start the training, you 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 we referred to it in the worker pool specification, we replace the directory slash with a dot (trainer.task) and dropped the file suffix (.py).

Package Assembly

In the following cells, you assemble the training 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        'tensorflow_datasets==1.3.0',\n\n    ],\n\n    packages=setuptools.find_packages())"
! echo "$setup_py" > custom/setup.py

pkg_info = "Metadata-Version: 1.0\n\nName: Boston Housing tabular regression\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

Task.py contents

In the next cell, you write the contents of the hyperparameter tuning script task.py. I won't go into detail, it's just there for you to browse. In summary:

• Parse the command line arguments for the hyperparameter settings for the current trial.
• Get the directory where to save the model artifacts from the command line (--model_dir), and if not specified, then from the environment variable AIP_MODEL_DIR.
• Download and preprocess the Boston Housing dataset.
• Build a DNN model.
• The number of units per dense layer and learning rate hyperparameter values are used during the build and compile of the model.
• A definition of a callback HPTCallback which obtains the validation loss at the end of each epoch (on_epoch_end()) and reports it to the hyperparameter tuning service using hpt.report_hyperparameter_tuning_metric().
• Train the model with the fit() method and specify a callback which report the validation loss back to the hyperparameter tuning service.

In [ ]:

%%writefile custom/trainer/task.py
# Custom Training for Boston Housing

import tensorflow_datasets as tfds
import tensorflow as tf
from tensorflow.python.client import device_lib
from hypertune import HyperTune
import numpy as np
import argparse
import os
import sys
tfds.disable_progress_bar()

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('--lr', dest='lr',
                    default=0.001, type=float,
                    help='Learning rate.')
parser.add_argument('--decay', dest='decay',
                    default=0.98, type=float,
                    help='Decay rate')
parser.add_argument('--units', dest='units',
                    default=64, type=int,
                    help='Number of units.')
parser.add_argument('--epochs', dest='epochs',
                    default=20, type=int,
                    help='Number of epochs.')
parser.add_argument('--steps', dest='steps',
                    default=200, type=int,
                    help='Number of steps per epoch.')
parser.add_argument('--param-file', dest='param_file',
                    default='/tmp/param.txt', type=str,
                    help='Output file for parameters')
parser.add_argument('--distribute', dest='distribute', type=str, default='single',
                    help='distributed training strategy')
args = parser.parse_args()

print('Python Version = {}'.format(sys.version))
print('TensorFlow Version = {}'.format(tf.__version__))
print('TF_CONFIG = {}'.format(os.environ.get('TF_CONFIG', 'Not found')))


def make_dataset():

  # Scaling Boston Housing data features
  def scale(feature):
    max = np.max(feature)
    feature = (feature / max).astype(np.float)
    return feature, max

  (x_train, y_train), (x_test, y_test) = tf.keras.datasets.boston_housing.load_data(
    path="boston_housing.npz", test_split=0.2, seed=113
  )
  params = []
  for _ in range(13):
    x_train[_], max = scale(x_train[_])
    x_test[_], _ = scale(x_test[_])
    params.append(max)

  # store the normalization (max) value for each feature
  with tf.io.gfile.GFile(args.param_file, 'w') as f:
    f.write(str(params))
  return (x_train, y_train), (x_test, y_test)

# Build the Keras model
def build_and_compile_dnn_model():
  model = tf.keras.Sequential([
      tf.keras.layers.Dense(args.units, activation='relu', input_shape=(13,)),
      tf.keras.layers.Dense(args.units, activation='relu'),
      tf.keras.layers.Dense(1, activation='linear')
  ])
  model.compile(
      loss='mse',
      optimizer=tf.keras.optimizers.RMSprop(learning_rate=args.lr, decay=args.decay))
  return model


model = build_and_compile_dnn_model()

# Instantiate the HyperTune reporting object
hpt = HyperTune()

# Reporting callback
class HPTCallback(tf.keras.callbacks.Callback):

    def on_epoch_end(self, epoch, logs=None):
        global hpt
        hpt.report_hyperparameter_tuning_metric(
        hyperparameter_metric_tag='val_loss',
        metric_value=logs['val_loss'],
        global_step=epoch)

# Train the model
BATCH_SIZE = 16
(x_train, y_train), (x_test, y_test) = make_dataset()
model.fit(x_train, y_train, epochs=args.epochs, batch_size=BATCH_SIZE, validation_split=0.1, callbacks=[HPTCallback()])
model.save(args.model_dir)

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_boston.tar.gz

Train a model

training.using-hyperparamter-tuning

Prepare your machine specification

Now define the machine specification for your custom training job. This tells Vertex what type of machine instance to provision for the training.

• 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 != None, you are using a GPU; otherwise you use a CPU.
• accelerator_count: The 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}

Prepare your disk specification

(optional) Now define the disk specification for your custom training job. This tells Vertex what type and size of disk to provision in each machine instance for the training.

• 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 = 200  # GB

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

Define the worker pool specification

Next, you define the worker pool specification for your custom training job. The worker pool specification consist of the following:

• replica_count: The number of instances to provision of this machine type.

• machine_spec: The hardware specification.

• disk_spec : (optional) The disk storage specification.

• python_package: The Python training package to install on the VM instance(s) and which Python module to invoke, along with command line arguments for the Python module.

Let's dive deeper now into the python package specification:

-executor_image_spec: This is the docker image which is configured for your custom training job.

-package_uris: This is a list of the locations (URIs) of your python training packages to install on the provisioned instance. The locations need to be in a Cloud Storage bucket. These can be either individual python files or a zip (archive) of an entire package. In the later case, the job service unzip (unarchive) the contents into the docker image.

-python_module: The Python module (script) to invoke for running the custom training job. In this example, you be invoking trainer.task.py -- note that it was not neccessary to append the .py suffix.

-args: The command line arguments to pass to the corresponding Pythom module. In this example, you be setting:

• "--model-dir=" + MODEL_DIR : The Cloud Storage location where to store the model artifacts. There are two ways to tell the training script where to save 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.
• "--epochs=" + EPOCHS: The number of epochs for training.
• "--steps=" + STEPS: The number of steps (batches) per epoch.
• "--distribute=" + TRAIN_STRATEGY" : The training distribution strategy to use for single or distributed training.
  • "single": single device.
  • "mirror": all GPU devices on a single compute instance.
  • "multi": all GPU devices on all compute instances.

In [ ]:

JOB_NAME = "custom_job_" + UUID
MODEL_DIR = "{}/{}".format(BUCKET_URI, JOB_NAME)

if not TRAIN_NGPU or TRAIN_NGPU < 2:
    TRAIN_STRATEGY = "single"
else:
    TRAIN_STRATEGY = "mirror"

EPOCHS = 20
STEPS = 100

DIRECT = True
if DIRECT:
    CMDARGS = [
        "--model-dir=" + MODEL_DIR,
        "--epochs=" + str(EPOCHS),
        "--steps=" + str(STEPS),
        "--distribute=" + TRAIN_STRATEGY,
    ]
else:
    CMDARGS = [
        "--epochs=" + str(EPOCHS),
        "--steps=" + str(STEPS),
        "--distribute=" + TRAIN_STRATEGY,
    ]

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_boston.tar.gz"],
            "python_module": "trainer.task",
            "args": CMDARGS,
        },
    }
]

training.create-custom-job

Create a custom job

Use the class CustomJob to create a custom job, such as for hyperparameter tuning, with the following parameters:

• display_name: A human readable name for the custom job.
• worker_pool_specs: The specification for the corresponding VM instances.

In [ ]:

job = aip.CustomJob(display_name="boston_" + UUID, worker_pool_specs=worker_pool_spec)

# print(job)

Create a hyperparameter tuning job

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

• display_name: A human readable name for the custom job.
• custom_job: 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.

In [ ]:

from google.cloud.aiplatform import hyperparameter_tuning as hpt

hpt_job = aip.HyperparameterTuningJob(
    display_name="boston_" + UUID,
    custom_job=job,
    metric_spec={
        "val_loss": "minimize",
    },
    parameter_spec={
        "lr": hpt.DoubleParameterSpec(min=0.001, max=0.1, scale="log"),
        "units": hpt.IntegerParameterSpec(min=4, max=128, scale="linear"),
    },
    max_trial_count=6,
    parallel_trial_count=1,
)

# print(hpt_job)

Run the hyperparameter tuning job

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

In [ ]:

hpt_job.run()

Example output:

INFO:google.cloud.aiplatform.jobs:Creating HyperparameterTuningJob
INFO:google.cloud.aiplatform.jobs:HyperparameterTuningJob created. Resource name: projects/759209241365/locations/us-central1/hyperparameterTuningJobs/760969798560514048
INFO:google.cloud.aiplatform.jobs:To use this HyperparameterTuningJob in another session:
INFO:google.cloud.aiplatform.jobs:hpt_job = aiplatform.HyperparameterTuningJob.get('projects/759209241365/locations/us-central1/hyperparameterTuningJobs/760969798560514048')
INFO:google.cloud.aiplatform.jobs:View HyperparameterTuningJob:
https://console.cloud.google.com/ai/platform/locations/us-central1/training/760969798560514048?project=759209241365
INFO:google.cloud.aiplatform.jobs:HyperparameterTuningJob projects/759209241365/locations/us-central1/hyperparameterTuningJobs/760969798560514048 current state:
JobState.JOB_STATE_RUNNING
INFO:google.cloud.aiplatform.jobs:HyperparameterTuningJob projects/759209241365/locations/us-central1/hyperparameterTuningJobs/760969798560514048 current state:
JobState.JOB_STATE_RUNNING
...
INFO:google.cloud.aiplatform.jobs:HyperparameterTuningJob projects/759209241365/locations/us-central1/hyperparameterTuningJobs/760969798560514048 current state:
JobState.JOB_STATE_SUCCEEDED
INFO:google.cloud.aiplatform.jobs:HyperparameterTuningJob run completed. Resource name: projects/759209241365/locations/us-central1/hyperparameterTuningJobs/760969798560514048

Display the hyperparameter tuning job trial results

After the hyperparameter tuning job has completed, the property trials return the results for each trial.

In [ ]:

print(hpt_job.trials)

Example output:

[id: "1"
state: SUCCEEDED
parameters {
  parameter_id: "lr"
  value {
    number_value: 0.010000000000000002
  }
}
parameters {
  parameter_id: "units"
  value {
    number_value: 66.0
  }
}
final_measurement {
  step_count: 19
  metrics {
    metric_id: "val_loss"
    value: 24.61802691948123
  }
}
start_time {
  seconds: 1629414344
  nanos: 232151761
}
end_time {
  seconds: 1629414825
}
, id: "2"
state: SUCCEEDED
parameters {
  parameter_id: "lr"
  value {
    number_value: 0.0036327633937958217
  }
}
parameters {
  parameter_id: "units"
  value {
    number_value: 94.0
  }
}
final_measurement {
  step_count: 19
  metrics {
    metric_id: "val_loss"
    value: 25.499705989186356
  }
}
start_time {
  seconds: 1629414960
  nanos: 618333580
}
end_time {
  seconds: 1629415021
}
...
, id: "6"
state: SUCCEEDED
parameters {
  parameter_id: "lr"
  value {
    number_value: 0.002775175422799751
  }
}
parameters {
  parameter_id: "units"
  value {
    number_value: 104.0
  }
}
final_measurement {
  step_count: 7
  metrics {
    metric_id: "val_loss"
    value: 24.655450123112377
  }
}
start_time {
  seconds: 1629415735
  nanos: 998711852
}
end_time {
  seconds: 1629415766
}
]

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

delete_bucket = False

# Delete the training job
try:
    job.delete()
except Exception as e:
    print(e)

# Delete the HPT job using the Vertex batch prediction object
hpt_job.delete()

if delete_bucket or os.getenv("IS_TESTING"):
    ! gsutil rm -r $BUCKET_URI