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 al tuo esperimento Vertex AI TensorBoard. Puoi utilizzare questa integrazione per monitorare l'addestramento quasi in tempo reale mentre Vertex AI TensorBoard trasmette i flussi di dati nei log di Vertex AI TensorBoard mentre vengono scritti in Cloud Storage.

Per la configurazione iniziale, consulta Configurazione per Vertex AI TensorBoard.

Modifiche allo script di addestramento

Lo script di addestramento deve essere configurato per scrivere i log di 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 questa operazione può essere eseguita fornendo os.environ['AIP_TENSORBOARD_LOG_DIR'] come directory di log alle API di scrittura di log open source TensorBoard. In genere, la posizione dell'elemento AIP_TENSORBOARD_LOG_DIR viene 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 TensorBoard verrà quindi 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 ulteriori esempi e dettagli aggiuntivi, consulta la documentazione di Vertex AI Pipelines.

crea un componente di addestramento

Pacchettizza il tuo codice di addestramento in un componente personalizzato, assicurandoti che il codice sia configurato per scrivere i log di TensorBoard in un bucket Cloud Storage. Per altri esempi, consulta Creare componenti della 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 creato specificando le specifiche del componente in create_custom_training_job_op_from_component. Imposta tensorboard_resource_name sull'istanza TensorBoard e staging_bucket sulla posizione per lo stage degli artefatti durante le chiamate API (inclusi i log di TensorBoard).

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

Per ulteriori esempi e informazioni, consulta Componenti di 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"
    )

Invia una pipeline Vertex AI

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

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: assegna un nome all'esperimento.
  • pipeline_job_display_name: il nome visualizzato del job della pipeline.
  • template_path: il percorso del modello di pipeline compilato.
  • pipeline_root: specifica un URI Cloud Storage a cui può accedere il tuo account di servizio delle pipeline. Gli artefatti delle esecuzioni della pipeline sono archiviati all'interno della radice della pipeline.
  • parameter_values: i parametri della pipeline da passare a questa esecuzione. Ad esempio, crea un elemento dict() con i nomi dei parametri come chiavi del dizionario e i valori dei parametri come valori del dizionario.
  • project: il tuo ID progetto. Il progetto Google Cloud in cui eseguire la pipeline. Puoi trovare i tuoi ID nella pagina di benvenuto della console Google Cloud.
  • location: la località in cui eseguire la pipeline. Questa dovrebbe essere la stessa posizione dell'istanza TensorBoard in uso.

Passaggi successivi