model_garden/notebooks/official/custom/custom_training_tensorboard_profiler.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.

Profile model training performance using Cloud Profiler

Open in Colab

Open in Colab Enterprise

Open in Workbench

View on GitHub

Overview

Cloud Profiler lets you monitor and optimize your model training performance by helping you understand the resource consumption of training operations. This tutorial demonstrates how to enable Cloud Profiler so you can debug model training performance for your custom training jobs.

Learn more about Cloud Profiler.

Objective

In this tutorial, you learn how to enable Cloud Profiler for custom training jobs.

This tutorial uses the following Google Cloud AI services:

• Vertex AI Training
• Vertex AI TensorBoard

The steps performed include:

• Setup a service account and a Cloud Storage bucket
• Create a TensorBoard instance
• Create and run a custom training job
• View the Cloud Profiler dashboard

Dataset

The dataset used for this tutorial is the mnist dataset from TensorFlow Datasets.

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 -p $PROJECT_ID $BUCKET_URI

Initialize Vertex AI SDK for Python

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

In [ ]:

from google.cloud import aiplatform

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

Setup service account and permissions

A service account will be used to create custom training jobs. If you do not want to use your project's Compute Engine service account, set SERVICE_ACCOUNT to another service account ID. You can create a service account by following the instructions.

In [ ]:

SERVICE_ACCOUNT = "[your-service-account]"  # @param {type:"string"}

In [ ]:

import sys

IS_COLAB = "google.colab" in sys.modules
if (
    SERVICE_ACCOUNT == ""
    or SERVICE_ACCOUNT is None
    or SERVICE_ACCOUNT == "[your-service-account]"
):
    # Get your service account from gcloud
    if not IS_COLAB:
        shell_output = !gcloud auth list 2>/dev/null
        SERVICE_ACCOUNT = shell_output[2].replace("*", "").strip()

    else:  # IS_COLAB:
        shell_output = ! gcloud projects describe  $PROJECT_ID
        project_number = shell_output[-1].split(":")[1].strip().replace("'", "")
        SERVICE_ACCOUNT = f"{project_number}-compute@developer.gserviceaccount.com"

    print("Service Account:", SERVICE_ACCOUNT)

Set service account access for Vertex AI Pipelines

Run the following commands to grant your service account access to read and write pipeline artifacts in the bucket that you created in the previous step. You only need to run this step once per service account.

In [ ]:

! gsutil iam ch serviceAccount:{SERVICE_ACCOUNT}:roles/storage.objectCreator $BUCKET_URI

! gsutil iam ch serviceAccount:{SERVICE_ACCOUNT}:roles/storage.objectViewer $BUCKET_URI

Enable Artifact Registry API

First, you must enable the Artifact Registry API service for your project.

Learn more about Enabling service.

Set project in colab environment (Colab only)

In [ ]:

if IS_COLAB:
    ! gcloud config set project $PROJECT_ID

In [ ]:

# Enable Artifact Registry API
! gcloud services enable artifactregistry.googleapis.com --quiet

Create a TensorBoard instance

A Vertex AI TensorBoard instance, which is a regionalized resource storing your Vertex AI TensorBoard experiments, must be created before the experiments can be visualized. You can create multiple instances in a project. You can use command gcloud ai tensorboards list to get a list of your existing TensorBoard instances.

Set your TensorBoard instance display name

In [ ]:

TENSORBOARD_NAME = "your-tensorboard-unique"  # @param {type:"string"}

Create a TensorBoard instance

If you don't have a TensorBoard instance, create one by running the following cell:

In [ ]:

tensorboard = aiplatform.Tensorboard.create(
    display_name=TENSORBOARD_NAME, project=PROJECT_ID, location=LOCATION
)

TENSORBOARD_INSTANCE_NAME = tensorboard.resource_name

print("TensorBoard instance name:", TENSORBOARD_INSTANCE_NAME)

Train a model

To train a model using your custom training code, choose one of the following options:

• Prebuilt container: Load your custom training code as a Python package to a prebuilt container image from Google Cloud.

• Custom container: Create your own container image that contains your custom training code.

In this tutorial, we will train a custom model using a custom container.

Create a private Docker repository

Your first step is to create your own Docker repository in Google Artifact Registry.

In [ ]:

import os

DOCKER_REPOSITORY = f"{PROJECT_ID}-repo-unique"

! gcloud services enable artifactregistry.googleapis.com

if os.getenv("IS_TESTING"):
    ! sudo apt-get update --yes && sudo apt-get --only-upgrade --yes install google-cloud-sdk-cloud-run-proxy google-cloud-sdk-harbourbridge google-cloud-sdk-cbt google-cloud-sdk-gke-gcloud-auth-plugin google-cloud-sdk-kpt google-cloud-sdk-local-extract google-cloud-sdk-minikube google-cloud-sdk-app-engine-java google-cloud-sdk-app-engine-go google-cloud-sdk-app-engine-python google-cloud-sdk-spanner-emulator google-cloud-sdk-bigtable-emulator google-cloud-sdk-nomos google-cloud-sdk-package-go-module google-cloud-sdk-firestore-emulator kubectl google-cloud-sdk-datastore-emulator google-cloud-sdk-app-engine-python-extras google-cloud-sdk-cloud-build-local google-cloud-sdk-kubectl-oidc google-cloud-sdk-anthos-auth google-cloud-sdk-app-engine-grpc google-cloud-sdk-pubsub-emulator google-cloud-sdk-datalab google-cloud-sdk-skaffold google-cloud-sdk google-cloud-sdk-terraform-tools google-cloud-sdk-config-connector
    ! gcloud components update --quiet

! gcloud artifacts repositories create {DOCKER_REPOSITORY} \
    --repository-format=docker \
    --location={LOCATION} \
    --description="Repository for TensorBoard Custom Training Job" \
    --quiet

! gcloud artifacts repositories list

Configure authentication to your private Docker repository

Before you push or pull container images, configure Docker to use the gcloud command-line tool to authenticate requests to Artifact Registry for your location.

In [ ]:

import sys

IS_COLAB = "google.colab" in sys.modules

if not IS_COLAB:
    ! gcloud auth configure-docker {LOCATION}-docker.pkg.dev --quiet

Create a custom container image and push to your private Docker repository

First, you create a training script file and a docker file.

Create a directory for all of your training code.

In [ ]:

PYTHON_PACKAGE_APPLICATION_DIR = "trainer"

!mkdir -p $PYTHON_PACKAGE_APPLICATION_DIR

Prepare the training script

Your training code must be configured to write TensorBoard logs to a Cloud Storage bucket, the location of which Vertex AI Training automatically makes available through a predefined environment variable, AIP_TENSORBOARD_LOG_DIR.

This can usually be done by providing os.environ['AIP_TENSORBOARD_LOG_DIR'] as the log directory to the open source TensorBoard log writing APIs.

For example, in TensorFlow 2.x, you can use following code to create a tensorboard_callback:

tensorboard_callback = tf.keras.callbacks.TensorBoard( 
  log_dir=os.environ['AIP_TENSORBOARD_LOG_DIR'], 
  histogram_freq=1) 

AIP_TENSORBOARD_LOG_DIR is in the BASE_OUTPUT_DIR that you provide when creating the custom training job.

To enable Cloud Profiler for your training job, add the following to your training script:

Add the cloud_profiler import at your top level imports:

from google.cloud.aiplatform.training_utils import cloud_profiler

Initialize the cloud_profiler plugin by adding:

cloud_profiler.init()

In [ ]:

%%writefile trainer/task.py

import tensorflow as tf
import argparse
import os
import sys, traceback
from google.cloud.aiplatform.training_utils import cloud_profiler

"""Train an mnist model and use cloud_profiler for profiling."""

def _create_model():
    model = tf.keras.models.Sequential(
        [
            tf.keras.layers.Flatten(input_shape=(28, 28)),
            tf.keras.layers.Dense(128, activation="relu"),
            tf.keras.layers.Dropout(0.2),
            tf.keras.layers.Dense(10),
        ]
    )
    return model


def main(args):
    print('Loading and preprocessing data ...')
    mnist = tf.keras.datasets.mnist

    (x_train, y_train), (x_test, y_test) = mnist.load_data()
    x_train, x_test = x_train / 255.0, x_test / 255.0

    print('Creating and training model ...')

    model = _create_model()
    model.compile(
      optimizer="adam",
      loss=tf.keras.losses.sparse_categorical_crossentropy,
      metrics=["accuracy"],
    )

    # Initialize the profiler.
    print('Initialize the profiler ...')
        
    try:
        cloud_profiler.init()
    except:
        ex_type, ex_value, ex_traceback = sys.exc_info()
        print("*** Unexpected:", ex_type.__name__, ex_value)
        traceback.print_tb(ex_traceback, limit=10, file=sys.stdout)
    
    print('The profiler initiated.')

    log_dir = "logs"
    if 'AIP_TENSORBOARD_LOG_DIR' in os.environ:
      log_dir = os.environ['AIP_TENSORBOARD_LOG_DIR']

    print('Setting up the TensorBoard callback ...')
    tensorboard_callback = tf.keras.callbacks.TensorBoard(
        log_dir=log_dir,
        histogram_freq=1)

    print('Training model ...')
    model.fit(
        x_train,
        y_train,
        epochs=args.epochs,
        verbose=0,
        callbacks=[tensorboard_callback],
    )
    print('Training completed.')

    print('Saving model ...')

    model_dir = "model"
    if 'AIP_MODEL_DIR' in os.environ:
      model_dir = os.environ['AIP_MODEL_DIR']
    tf.saved_model.save(model, model_dir)

    print('Model saved at ' + model_dir)


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--epochs", type=int, default=100, help="Number of epochs to run model."
    )
    
    args = parser.parse_args()
    main(args)

Prepare the Dockerfile

In [ ]:

%%writefile Dockerfile
# Specifies base image and tag
FROM us-docker.pkg.dev/vertex-ai/training/tf-cpu.2-9:latest
WORKDIR /root

# Installs additional packages as you need.
RUN pip3 install google-cloud-aiplatform[cloud_profiler]
RUN pip3 install protobuf==3.20.2

# Copies the trainer code to the docker image.
RUN mkdir /root/trainer
COPY trainer/task.py /root/trainer/task.py

# Sets up the entry point to invoke the trainer.
ENTRYPOINT ["python", "-m", "trainer.task"]

Build a custom container image and push to your private Docker repository

In [ ]:

IMAGE_NAME = "tensorboard-custom-container"
IMAGE_URI = f"{LOCATION}-docker.pkg.dev/{PROJECT_ID}/{DOCKER_REPOSITORY}/{IMAGE_NAME}"

! gcloud builds submit --project {PROJECT_ID} --region={LOCATION} --tag {IMAGE_URI} --timeout=60m --quiet

Create and run the custom training job

Configure a custom job with the custom container image.

In [ ]:

JOB_NAME = "tensorboard-job-unique"

job = aiplatform.CustomContainerTrainingJob(
    display_name=JOB_NAME, container_uri=IMAGE_URI
)

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:

• args: The command-line arguments to pass to the training script.
  • --epochs : The number of epochs for training.
• 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.
• tensorboard: The TensorBoard instance.
• service_account: The service account.
• sync: Whether to block until completion of the job.

In [ ]:

base_output_dir = "{}/{}".format(BUCKET_URI, JOB_NAME)
MACHINE_TYPE = "n1-standard-4"
EPOCHS = 2
training_args = [
    "--epochs=" + str(EPOCHS),
]

job.run(
    args=training_args,
    replica_count=1,
    machine_type=MACHINE_TYPE,
    base_output_dir=base_output_dir,
    tensorboard=TENSORBOARD_INSTANCE_NAME,
    service_account=SERVICE_ACCOUNT,
)

View the Cloud Profiler dashboard

When the custom job state switches to Running, you can access the Cloud Profiler dashboard through the Custom jobs page or the Experiments page on the Google Cloud console.

The Google Cloud guide to Profile model training performance using Cloud Profiler provides detailed instructions for accessing the Cloud Profiler dashboard and capturing a profiling session.

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:

• Docker repository
• Training job
• TensorBoard instance
• Cloud Storage bucket

In [ ]:

delete_tensorboard = True
delete_bucket = False
delete_generated_files_after_execution = False

# Delete docker repository.
! gcloud artifacts repositories delete $DOCKER_REPOSITORY --project {PROJECT_ID} --location {LOCATION} --quiet

job.delete()

if delete_tensorboard:
    tensorboard.delete()

if delete_bucket and "BUCKET_URI" in globals():
    ! gsutil -m rm -r $BUCKET_URI

if delete_generated_files_after_execution:
    ! rm -rf $PYTHON_PACKAGE_APPLICATION_DIR Dockerfile Dockerfile