Vertex AI TensorBoard mit Vertex AI Pipelines verwenden

Ihr Trainingscode kann in eine benutzerdefinierte Trainingskomponente verpackt und in einem Pipelinejob ausgeführt werden. TensorBoard-Logs werden automatisch an Ihren Vertex AI TensorBoard-Test gestreamt. Mit dieser Integration können Sie Ihr Training nahezu in Echtzeit überwachen, da Vertex AI TensorBoard-Logs in Vertex AI TensorBoard zu dem Zeitpunkt gestreamt werden, in dem sie in Cloud Storage geschrieben werden.

Informationen zur Ersteinrichtung finden Sie unter Für Vertex AI TensorBoard einrichten.

Änderungen an Ihrem Trainingsskript

Ihr Trainingsskript muss so konfiguriert sein, dass TensorBoard-Logs in den Cloud Storage-Bucket geschrieben werden. Dies ist der Speicherort, der vom Vertex AI-Trainingsdienst automatisch über die vordefinierte Umgebungsvariable AIP_TENSORBOARD_LOG_DIR verfügbar gemacht wird.

Dazu können Sie in der Regel os.environ['AIP_TENSORBOARD_LOG_DIR'] als Logverzeichnis für die Open-Source-TensorBoard-APIs zum Schreiben von Logs bereitstellen. Der Speicherort von AIP_TENSORBOARD_LOG_DIR wird normalerweise mit der Variable staging_bucket festgelegt.

Erstellen Sie zum Konfigurieren Ihres Trainingsskripts in TensorFlow 2.x einen TensorBoard-Callback und legen Sie die log_dir Variable auf os.environ['AIP_TENSORBOARD_LOG_DIR'] fest. Der TensorBoard-Callback ist dann in der TensorFlow model.fit Callback-Liste enthalten.

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

Weitere Informationen darüber, wie Vertex AI Umgebungsvariablen in Ihrer benutzerdefinierten Trainingsumgebung festlegt, erhalten Sie in diesem Artikel.

Pipeline erstellen und ausführen

Das folgende Beispiel zeigt, wie Sie eine Pipeline mit dem Kubeflow Pipelines-DSL-Paket erstellen und ausführen. Weitere Beispiele und zusätzliche Informationen finden Sie in der Dokumentation zu Vertex AI Pipelines.

Trainingskomponente erstellen

Verpacken Sie Ihren Trainingscode in eine benutzerdefinierte Komponente und achten Sie darauf, dass der Code so konfiguriert ist, dass TensorBoard-Logs in einen Cloud Storage-Bucket geschrieben werden. Weitere Beispiele finden Sie unter Eigene Pipeline-Komponenten erstellen.

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 erstellen und kompilieren

Erstellen Sie einen benutzerdefinierten Trainingsjob aus der von Ihnen erstellten Komponente, indem Sie die Komponentenspezifikation in create_custom_training_job_op_from_component angeben. Setzen Sie tensorboard_resource_name auf Ihre TensorBoard-Instanz und staging_bucket auf den Ort, an dem Artefakte bei API-Aufrufen (einschließlich TensorBoard-Logs) bereitgestellt werden.

Erstellen Sie dann eine Pipeline, die diesen Job enthält, und kompilieren Sie die Pipeline in eine JSON-Datei.

Weitere Beispiele und Informationen finden Sie unter Benutzerdefinierte Jobkomponenten und Pipeline erstellen.

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-Pipeline senden

Senden Sie die Pipeline mit dem Vertex AI SDK für Python. Weitere Informationen finden Sie unter Pipeline ausführen.

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: Geben Sie einen Namen für den Test an.
  • pipeline_job_display_name: Der Anzeigename für den Pipelinejob.
  • template_path: Der Pfad zur kompilierten Pipelinevorlage.
  • pipeline_root: Geben Sie einen Cloud Storage-URI an, auf den das Pipelines-Dienstkonto zugreifen kann. Die Artefakte Ihrer Pipelineausführungen werden im Pipeline-Stammverzeichnis gespeichert.
  • parameter_values: Die Pipelineparameter, die an diese Ausführung übergeben werden sollen. Erstellen Sie beispielsweise ein dict(), das die Parameternamen als Wörterbuchschlüssel und die Parameterwerte als Wörterbuchwerte enthält.
  • project: Ihre Projekt-ID. Das Google Cloud-Projekt, in dem die Pipeline ausgeführt wird. Sie finden Ihre IDs auf der Begrüßungsseite der Google Cloud Console.
  • location: Der Standort, an dem die Pipeline ausgeführt wird. Dies sollte derselbe Standort sein wie der der TensorBoard-Instanz, die Sie verwenden.

Nächste Schritte