Usar Vertex AI TensorBoard con Vertex AI Pipelines

Tu código de entrenamiento se puede empaquetar en un componente de entrenamiento personalizado y ejecutarlo en una tarea de flujo de trabajo. Los registros de TensorBoard se envían automáticamente a tu experimento de Vertex AI TensorBoard. Puedes usar esta integración para monitorizar tu entrenamiento casi en tiempo real, ya que Vertex AI TensorBoard transmite los registros de Vertex AI TensorBoard a medida que se escriben en Cloud Storage.

Para la configuración inicial, consulta el artículo Configurar Vertex AI TensorBoard.

Cambios en tu guion de entrenamiento

Tu secuencia de comandos de entrenamiento debe configurarse para escribir registros de TensorBoard en el contenedor de Cloud Storage, cuya ubicación el servicio de entrenamiento de Vertex AI pondrá automáticamente a disposición a través de una variable de entorno predefinida AIP_TENSORBOARD_LOG_DIR.

Para ello, normalmente 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. La ubicación de AIP_TENSORBOARD_LOG_DIR se suele definir con la variable staging_bucket.

Para configurar tu script de entrenamiento en TensorFlow 2.x, crea una retrollamada de TensorBoard y asigna el valor os.environ['AIP_TENSORBOARD_LOG_DIR'] a la variable log_dir. La retrollamada de TensorBoard se incluye en la lista de retrollamadas de TensorFlow model.fit.

  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],
  )

Consulta más información sobre cómo Vertex AI define las variables de entorno en tu entorno de entrenamiento personalizado.

Crear y ejecutar un flujo de procesamiento

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

Crear un componente de entrenamiento

Empaqueta el código de entrenamiento en un componente personalizado y asegúrate de que esté configurado para escribir registros de TensorBoard en un segmento de Cloud Storage. Para ver más ejemplos, consulta Crea tus propios componentes de flujo de procesamiento.

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],
    )

Compilar un flujo de procesamiento

Crea una tarea de entrenamiento personalizada a partir del componente que has creado especificando la especificación del componente en create_custom_training_job_op_from_component. Define tensorboard_resource_name en tu instancia de TensorBoard y staging_bucket en la ubicación en la que se almacenarán los artefactos durante las llamadas a la API (incluidos los registros de TensorBoard).

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

Para ver más ejemplos e información, consulta Componentes de trabajos personalizados y Crear 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"
    )

Enviar un flujo de procesamiento de Vertex AI

Envía tu flujo de trabajo con el SDK de Vertex AI para Python. Para obtener más información, consulta Ejecutar 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: ponle un nombre al experimento.
  • pipeline_job_display_name: nombre visible del trabajo de flujo de procesamiento.
  • template_path: ruta a la plantilla de flujo de procesamiento 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 tu canalización se almacenan en la raíz de la canalización.
  • parameter_values: los parámetros de la canalización que se van a pasar 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 proyecto de Google Cloud en el que se ejecutará la canalización. Puedes encontrar tus IDs en la página Bienvenido de la Google Cloud consola.
  • location: la ubicación en la que se ejecutará la canalización. Debe ser la misma ubicación que la instancia de TensorBoard que estés usando.

Siguientes pasos