Utiliser Vertex AI TensorBoard avec Vertex AI Pipelines

Votre code d'entraînement peut être empaqueté dans un composant d'entraînement personnalisé et s'exécuter dans un job de pipeline. Les journaux TensorBoard sont automatiquement diffusés vers votre test Vertex AI TensorBoard. Vous pouvez utiliser cette intégration pour surveiller votre entraînement quasiment en temps réel, car Vertex AI TensorBoard diffuse les journaux Vertex AI TensorBoard à mesure qu'ils sont écrits dans Cloud Storage.

Pour la configuration initiale, consultez la page Configurer Vertex AI TensorBoard.

Modifier le script d'entraînement

Votre script d'entraînement doit être configuré pour écrire des journaux TensorBoard dans le bucket Cloud Storage, dont l'emplacement sera automatiquement mis à disposition par le service d'entraînement Vertex AI via une variable d'environnement prédéfinie AIP_TENSORBOARD_LOG_DIR.

Pour ce faire, vous devez généralement fournir os.environ['AIP_TENSORBOARD_LOG_DIR'] comme répertoire de journal aux API d'écriture de journaux TensorBoard Open Source. L'emplacement de AIP_TENSORBOARD_LOG_DIR est généralement défini avec la variable staging_bucket.

Pour configurer votre script d'entraînement dans TensorFlow 2.x, créez un rappel TensorBoard et définissez la variable log_dir sur os.environ['AIP_TENSORBOARD_LOG_DIR']. Le rappel TensorBoard est ensuite inclus dans la liste des rappels model.fit 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],
  )
  

Découvrez comment Vertex AI définit les variables d'environnement dans votre environnement d'entraînement personnalisé.

Créer et exécuter un pipeline

L'exemple suivant montre comment créer et exécuter un pipeline à l'aide du package DSL de Kubeflow Pipelines. Pour plus d'exemples et d'informations, consultez la documentation sur Vertex AI Pipelines.

Créer un composant d'entraînement

Empaquetez votre code d'entraînement dans un composant personnalisé, en vous assurant que le code est configuré pour écrire des journaux TensorBoard dans un bucket Cloud Storage. Pour plus d'exemples, consultez la section Créer vos propres composants de 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],
    )

Créer et compiler un pipeline

Créez un job d'entraînement personnalisée à partir du composant que vous avez créé en spécifiant la spécification de composant dans create_custom_training_job_op_from_component. Définissez tensorboard_resource_name sur votre instance TensorBoard et staging_bucket sur l'emplacement pour organiser les artefacts lors des appels d'API (y compris les journaux TensorBoard).

Ensuite, créez un pipeline pour inclure ce job et compilez le pipeline dans un fichier JSON.

Pour obtenir plus d'exemples et d'informations, consultez les pages Composants de jobs personnalisés et Créer un 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"
    )

Envoyer un pipeline Vertex AI

Envoyez votre pipeline à l'aide du SDK Vertex AI pour Python. Pour en savoir plus, consultez la page Exécuter un pipeline.

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)

Étapes suivantes

• Affichez vos résultats : Afficher le TensorBoard pour Vertex AI Pipelines.
• Découvrez comment optimiser les performances de vos tâches d'entraînement personnalisé à l'aide de Vertex AI TensorBoard Profiler.