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 in streaming automaticamente all'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, consulta Configurazione per Vertex AI TensorBoard.

Modifiche allo script di allenamento

Lo script di addestramento deve essere configurato per scrivere i log TensorBoard nel bucket Cloud Storage, la località in cui il servizio Vertex AI Training metterà automaticamente a disposizione tramite una variabile di ambiente predefinita AIP_TENSORBOARD_LOG_DIR.

In genere, questo può essere fatto fornendo os.environ['AIP_TENSORBOARD_LOG_DIR'] come directory dei log alle 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.

Creare ed eseguire 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],
    )

Creare e compilare 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 compilala in un file JSON.

Per altri esempi e informazioni, consulta Componenti dei job personalizzati e Creare 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"
    )

Inviare 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 il rendimento dei job di addestramento personalizzati utilizzando Cloud Profiler.