Vertex AI TensorBoard mit Vertex AI Pipelines verwenden

Der Trainingscode kann in einer benutzerdefinierten Trainingskomponente verpackt und in einem Pipeline-Job 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.

Eine Anleitung zur Ersteinrichtung finden Sie unter Einrichtung für Vertex AI TensorBoard.

Ä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 der AIP_TENSORBOARD_LOG_DIR wird in der Regel mit der Variablen 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

Im folgenden Beispiel wird gezeigt, wie Sie eine Pipeline mit dem Kubeflow Pipelines-DSL-Paket erstellen und ausführen. Weitere Beispiele und weitere Informationen finden Sie in der Dokumentation zu Vertex AI Pipelines.

Trainingskomponente erstellen

Verpacken Sie Ihren Trainingscode in einer benutzerdefinierten Komponente und achten Sie darauf, dass der Code so konfiguriert ist, dass TensorBoard-Protokolle 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. Legen Sie tensorboard_resource_name auf Ihre TensorBoard-Instanz und staging_bucket auf den Speicherort fest, an dem Artefakte während API-Aufrufen bereitgestellt werden sollen (einschließlich TensorBoard-Protokolle).

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 einreichen

Reichen Sie Ihre Pipeline mit dem Vertex AI SDK für Python ein. 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 Pipeline-Vorlage.
  • 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 werden soll. 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