Vertex AI Pipelines で Vertex AI TensorBoard を使用する

トレーニングコードは、カスタムトレーニングコンポーネントにパッケージ化して、パイプラインジョブで実行できます。TensorBoard のログは、Vertex AI TensorBoard テストに自動的にストリーミングされます。このインテグレーションを使用すると、Vertex AI TensorBoard ログが Cloud Storage に書き込まれる際に、そのログの Vertex AI TensorBoard ストリームとして、ほぼリアルタイムでトレーニングをモニタリングできます。

初期設定については、Vertex AI TensorBoard を設定するをご覧ください。

トレーニングスクリプトの変更点

トレーニングスクリプトは、TensorBoard ログを Cloud Storage バケットに書き込むように構成する必要があります。Cloud Strage バケットの場所は、Vertex AI Training サービスにより事前定義の環境変数 AIP_TENSORBOARD_LOG_DIR で自動的に利用可能にされます。

通常これは、オープンソースの TensorBoard ログ書き込み API に、ログディレクトリとして os.environ['AIP_TENSORBOARD_LOG_DIR'] を指定することで実現できます。AIP_TENSORBOARD_LOG_DIR の場所は、staging_bucket 変数で設定します。

TensorFlow 2.x でトレーニングスクリプトを構成するには、TensorBoard コールバックを作成して、log_dir 変数を os.environ['AIP_TENSORBOARD_LOG_DIR'] に設定します。これにより、TensorBoard コールバックが、TensorFlow model.fit コールバックのリストに含まれます。

  tensorboard_callback = tf.keras.callbacks.TensorBoard(
       log_dir=os.environ['AIP_TENSORBOARD_LOG_DIR'],
       histogram_freq=1
  )

  model.fit(
       x=x_train,
       y=y_train,
       epochs=epochs,
       validation_data=(x_test, y_test),
       callbacks=[tensorboard_callback],
  )

詳しくは、Vertex AI がカスタムトレーニング環境で環境変数を設定する方法をご覧ください。

パイプラインを構築して実行する

次の例では、Kubeflow Pipelines DSL パッケージでパイプラインを構築して実行する方法を示します。その他の例と詳細については、Vertex AI Pipelines のドキュメントをご覧ください。

トレーニングコンポーネントを作成する

トレーニングコードをカスタムコンポーネントにパッケージ化します。そのコードでは、TensorBoard ログを Cloud Storage バケットに書き込むように構成してください。他の例については、独自のパイプラインコンポーネントの構築をご覧ください。

from kfp.v2.dsl import component

@component(
    base_image="tensorflow/tensorflow:latest",
    packages_to_install=["tensorflow_datasets"],
)
def train_tensorflow_model_with_tensorboard():
    import datetime, os
    import tensorflow as tf

    (x_train, y_train), (x_test, y_test) = tf.keras.datasets.mnist.load_data()
    x_train, x_test = x_train / 255.0, x_test / 255.0

    def create_model():
        return tf.keras.models.Sequential(
            [
                tf.keras.layers.Flatten(input_shape=(28, 28)),
                tf.keras.layers.Dense(512, activation="relu"),
            ]
        )

    model = create_model()
    model.compile(
        optimizer="adam",
        loss="sparse_categorical_crossentropy",
        metrics=["accuracy"]
    )

    tensorboard_callback = tf.keras.callbacks.TensorBoard(
        log_dir=os.environ['AIP_TENSORBOARD_LOG_DIR'],
        histogram_freq=1
    )

    model.fit(
        x=x_train,
        y=y_train,
        epochs=5,
        validation_data=(x_test, y_test),
        callbacks=[tensorboard_callback],
    )

パイプラインを構築してコンパイルする

create_custom_training_job_op_from_component でコンポーネント仕様を指定し、作成したコンポーネントからカスタムトレーニングジョブを作成します。tensorboard_resource_name を TensorBoard インスタンスに設定し、API 呼び出し中にアーティファクト（TensorBoard ログを含む）をステージングする場所に staging_bucket を設定します。

次に、このジョブを含めるパイプラインを構築し、そのパイプラインを JSON ファイルにコンパイルします。

他の例と詳細については、CustomJob コンポーネントとパイプラインの構築をご覧ください。

from kfp.v2 import compiler
from google_cloud_pipeline_components.v1.custom_job.utils import \
    create_custom_training_job_op_from_component
from kfp.v2 import dsl

def create_tensorboard_pipeline_sample(
    project, location, staging_bucket, display_name, service_account, experiment, tensorboard_resource_name
):

    @dsl.pipeline(
        pipeline_root=f"{staging_bucket}/pipeline_root",
        name=display_name,
    )
    def pipeline():
        custom_job_op = create_custom_training_job_op_from_component(
            component_spec=train_tensorflow_model_with_tensorboard,
            tensorboard=tensorboard_resource_name,
            base_output_directory=staging_bucket,
            service_account=service_account,
        )
        custom_job_op(project=project, location=location)

    compiler.Compiler().compile(
        pipeline_func=pipeline, package_path=f"{display_name}.json"
    )

Vertex AI パイプラインを送信する

パイプラインは、Vertex AI SDK for Python を使用して送信します。詳細については、パイプラインを実行するをご覧ください。

Python

def log_pipeline_job_to_experiment_sample(
    experiment_name: str,
    pipeline_job_display_name: str,
    template_path: str,
    pipeline_root: str,
    project: str,
    location: str,
    parameter_values: Optional[Dict[str, Any]] = None,
):
    aiplatform.init(project=project, location=location)

    pipeline_job = aiplatform.PipelineJob(
        display_name=pipeline_job_display_name,
        template_path=template_path,
        pipeline_root=pipeline_root,
        parameter_values=parameter_values,
    )

    pipeline_job.submit(experiment=experiment_name)

experiment_name: テストの名前を指定します。
pipeline_job_display_name: パイプラインジョブの表示名。
template_path: コンパイルされたパイプラインテンプレートへのパス。
pipeline_root: パイプラインサービスアカウントがアクセスできる Cloud Storage URI を指定します。パイプライン実行のアーティファクトはパイプラインルート内に保存されます。
parameter_values: （省略可）この実行に渡すパイプラインパラメータ。たとえば、パラメータ名をディクショナリキー、パラメータ値をディクショナリ値として dict() を作成します。
project: 実際のプロジェクト ID。パイプラインを実行する Google Cloud プロジェクト。ID は、Google Cloud コンソールの [ようこそ] ページで確認できます。
location: パイプラインを実行するリージョン。これは、使用している TensorBoard インスタンスと同じリージョンにする必要があります。

次のステップ

結果を表示する（Vertex AI Pipelines の TensorBoard を表示する）。
Vertex AI TensorBoard Profiler を使用して、カスタムトレーニングジョブのパフォーマンスを最適化する方法を確認する。