Utilizzare Vertex AI TensorBoard con le pipeline della piattaforma di agenti Gemini Enterprise

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 all'esperimento di Vertex AI TensorBoard. Puoi utilizzare questa integrazione per monitorare l'addestramento quasi in tempo reale man mano che Vertex AI TensorBoard trasmette in streaming i log di Vertex AI TensorBoard mentre vengono scritti in Cloud Storage.

Per la configurazione iniziale, vedi Configurare 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 cui posizione verrà resa disponibile automaticamente dal servizio di addestramento di Agent Platform Gemini Enterprise tramite una variabile di ambiente predefinita AIP_TENSORBOARD_LOG_DIR.

In genere, puoi farlo fornendo os.environ['AIP_TENSORBOARD_LOG_DIR'] come directory dei log alle API di scrittura dei log di TensorBoard open source. La posizione di 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 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 Agent Platform imposta le variabili di ambiente nell'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 aggiuntivi, consulta la documentazione sulle pipeline di Agent Platform Gemini Enterprise.

Creare un componente di addestramento

Pacchettizza il 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, vedi Creare componenti della pipeline personalizzati.

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 di TensorBoard e staging_bucket sulla posizione in cui eseguire lo staging 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 altri esempi e informazioni, vedi 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 di Agent Platform Gemini Enterprise

Invia la pipeline utilizzando l'SDK Agent Platform Python. Per saperne di più, 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: fornisci un nome per l'esperimento.
  • pipeline_job_display_name: il nome visualizzato del job della pipeline.
  • template_path: il percorso del modello della pipeline compilato.
  • pipeline_root: specifica un URI Cloud Storage a cui può accedere l'account di servizio delle pipeline. Gli artefatti delle esecuzioni della pipeline vengono archiviati nella radice della pipeline.
  • parameter_values: i parametri della pipeline da passare a questa esecuzione. Ad esempio, crea un dict() con i nomi dei parametri come chiavi del dizionario e i valori dei parametri come valori del dizionario.
  • project: . Il progetto Google Cloud in cui eseguire la pipeline. Puoi trovare i tuoi ID nella Google Cloud console pagina di benvenuto.
  • location: la località in cui eseguire la pipeline. Deve essere la stessa località dell'istanza di TensorBoard che stai utilizzando.

Passaggi successivi