搭配使用 Vertex AI TensorBoard 和 Vertex AI Pipelines

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

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

訓練指令碼異動

訓練指令碼必須設為將 TensorBoard 記錄寫入 Cloud Storage 值區,Vertex AI 訓練服務會透過預先定義的環境變數 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 適用的 Vertex AI SDK

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

後續步驟