搭配使用 Vertex AI TensorBoard 與 Vertex AI Pipelines

您可以將訓練程式碼封裝到自訂訓練元件中，並在管道工作中執行。TensorBoard 記錄檔會自動串流至 Vertex AI TensorBoard 實驗。您可以透過這項整合功能，在 Vertex AI TensorBoard 將 Vertex AI TensorBoard 記錄檔寫入 Cloud Storage 時，近乎即時地監控訓練作業。

如需初始設定，請參閱「設定 Vertex AI TensorBoard」。

訓練指令碼變更

訓練指令碼必須設定為將 TensorBoard 記錄寫入 Cloud Storage 值區，而 Vertex AI Training 服務會透過預先定義的環境變數 AIP_TENSORBOARD_LOG_DIR，自動提供該值區的位置。

通常只要將 os.environ['AIP_TENSORBOARD_LOG_DIR'] 提供給開放原始碼 TensorBoard 記錄寫入 API 做為記錄目錄即可。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 執行個體，並將 staging_bucket 設為在 API 呼叫期間 (包括 TensorBoard 記錄) 暫存構件的位置。

接著，建構包含這項工作的管道，並將管道編譯為 JSON 檔案。

如需更多範例和資訊，請參閱「自訂工作元件」和「建構管道」。

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 管道

使用 Python 適用的 Vertex AI SDK 提交管道。詳情請參閱「執行管道」。

Python

from typing import Any, Dict, Optional

from google.cloud import aiplatform


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：。用於執行管道的 Google Cloud 專案。您可以在 Google Cloud 控制台 的歡迎頁面中找到 ID。
  • location：執行管道的位置。這個位置應與您使用的 TensorBoard 執行個體相同。

後續步驟