使用 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.adminroles/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

笔记本

后续步骤