使用 Profiler 分析模型训练性能

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

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

准备工作

启用 Vertex AI TensorBoard Profiler

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

  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 Profiler 信息中心

将训练脚本配置为启用 Vertex AI TensorBoard Profiler 后，使用 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 Console 访问 Vertex AI TensorBoard Profiler 信息中心：

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

通过“自定义作业”页面访问 Profiler 信息中心

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

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

    转到自定义作业

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

  3. 点击打开 TensorBoard。

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

通过“实验”页面访问 Profiler 信息中心

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

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

    转到 Vertex AI 实验

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

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

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

捕获性能剖析会话

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

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

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

    workerpool0-0

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

  4. 点击 Capture。

笔记本

后续步骤