Usa Vertex AI TensorBoard con Vertex AI Pipelines

Tu código de entrenamiento se puede empaquetar en un componente de entrenamiento personalizado y ejecutarse en un trabajo de canalización. Los registros de TensorBoard se transmiten de forma automática a tu experimento de Vertex AI TensorBoard. Puedes usar esta integración para supervisar tu entrenamiento casi en tiempo real, ya que Vertex AI TensorBoard permite ver los registros de Vertex AI TensorBoard a medida que se escriben en Cloud Storage.

Para la configuración inicial, consulta Configura Vertex AI TensorBoard.

Cambios en la secuencia de comandos de entrenamiento

La secuencia de comandos de entrenamiento debe configurarse para escribir registros de TensorBoard en el bucket de Cloud Storage, cuya ubicación pondrá a disposición Vertex AI Training Service de forma automática a través de una variable de entorno predefinida AIP_TENSORBOARD_LOG_DIR.

Por lo general, esto se puede hacer si se proporciona os.environ['AIP_TENSORBOARD_LOG_DIR'] como directorio de registro a las APIs de escritura de registros de TensorBoard de código abierto. Por lo general, la ubicación de AIP_TENSORBOARD_LOG_DIR se establece con la variable staging_bucket.

Para configurar la secuencia de comandos de entrenamiento en TensorFlow 2.x, crea una devolución de llamada de TensorBoard y establece la variable log_dir en os.environ['AIP_TENSORBOARD_LOG_DIR']La devolución de llamada de TensorBoard se incluye en la lista de devoluciones de llamadas model.fit de 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],
  )

Obtén más información sobre cómo Vertex AI establece variables de entorno en el entorno de entrenamiento personalizado.

Compila y ejecuta una canalización

En el siguiente ejemplo, se muestra cómo compilar y ejecutar una canalización con el paquete DSL de Kubeflow Pipelines. Para obtener más ejemplos y detalles adicionales, consulta la documentación de Vertex AI Pipelines.

Crea un componente de entrenamiento

Empaqueta tu código de entrenamiento en un componente personalizado y asegúrate de que el código esté configurado para escribir registros de TensorBoard en un bucket de Cloud Storage. Para obtener más ejemplos, consulta Compila tus propios componentes de canalización.

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 y compila una canalización

Crea un trabajo de entrenamiento personalizado a partir del componente que creaste a través de la especificación de la especificación del componente en create_custom_training_job_op_from_component. Establece tensorboard_resource_name en tu instancia de TensorBoard y staging_bucket en la ubicación para almacenar en etapa intermedia los artefactos durante las llamadas a la API (incluidos los registros de TensorBoard).

Luego, compila una canalización para incluir este trabajo y compila la canalización en un archivo JSON.

Para obtener más información y ejemplos, consulta Componentes de trabajo personalizado y Compila una canalización.

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

Envía una canalización de Vertex AI

Envía tu canalización con el SDK de Vertex AI para Python. Para obtener más información, consulta Ejecuta una canalización.

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: Proporciona un nombre para tu experimento.
  • pipeline_job_display_name: el nombre visible del trabajo de canalización.
  • template_path: la ruta a la plantilla de canalización compilada.
  • pipeline_root: Especifica un URI de Cloud Storage al que pueda acceder la cuenta de servicio de tus canalizaciones. Los artefactos de las ejecuciones de tus canalizaciones se almacenan en la raíz de la canalización.
  • parameter_values: los parámetros de canalización que se pasarán a esta ejecución. Por ejemplo, crea un dict() con los nombres de los parámetros como claves del diccionario y los valores de los parámetros como valores del diccionario.
  • project: El ID del proyecto. El proyecto de Google Cloud en el que se ejecutará la canalización. Puedes encontrar los IDs en la página de bienvenida de la consola de Google Cloud.
  • location: la ubicación en la que se ejecutará la canalización. Debe ser la misma ubicación que la instancia de TensorBoard que usas.

¿Qué sigue?