使用 Vertex AI TensorBoard 性能分析器分析模型训练性能

本页面介绍如何启用 Vertex AI TensorBoard 性能分析器，以便为自定义训练作业调试模型训练性能。

训练模型的计算开销可能很大。Vertex AI TensorBoard 性能分析器可以帮助您了解训练操作的资源消耗情况，从而使您可以监控和优化模型训练性能。利用这些信息，您可以精确地找出并解决性能瓶颈，以更快的速度、更低的费用训练模型。

准备工作

启用 Vertex AI TensorBoard 性能分析器

如需为训练作业启用 Vertex AI TensorBoard 性能分析器，请将以下内容添加到训练脚本中：

  1. 在顶级导入中添加 cloud_profiler 导入：

    from google.cloud.aiplatform.training_utils import cloud_profiler

  2. 通过添加以下代码初始化 cloud_profiler 插件：

    cloud_profiler.init()

示例

下面是一个训练脚本示例：

#!/usr/bin/env python

import tensorflow as tf
import argparse
import os
from google.cloud.aiplatform.training_utils import cloud_profiler
import time

"""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):
    strategy = None
    if args.distributed:
        strategy = tf.distribute.MultiWorkerMirroredStrategy()

    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

    if args.distributed:
        strategy = tf.distribute.MultiWorkerMirroredStrategy()
        with strategy.scope():
            model = _create_model()
            model.compile(
                optimizer="adam",
                loss=tf.keras.losses.sparse_categorical_crossentropy,
                metrics=["accuracy"],
            )
    else:
        model = _create_model()
        model.compile(
            optimizer="adam",
            loss=tf.keras.losses.sparse_categorical_crossentropy,
            metrics=["accuracy"],
        )

    # Initialize the profiler.
    cloud_profiler.init()

    # Use AIP_TENSORBOARD_LOG_DIR to update where logs are written to.
    tensorboard_callback = tf.keras.callbacks.TensorBoard(
        log_dir=os.environ["AIP_TENSORBOARD_LOG_DIR"], histogram_freq=1
    )

    model.fit(
        x_train,
        y_train,
        epochs=args.epochs,
        verbose=0,
        callbacks=[tensorboard_callback],
    )
if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--epochs", type=int, default=100, help="Number of epochs to run model."
    )
    parser.add_argument(
        "--distributed", action="store_true", help="Use MultiWorkerMirroredStrategy"
    )
    args = parser.parse_args()
    main(args)

访问 Vertex AI TensorBoard 性能分析器信息中心

将训练脚本配置为启用 Vertex AI TensorBoard 性能分析器后，使用 Vertex AI TensorBoard 实例运行训练脚本。

在训练脚本中，确保以下配置：

  • 将 BASE_OUTPUT_DIR: 设置为您要在其中存储训练脚本生成的 Vertex AI TensorBoard 日志的 Cloud Storage 存储桶。
  • 将 'serviceAccount': 设置为您使用 roles/storage.admin 和 roles/aiplatform.user 角色创建的服务账号。

  • 将 'tensorboard': 设置为要用于此训练作业的 Vertex AI TensorBoard 实例的完全限定名称。完全限定的名称采用以下格式：

    projects/PROJECT_NUMBER_OR_ID/locations/REGION/tensorboards/TENSORBOARD_INSTANCE_ID

您可以通过两种方法从 Google Cloud 控制台访问 Vertex AI TensorBoard 性能分析器信息中心：

  • 从“自定义作业”页面。
  • 在“实验”页面。

通过“自定义作业”页面访问性能分析器信息中心

即使训练作业处于已完成状态，也可以使用此方法访问 Vertex AI TensorBoard 性能分析器信息中心。

  1. 在 Google Cloud 控制台中，转到训练页面上的自定义作业标签页。

    转到自定义作业

  2. 点击您刚刚创建的训练作业的名称，以转到作业详情页面。

  3. 点击打开 TensorBoard。

  4. 点击性能分析标签页。

通过“实验”页面访问性能分析器信息中心

仅当训练作业处于正在运行状态时，才能使用此方法访问 Vertex AI TensorBoard 性能分析器信息中心。

  1. 在 Google Cloud 控制台中，转到“Vertex AI 实验”页面。

    转到 Vertex AI 实验

  2. 选择刚创建的训练作业的区域。

  3. 点击训练作业名称旁边的打开 TensorBoard。

  4. 点击性能分析标签页。

捕获性能剖析会话

如需捕获性能剖析会话，您的训练作业必须处于正在运行状态。在 Vertex AI TensorBoard 实例的性能分析标签页中，执行以下步骤：

  1. 点击捕获性能剖析文件。

  2. 在配置文件服务网址或 TPU 名称字段中，输入：

    workerpool0-0

  3. 对于地址类型，选择 IP 地址。

  4. 点击 Capture。

笔记本

后续步骤