搭配 Gemini Enterprise Agent Platform Pipelines 使用 Vertex AI TensorBoard

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

如要瞭解初始設定，請參閱「設定 Vertex AI TensorBoard」。

訓練指令碼異動

您必須設定訓練指令碼，將 TensorBoard 記錄寫入 Cloud Storage bucket，Gemini Enterprise Agent Platform 訓練服務會透過預先定義的環境變數 AIP_TENSORBOARD_LOG_DIR 自動提供該 bucket 的位置。

通常只要將 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],
  )

進一步瞭解 Agent Platform 如何在自訂訓練環境中設定環境變數。

建構及執行管道

以下範例說明如何使用 Kubeflow Pipelines DSL 套件建構及執行管道。如需更多範例和詳細資料，請參閱 Agent Platform Pipelines 說明文件。

建立訓練元件

將訓練程式碼封裝至自訂元件，並確保程式碼已設定為將 TensorBoard 記錄寫入 Cloud Storage bucket。如需更多範例，請參閱「建構自己的管道元件」。

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],
    )

建構及編譯 pipeline

在 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"
    )

提交 Gemini Enterprise Agent Platform 管道

使用 Agent Platform SDK for Python 提交管道。詳情請參閱「執行管道」。

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：執行 pipeline 的位置。這個位置應與您使用的 TensorBoard 執行個體相同。

後續步驟