feat: Reorganize all TensorBoard notebooks together under directory 'tensorboard" (#1296)

**REQUIRED:** Add a summary of your PR here, typically including why the change is needed and what was changed. Include any design alternatives for discussion purposes.

<br>
This is the first step of two to move the TensorBoard profiler to tensorboard directory.
<br><br><br>

**REQUIRED:** Fill out the below checklists or remove if irrelevant
1. If you are opening a PR for `Official Notebooks` under the [notebooks/official](https://togithub.com/GoogleCloudPlatform/vertex-ai-samples/tree/main/notebooks/official) folder, follow this mandatory checklist:
- [ ] Use the [notebook template](https://togithub.com/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/notebook_template.ipynb) as a starting point.
- [ ] Follow the style and grammar rules outlined in the above notebook template.
- [ ] Verify the notebook runs successfully in Colab since the automated tests cannot guarantee this even when it passes.
- [ ] Passes all the required automated checks. You can locally test for formatting and linting with these [instructions](https://togithub.com/GoogleCloudPlatform/vertex-ai-samples/blob/main/CONTRIBUTING.md#code-quality-checks).
- [ ] You have consulted with a tech writer to see if tech writer review is necessary. If so, the notebook has been reviewed by a tech writer, and they have approved it.
- [ ] This notebook has been added to the [CODEOWNERS](https://togithub.com/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/official/CODEOWNERS) file under the `Official Notebooks` section, pointing to the author or the author's team.
- [ ] The Jupyter notebook cleans up any artifacts it has created (datasets, ML models, endpoints, etc) so as not to eat up unnecessary resources.

<br>

2. If you are opening a PR for `Community Notebooks` under the [notebooks/community](https://togithub.com/GoogleCloudPlatform/vertex-ai-samples/tree/main/notebooks/community) folder:
- [ ] This notebook has been added to the [CODEOWNERS](https://togithub.com/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/community/CODEOWNERS) file under the `Community Notebooks` section, pointing to the author or the author's team.
- [ ] Passes all the required formatting and linting checks. You can locally test with these [instructions](https://togithub.com/GoogleCloudPlatform/vertex-ai-samples/blob/main/CONTRIBUTING.md#code-quality-checks).

<br>

3. If you are opening a PR for `Community Content` under the [community-content](https://togithub.com/GoogleCloudPlatform/vertex-ai-samples/tree/main/community-content) folder:
- [ ] Make sure your main `Content Directory Name` is descriptive, informative, and includes some of the key products and attributes of your content, so that it is differentiable from other content
- [ ] The main content directory has been added to the [CODEOWNERS](https://togithub.com/GoogleCloudPlatform/vertex-ai-samples/blob/main/community-content/CODEOWNERS) file under the `Community Content` section, pointing to the author or the author's team.
- [ ] Passes all the required formatting and linting checks. You can locally test with these [instructions](https://togithub.com/GoogleCloudPlatform/vertex-ai-samples/blob/main/CONTRIBUTING.md#code-quality-checks).
This commit is contained in:
gericdong
2022-11-23 16:48:15 +00:00
committed by GitHub
parent 1ff7fc70ed
commit ad1a827824
2 changed files with 953 additions and 0 deletions
+11
View File
@@ -20,3 +20,14 @@ The steps performed include:
* Create a custom container image with your customized training code.
* Setup service account and Google Cloud Storage buckets.
* Create & launch your custom training job with your custom container.
[Profile model training performance using Profiler](https://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/official/tensorboard/tensorboard_profiler_custom_training.ipynb)
Learn how to enable Vertex AI TensorBoard Profiler for custom training jobs.
The steps performed include:
* Setup a service account and a Cloud Storage bucket
* Create a TensorBoard instance
* Create and run a custom training job that enables TensorBoard Profiler
* View the TensorBoard Profiler dashboard to debug your model training performance
@@ -0,0 +1,942 @@
{
"cells": [
{
"cell_type": "markdown",
"metadata": {
"id": "l2mMvIUG9meX"
},
"source": [
"# Profile model training performance using Profiler\n",
"\n",
"<table align=\"left\">\n",
"\n",
" <td>\n",
" <a href=\"https://colab.research.google.com/github/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/official/tensorboard/tensorboard_profiler_custom_training.ipynb\">\n",
" <img src=\"https://cloud.google.com/ml-engine/images/colab-logo-32px.png\" alt=\"Colab logo\"> Run in Colab\n",
" </a>\n",
" </td>\n",
" <td>\n",
" <a href=\"https://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/official/tensorboard/tensorboard_profiler_custom_training.ipynb\">\n",
" <img src=\"https://cloud.google.com/ml-engine/images/github-logo-32px.png\" alt=\"GitHub logo\">\n",
" View on GitHub\n",
" </a>\n",
" </td>\n",
" <td>\n",
" <a href=\"https://console.cloud.google.com/vertex-ai/workbench/deploy-notebook?download_url=https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/main/notebooks/official/tensorboard/tensorboard_profiler_custom_training.ipynb\">\n",
" <img src=\"https://lh3.googleusercontent.com/UiNooY4LUgW_oTvpsNhPpQzsstV5W8F7rYgxgGBD85cWJoLmrOzhVs_ksK_vgx40SHs7jCqkTkCk=e14-rj-sc0xffffff-h130-w32\" alt=\"Vertex AI logo\">\n",
" Open in Vertex AI Workbench\n",
" </a>\n",
" </td> \n",
"</table>"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "ur8xi4C7S06n"
},
"outputs": [],
"source": [
"# Copyright 2022 Google LLC\n",
"#\n",
"# Licensed under the Apache License, Version 2.0 (the \"License\");\n",
"# you may not use this file except in compliance with the License.\n",
"# You may obtain a copy of the License at\n",
"#\n",
"# https://www.apache.org/licenses/LICENSE-2.0\n",
"#\n",
"# Unless required by applicable law or agreed to in writing, software\n",
"# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
"# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
"# See the License for the specific language governing permissions and\n",
"# limitations under the License."
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "tvgnzT1CKxrO"
},
"source": [
"## Overview\n",
"\n",
"Vertex AI TensorBoard 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 Vertex AI TensorBoard Profiler so you can debug model training performance for your custom training jobs.\n"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "dmfmQL6w84pS"
},
"source": [
"### Objective\n",
"\n",
"In this tutorial, you learn how to enable Vertex AI TensorBoard Profiler for custom training jobs.\n",
"\n",
"This tutorial uses the following Google Cloud AI services:\n",
"\n",
"- `Vertex AI Training`\n",
"- `Vertex AI TensorBoard`\n",
"\n",
"The steps performed include:\n",
"\n",
"- Setup a service account and a Cloud Storage bucket\n",
"- Create a TensorBoard instance\n",
"- Create and run a custom training job that enables TensorBoard Profiler\n",
"- View the TensorBoard Profiler dashboard to debug your model training performance\n"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "zfXf0r-K81Y-"
},
"source": [
"### Dataset\n",
"\n",
"The dataset used for this tutorial is the [mnist dataset](https://www.tensorflow.org/datasets/catalog/mnist) from [TensorFlow Datasets](https://www.tensorflow.org/datasets/catalog/overview).\n"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "I3KFLvpq87rs"
},
"source": [
"### Costs \n",
"\n",
"This tutorial uses billable components of Google Cloud:\n",
"\n",
"* Vertex AI\n",
"* Cloud Storage\n",
"\n",
"Learn about [Vertex AI\n",
"pricing](https://cloud.google.com/vertex-ai/pricing) and [Cloud Storage\n",
"pricing](https://cloud.google.com/storage/pricing), and use the [Pricing\n",
"Calculator](https://cloud.google.com/products/calculator/)\n",
"to generate a cost estimate based on your projected usage."
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "ze4-nDLfK4pw"
},
"source": [
"## Installation\n",
"\n",
"Install the following packages required to execute this notebook. "
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "2b4ef9b72d43"
},
"outputs": [],
"source": [
"! pip3 install --user --upgrade google-cloud-aiplatform --quiet"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "aUw6ibN-n5Za"
},
"source": [
"### Colab only: Uncomment the following cell to restart the kernel."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "FM12wbWhn7w0"
},
"outputs": [],
"source": [
"# Automatically restart kernel after installs so that your environment can access the new packages\n",
"# import IPython\n",
"\n",
"# app = IPython.Application.instance()\n",
"# app.kernel.do_shutdown(True)"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "LgFWLeJfoGQu"
},
"source": [
"## Before you begin\n",
"\n",
"### Set up your Google Cloud project\n",
"\n",
"**The following steps are required, regardless of your notebook environment.**\n",
"\n",
"1. [Select or create a Google Cloud project](https://console.cloud.google.com/cloud-resource-manager). When you first create an account, you get a $300 free credit towards your compute/storage costs.\n",
"\n",
"2. [Make sure that billing is enabled for your project](https://cloud.google.com/billing/docs/how-to/modify-project).\n",
"\n",
"3. [Enable the Vertex AI API](https://console.cloud.google.com/flows/enableapi?apiid=aiplatform.googleapis.com). \n",
"\n",
"4. If you are running this notebook locally, you need to install the [Cloud SDK](https://cloud.google.com/sdk)."
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "8ckyxpX_oSzD"
},
"source": [
"#### Set your project ID\n",
"\n",
"**If you don't know your project ID**, try the following:\n",
"* Run `gcloud config list`.\n",
"* Run `gcloud projects list`.\n",
"* See the support page: [Locate the project ID](https://support.google.com/googleapi/answer/7014113)"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "zY8DKBoVoVy3"
},
"outputs": [],
"source": [
"PROJECT_ID = \"[your-project-id]\" # @param {type:\"string\"}\n",
"\n",
"# Set the project id\n",
"! gcloud config set project {PROJECT_ID}"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "mSQjVQmMosMl"
},
"source": [
"#### Region\n",
"\n",
"You can also change the `REGION` variable used by Vertex AI. Learn more about [Vertex AI regions](https://cloud.google.com/vertex-ai/docs/general/locations)."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "Se9FWWhLotvB"
},
"outputs": [],
"source": [
"REGION = \"us-central1\" # @param {type:\"string\"}"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "IfJRIMBpo5Pg"
},
"source": [
"### Authenticate your Google Cloud account\n",
"\n",
"Depending on your Jupyter environment, you may have to manually authenticate. Follow the relevant instructions below."
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "acFN0s3So9-Y"
},
"source": [
"**1. Vertex AI Workbench**\n",
"* Do nothing as you are already authenticated."
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "dQ_mNwuapE5T"
},
"source": [
"**2. Local JupyterLab instance, uncomment and run:**"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "cR_MzpknpGgM"
},
"outputs": [],
"source": [
"# ! gcloud auth login"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "h-MuVI_ypJfw"
},
"source": [
"**3. Colab, uncomment and run:**"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "BeaQlCwMpQUT"
},
"outputs": [],
"source": [
"# from google.colab import auth\n",
"# auth.authenticate_user()"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "3ivZkPUjpaFz"
},
"source": [
"**4. Setup service account and permissions**\n",
"\n",
"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](https://cloud.google.com/iam/docs/creating-managing-service-accounts#creating)."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "vYE3b942wza4"
},
"outputs": [],
"source": [
"SERVICE_ACCOUNT = \"[your-service-account]\" # @param {type:\"string\"}"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "WWIxsCJFCg5Z"
},
"outputs": [],
"source": [
"# Grant Cloud Storage permission.\n",
"! gcloud projects add-iam-policy-binding $PROJECT_ID \\\n",
" --member=\"serviceAccount:$SERVICE_ACCOUNT\" \\\n",
" --role=\"roles/storage.admin\" \\\n",
" --quiet"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "S_8_5jm-Gk6w"
},
"outputs": [],
"source": [
"# Grant AI Platform permission.\n",
"! gcloud projects add-iam-policy-binding $PROJECT_ID \\\n",
" --member=\"serviceAccount:$SERVICE_ACCOUNT\" \\\n",
" --role=\"roles/aiplatform.user\" \\\n",
" --quiet"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "OKtKGmr9pfr6"
},
"source": [
"### Create a Cloud Storage bucket\n",
"\n",
"Create a storage bucket to store intermediate artifacts such as datasets."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "In3aQanwYjFB"
},
"outputs": [],
"source": [
"BUCKET_URI = \"gs://your-bucket-name-unique\" # @param {type:\"string\"}"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "GOaOsIjxp0oB"
},
"source": [
"**Only if your bucket doesn't already exist**: Run the following cell to create your Cloud Storage bucket."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "Wn5QiIl2p16e"
},
"outputs": [],
"source": [
"! gsutil mb -l $REGION -p $PROJECT_ID $BUCKET_URI"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "ankcS-vtp7Wv"
},
"source": [
"### Import libraries"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "WffSImMvp-Po"
},
"outputs": [],
"source": [
"from google.cloud import aiplatform"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "OMrAJ8RGqBQu"
},
"source": [
"### Initialize Vertex AI SDK for Python\n",
"\n",
"Initialize the Vertex AI SDK for Python for your project and corresponding bucket."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "AWRzBFExqERG"
},
"outputs": [],
"source": [
"aiplatform.init(project=PROJECT_ID, location=REGION, staging_bucket=BUCKET_URI)"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "MYzX9ZgjsyyY"
},
"source": [
"### Enable Artifact Registry API\n",
"\n",
"First, you must enable the Artifact Registry API service for your project.\n",
"\n",
"Learn more about [Enabling service](https://cloud.google.com/artifact-registry/docs/enable-service)."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "bG4Rpt-As55f"
},
"outputs": [],
"source": [
"! gcloud services enable artifactregistry.googleapis.com --quiet"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "-ayTbNdi62_t"
},
"source": [
"### Create a TensorBoard instance\n",
"\n",
"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."
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "9c3QrDTZdaxk"
},
"source": [
"#### Set your TensorBoard instance display name\n"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "azlwb__AX8gs"
},
"outputs": [],
"source": [
"TENSORBOARD_NAME = \"your-tensorboard-unique\" # @param {type:\"string\"}"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "vJrWKK0mY7H7"
},
"source": [
"#### Create a TensorBoard instance\n",
"\n",
"If you don't have a TensorBoard instance, create one by running the following cell:"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "JqVNsRFrc_78"
},
"outputs": [],
"source": [
"tensorboard = aiplatform.Tensorboard.create(\n",
" display_name=TENSORBOARD_NAME, project=PROJECT_ID, location=REGION\n",
")\n",
"\n",
"TENSORBOARD_INSTANCE_NAME = tensorboard.resource_name\n",
"\n",
"print(\"TensorBoard instance name:\", TENSORBOARD_INSTANCE_NAME)"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "yoR29gW2S24w"
},
"source": [
"## Train a model\n",
"\n",
"To train a model using your custom training code, choose one of the following options:\n",
"\n",
"- **Prebuilt container**: Load your custom training code as a Python package to a prebuilt container image from Google Cloud.\n",
"\n",
"- **Custom container**: Create your own container image that contains your custom training code.\n",
"\n",
"In this tutorial, we will train a custom model using a custom container."
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "OrpUIkAIs_uQ"
},
"source": [
"### Create a private Docker repository\n",
"\n",
"Your first step is to create your own Docker repository in Google Artifact Registry."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "0amu4063tDnG"
},
"outputs": [],
"source": [
"DOCKER_REPOSITORY = f\"{PROJECT_ID}-repo-unique\"\n",
"\n",
"! gcloud artifacts repositories create {DOCKER_REPOSITORY} \\\n",
" --repository-format=docker \\\n",
" --location={REGION} \\\n",
" --description=\"Repository for TensorBoard Custom Training Job\" \\\n",
" --quiet\n",
"\n",
"! gcloud artifacts repositories list"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "MPqbvhO6tK_e"
},
"source": [
"### Configure authentication to your private Docker repository\n",
"\n",
"Before you push or pull container images, configure Docker to use the `gcloud` command-line tool to authenticate requests to `Artifact Registry` for your region."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "yAeCXZiStPCX"
},
"outputs": [],
"source": [
"import sys\n",
"\n",
"IS_COLAB = \"google.colab\" in sys.modules\n",
"\n",
"if not IS_COLAB:\n",
" ! gcloud auth configure-docker {REGION}-docker.pkg.dev --quiet"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "pSXrzCWi-bTD"
},
"source": [
"### Create a custom container image and push to your private Docker repository\n",
"\n",
"First, you create a training script file and a docker file.\n",
"\n",
"Create a directory for all of your training code."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "ZYO55-ZK-lm2"
},
"outputs": [],
"source": [
"PYTHON_PACKAGE_APPLICATION_DIR = \"trainer\"\n",
"\n",
"!mkdir -p $PYTHON_PACKAGE_APPLICATION_DIR"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "hyAwgsoQmaYI"
},
"source": [
"#### Prepare the training script\n",
"\n",
"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`.\n",
"\n",
"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. \n",
"\n",
"For example, in TensorFlow 2.x, you can use following code to create a tensorboard_callback: \n",
"\n",
" tensorboard_callback = tf.keras.callbacks.TensorBoard( \n",
" log_dir=os.environ['AIP_TENSORBOARD_LOG_DIR'], \n",
" histogram_freq=1) \n",
"`AIP_TENSORBOARD_LOG_DIR` is in the `BASE_OUTPUT_DIR` that you provide when creating the custom training job.\n",
"\n",
"To enable Vertex AI TensorBoard Profiler for your training job, add the following to your training script:\n",
"\n",
"Add the cloud_profiler import at your top level imports:\n",
"\n",
" from google.cloud.aiplatform.training_utils import cloud_profiler\n",
"\n",
"\n",
"Initialize the cloud_profiler plugin by adding:\n",
"\n",
"\n",
" cloud_profiler.init()"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "8JCgWW7Au1w8"
},
"outputs": [],
"source": [
"%%writefile trainer/task.py\n",
"\n",
"import tensorflow as tf\n",
"import argparse\n",
"import os\n",
"import sys, traceback\n",
"from google.cloud.aiplatform.training_utils import cloud_profiler\n",
"\n",
"\"\"\"Train an mnist model and use cloud_profiler for profiling.\"\"\"\n",
"\n",
"def _create_model():\n",
" model = tf.keras.models.Sequential(\n",
" [\n",
" tf.keras.layers.Flatten(input_shape=(28, 28)),\n",
" tf.keras.layers.Dense(128, activation=\"relu\"),\n",
" tf.keras.layers.Dropout(0.2),\n",
" tf.keras.layers.Dense(10),\n",
" ]\n",
" )\n",
" return model\n",
"\n",
"\n",
"def main(args):\n",
" print('Loading and preprocessing data ...')\n",
" mnist = tf.keras.datasets.mnist\n",
"\n",
" (x_train, y_train), (x_test, y_test) = mnist.load_data()\n",
" x_train, x_test = x_train / 255.0, x_test / 255.0\n",
"\n",
" print('Creating and training model ...')\n",
"\n",
" model = _create_model()\n",
" model.compile(\n",
" optimizer=\"adam\",\n",
" loss=tf.keras.losses.sparse_categorical_crossentropy,\n",
" metrics=[\"accuracy\"],\n",
" )\n",
"\n",
" # Initialize the profiler.\n",
" print('Initialize the profiler ...')\n",
" \n",
" try:\n",
" cloud_profiler.init()\n",
" except:\n",
" ex_type, ex_value, ex_traceback = sys.exc_info()\n",
" print(\"*** Unexpected:\", ex_type.__name__, ex_value)\n",
" traceback.print_tb(ex_traceback, limit=10, file=sys.stdout)\n",
" \n",
" print('The profiler initiated.')\n",
"\n",
" log_dir = \"logs\"\n",
" if 'AIP_TENSORBOARD_LOG_DIR' in os.environ:\n",
" log_dir = os.environ['AIP_TENSORBOARD_LOG_DIR']\n",
"\n",
" print('Setting up the TensorBoard callback ...')\n",
" tensorboard_callback = tf.keras.callbacks.TensorBoard(\n",
" log_dir=log_dir,\n",
" histogram_freq=1)\n",
"\n",
" print('Training model ...')\n",
" model.fit(\n",
" x_train,\n",
" y_train,\n",
" epochs=args.epochs,\n",
" verbose=0,\n",
" callbacks=[tensorboard_callback],\n",
" )\n",
" print('Training completed.')\n",
"\n",
" print('Saving model ...')\n",
"\n",
" model_dir = \"model\"\n",
" if 'AIP_MODEL_DIR' in os.environ:\n",
" model_dir = os.environ['AIP_MODEL_DIR']\n",
" tf.saved_model.save(model, model_dir)\n",
"\n",
" print('Model saved at ' + model_dir)\n",
"\n",
"\n",
"if __name__ == \"__main__\":\n",
" parser = argparse.ArgumentParser()\n",
" parser.add_argument(\n",
" \"--epochs\", type=int, default=100, help=\"Number of epochs to run model.\"\n",
" )\n",
" \n",
" args = parser.parse_args()\n",
" main(args)"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "LiOEvgkA_90f"
},
"source": [
"#### Prepare the Dockerfile\n"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "aKcS7yzcAXWf"
},
"outputs": [],
"source": [
"%%writefile Dockerfile\n",
"# Specifies base image and tag\n",
"FROM us-docker.pkg.dev/vertex-ai/training/tf-cpu.2-9:latest\n",
"WORKDIR /root\n",
"\n",
"# Installs additional packages as you need.\n",
"RUN pip3 install google-cloud-aiplatform[cloud_profiler]\n",
"\n",
"# Copies the trainer code to the docker image.\n",
"RUN mkdir /root/trainer\n",
"COPY trainer/task.py /root/trainer/task.py\n",
"\n",
"# Sets up the entry point to invoke the trainer.\n",
"ENTRYPOINT [\"python\", \"-m\", \"trainer.task\"]"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "ihYFahRAr6sj"
},
"source": [
"#### Build a custom container image and push to your private Docker repository"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "pDs1qDLDA27F"
},
"outputs": [],
"source": [
"IMAGE_NAME = \"tensorboard-custom-container\"\n",
"IMAGE_URI = f\"{REGION}-docker.pkg.dev/{PROJECT_ID}/{DOCKER_REPOSITORY}/{IMAGE_NAME}\"\n",
"\n",
"! gcloud builds submit --project {PROJECT_ID} --region={REGION} --tag {IMAGE_URI} --timeout=60m --quiet"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "k4e6OYmimqTR"
},
"source": [
"### Create and run the custom training job\n",
"\n",
"Configure a [custom job](https://cloud.google.com/vertex-ai/docs/training/create-custom-job) with the custom container image."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "jd6P-lQhFBPg"
},
"outputs": [],
"source": [
"JOB_NAME = \"tensorboard-job-unique\"\n",
"\n",
"job = aiplatform.CustomContainerTrainingJob(\n",
" display_name=JOB_NAME, container_uri=IMAGE_URI\n",
")"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "51hKGTbU32Eg"
},
"source": [
"#### Run the custom training job\n",
"\n",
"Next, you run the custom job to start the training job by invoking the method `run`, with the following parameters:\n",
"\n",
"- `args`: The command-line arguments to pass to the training script.\n",
" - `--epochs` : The number of epochs for training.\n",
"- `replica_count`: The number of compute instances for training (replica_count = 1 is single node training).\n",
"- `machine_type`: The machine type for the compute instances.\n",
"- `tensorboard`: The TensorBoard instance.\n",
"- `service_account`: The service account.\n",
"- `sync`: Whether to block until completion of the job."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "OC6_4KeI4Fit"
},
"outputs": [],
"source": [
"base_output_dir = \"{}/{}\".format(BUCKET_URI, JOB_NAME)\n",
"MACHINE_TYPE = \"n1-standard-4\"\n",
"EPOCHS = 2\n",
"training_args = [\n",
" \"--epochs=\" + str(EPOCHS),\n",
"]\n",
"\n",
"job.run(\n",
" args=training_args,\n",
" replica_count=1,\n",
" machine_type=MACHINE_TYPE,\n",
" base_output_dir=base_output_dir,\n",
" tensorboard=TENSORBOARD_INSTANCE_NAME,\n",
" service_account=SERVICE_ACCOUNT,\n",
")"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "JkEe2Nb_85UD"
},
"source": [
"## View the TensorBoard Profiler dashboard\n",
"\n",
"When the custom job state switches to `Running`, you can access the Vertex AI TensorBoard Profiler dashboard through the Custom jobs page or the Experiments page on the Google Cloud console. \n",
"\n",
"The Google Cloud guide to [Profile model training performance using Profiler](https://cloud.google.com/vertex-ai/docs/experiments/tensorboard-profiler) provides detailed instructions for accessing the Vertex AI TensorBoard Profiler dashboard and capturing a profiling session. \n"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "TpV-iwP9qw9c"
},
"source": [
"## Cleaning up\n",
"\n",
"To clean up all Google Cloud resources used in this project, you can [delete the Google Cloud\n",
"project](https://cloud.google.com/resource-manager/docs/creating-managing-projects#shutting_down_projects) you used for the tutorial.\n",
"\n",
"Otherwise, you can delete the individual resources you created in this tutorial:\n",
"\n",
"- Docker repository\n",
"- Training job\n",
"- TensorBoard instance\n",
"- Cloud Storage bucket\n"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "WR-ZhQ9XwpRI"
},
"outputs": [],
"source": [
"delete_tensorboard = True\n",
"delete_bucket = False\n",
"\n",
"# Delete docker repository.\n",
"! gcloud artifacts repositories delete $DOCKER_REPOSITORY --project {PROJECT_ID} --location {REGION} --quiet\n",
"\n",
"job.delete()\n",
"\n",
"if delete_tensorboard:\n",
" tensorboard.delete()\n",
"\n",
"if delete_bucket and \"BUCKET_URI\" in globals():\n",
" ! gsutil -m rm -r $BUCKET_URI"
]
}
],
"metadata": {
"colab": {
"name": "tensorboard_profiler_custom_training.ipynb",
"toc_visible": true
},
"kernelspec": {
"display_name": "Python 3",
"name": "python3"
}
},
"nbformat": 4,
"nbformat_minor": 0
}