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

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

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

訓練指令碼變更

訓練指令碼必須設定為將 TensorBoard 記錄寫入 Cloud Storage 值區，Gemini Enterprise Agent Platform Training Service 會透過預先定義的環境變數 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],
  )

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

建立及執行管道

以下範例說明如何使用 Kubeflow Pipelines DSL 套件建構及執行管道。如需更多範例和詳細資訊，請參閱 Agent Platform 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"
    )

提交 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：要在哪個位置執行管道。這應與您使用的 TensorBoard 執行個體位於相同位置。

後續步驟