Utilizzare Vertex AI TensorBoard con Vertex AI Pipelines

Il codice di addestramento può essere pacchettizzato in un componente di addestramento personalizzato ed eseguito in un job della pipeline. I log di TensorBoard vengono trasmessi automaticamente in streaming Esperimento Vertex AI TensorBoard. Puoi utilizzare questa integrazione per monitorare l'addestramento quasi in tempo reale mentre Vertex AI TensorBoard genera stream nei log di Vertex AI TensorBoard man mano che vengono scritti in Cloud Storage.

Per la configurazione iniziale, vedi Configura Vertex AI TensorBoard.

Modifiche allo script di allenamento

Lo script di addestramento deve essere configurato per scrivere i log di TensorBoard nell'interfaccia Bucket Cloud Storage, la località in cui viene eseguito il Vertex AI Training Il servizio sarà automaticamente disponibile tramite un ambiente predefinito la variabile AIP_TENSORBOARD_LOG_DIR.

In genere, questo può essere fatto specificando os.environ['AIP_TENSORBOARD_LOG_DIR'] come directory dei log per le API di scrittura dei log di TensorBoard open source. La posizione del AIP_TENSORBOARD_LOG_DIR viene in genere impostata con la variabile staging_bucket.

Per configurare lo script di addestramento in TensorFlow 2.x, crea un callback di TensorBoard e imposta la variabile log_dir su os.environ['AIP_TENSORBOARD_LOG_DIR']. Il callback di TensorBoard viene poi incluso nell'elenco dei callback model.fit di TensorFlow.

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

Scopri di più su come Vertex AI imposta le variabili di ambiente nel tuo ambiente di addestramento personalizzato.

Crea ed esegui una pipeline

L'esempio seguente mostra come creare ed eseguire una pipeline utilizzando il pacchetto DSL Kubeflow Pipelines. Per altri esempi e dettagli, consulta la documentazione di Vertex AI Pipelines.

Creare un componente di addestramento

Impacchetta il codice di addestramento in un componente personalizzato, assicurandoti che sia configurato per scrivere i log di TensorBoard in un bucket Cloud Storage. Per altri esempi, consulta Creare i tuoi componenti di pipeline.

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

Crea e compila una pipeline

Crea un job di addestramento personalizzato dal componente che hai creato specificando la specifica del componente in create_custom_training_job_op_from_component. Imposta tensorboard_resource_name sull'istanza TensorBoard e staging_bucket sulla posizione in cui eseguire il commit degli elementi durante le chiamate all'API (inclusi i log di TensorBoard).

Quindi, crea una pipeline per includere questo job e compila la pipeline in JSON.

Per ulteriori esempi e informazioni, vedi Componenti personalizzati dei job e Crea una pipeline.

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

Invia una pipeline Vertex AI

Invia la pipeline utilizzando l'SDK Vertex AI per Python. Per ulteriori informazioni, consulta Eseguire una pipeline.

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)

Passaggi successivi

• Visualizza i risultati: Visualizza TensorBoard per Vertex AI Pipelines.
• Scopri come ottimizzare le prestazioni dei tuoi job di addestramento personalizzato utilizzando Cloud Profiler.