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

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

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

始める前に

プロファイラを有効にする

トレーニング ジョブで 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)

Profiler ダッシュボードにアクセスする

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 コンソールから Profiler ダッシュボードにアクセスする方法は 2 つあります。

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

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

トレーニング ジョブが「完了」状態でも、以下の方法で Profiler ダッシュボードにアクセスできます。

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

    [カスタムジョブ] に移動

  2. 作成したトレーニング ジョブの名前をクリックして、ジョブの詳細ページに移動します。

  3. [TensorBoard を開く] をクリックします。

  4. [Profile] タブをクリックします

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

トレーニング ジョブが「実行中」状態の場合にのみ、この方法を使用して Profiler ダッシュボードにアクセスできます。

  1. Google Cloud コンソールで、[Vertex AI Experiments] ページに移動します。

    [Vertex AI Experiments] に移動

  2. 作成したトレーニング ジョブのリージョンを選択します。

  3. トレーニング ジョブの名前の横にある [TensorBoard を開く] をクリックします。

  4. [プロファイル] タブをクリックします

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

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

  1. [Capture profile] をクリックします。
  2. [Profile Service URL(s) or TPU name] フィールドに、次のように入力します。

    workerpool0-0
    
  3. [Address type] では [IP address] を選択します。

  4. [Capture] をクリックします

ノートブック

次のステップ

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