Vertex AI TensorBoard Profiler を使用してモデルのトレーニングパフォーマンスのプロファイルを作成する

このページでは、Vertex AI TensorBoard Profiler を有効にして、カスタムトレーニングジョブでモデルのトレーニングのパフォーマンスをデバッグする方法について説明します。

モデルのトレーニングは、計算費用が高額になることがあります。Vertex AI TensorBoard Profiler を使用すると、トレーニングオペレーションのリソース消費量を把握し、モデルのトレーニングのパフォーマンスをモニタリングして最適化できます。この情報により、パフォーマンスのボトルネックを特定して修正し、モデルをより早く低コストでトレーニングできます。

始める前に

TensorFlow 2.4 以降を使用していることを確認します。
cloud_profiler プラグインを含む Vertex AI SDK をインストールします。ローカルの Docker コンテナから、次のコマンドを実行します。
```
pip install google-cloud-aiplatform[cloud_profiler]
```
Vertex AI TensorBoard インスタンスを用意する必要があります。手順については、Vertex AI TensorBoard インスタンスを作成するをご覧ください。
roles/storage.admin ロールと roles/aiplatform.user ロールを持つサービスアカウントを用意する必要があります。手順については、必要な権限が付与されたサービスアカウントを作成するをご覧ください。
Vertex AI TensorBoard のログを保存する Cloud Storage バケットを用意する必要があります。手順については、Vertex AI TensorBoard のログを保存する Cloud Storage バケットを作成するをご覧ください。

Vertex AI TensorBoard Profiler を有効にする

トレーニングジョブで Vertex AI TensorBoard Profiler を有効にするには、トレーニングスクリプトに次の内容を追加します。

最上位のインポートに cloud_profiler のインポートを追加します。
```
from google.cloud.aiplatform.training_utils import cloud_profiler
```
次の行を追加して 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 コンソールから Vertex AI TensorBoard Profiler ダッシュボードにアクセスする方法は 2 つあります。

[カスタムジョブ] ページから。
[テスト] ページから。

[カスタムジョブ] ページから Profiler ダッシュボードにアクセスする

トレーニングジョブが [完了] ステータスでも、以下の方法で Vertex AI TensorBoard Profiler ダッシュボードにアクセスできます。

Google Cloud コンソールで、[トレーニング] ページの [カスタムジョブ] タブに移動します。

[カスタムジョブ] に移動
作成したトレーニングジョブの名前をクリックして、ジョブの詳細ページに移動します。
[TensorBoard を開く] をクリックします。
[プロファイル] タブをクリックします

[テスト] ページから Profiler ダッシュボードにアクセスする

トレーニングジョブが [実行中] ステータスの場合にのみ、以下の方法を使用して Vertex AI TensorBoard Profiler ダッシュボードにアクセスできます。

Google Cloud コンソールで、[Vertex AI テスト] ページに移動します。

[Vertex AI テスト] に移動
作成したトレーニングジョブのリージョンを選択します。
トレーニングジョブの名前の横にある [TensorBoard を開く] をクリックします。
[プロファイル] タブをクリックします

プロファイリングセッションをキャプチャする

プロファイリングセッションをキャプチャするには、トレーニングジョブが Running の状態である必要があります。Vertex AI TensorBoard インスタンスの [Profile] タブで、次の操作を行います。

[Capture profile] をクリックします。
[Profile Service URL(s) or TPU name] フィールドに、次のように入力します。
```
workerpool0-0
```
[Address type] では [IP address] を選択します。
[Capture] をクリックします

ノートブック

次のステップ

Profiler ツールの詳細と、これらのツールを使用してモデルのパフォーマンスを最適化する方法については、Tensorflow Profiler のドキュメントをご覧ください。

Vertex AI TensorBoard Profiler を使用してモデルのトレーニング パフォーマンスのプロファイルを作成する