Visualiza trabajos con Vertex AI TensorBoard

Si te interesa el entrenamiento administrado de Gemini Enterprise Agent Platform, comunícate con tu representante de ventas para obtener acceso.

Con el entrenamiento administrado, puedes visualizar tus registros de entrenamiento casi en tiempo real con Vertex AI TensorBoard. Solo debes configurar tu carga de trabajo para guardar los registros en un bucket de Cloud Storage, y se transmitirán automáticamente a la interfaz de TensorBoard para su análisis.

Requisitos previos

Antes de comenzar, asegúrate de tener lo siguiente:

  • Un clúster de entrenamiento administrado en ejecución
  • Un bucket de Cloud Storage para almacenar tus registros de TensorBoard (este bucket debe estar en la misma región que tu instancia de TensorBoard) (para obtener instrucciones de configuración, consulta Crea un bucket de Cloud Storage)
  • Una instancia de Vertex AI TensorBoard (para obtener instrucciones de creación, consulta Crea una instancia de Vertex AI TensorBoard)
  • Los permisos de IAM correctos (para permitir que Cloud Storage FUSE lea y escriba en el bucket de almacenamiento, la cuenta de servicio que usan las VMs de tu clúster requiere el rol Usuario de objetos de almacenamiento (roles/storage.objectUser))

Cómo habilitar la carga de TensorBoard

Para configurar la integración de TensorBoard para tu trabajo, pasa los siguientes argumentos con la marca --extra en el envío de tu trabajo de Slurm:

  • tensorboard_base_output_dir: Especifica la ruta de acceso de Cloud Storage para subir registros. Por ejemplo, gs://my-bucket/my-logs.

  • tensorboard_url: Especifica la URL de la instancia, el experimento o la ejecución de Vertex AI TensorBoard. Si solo se proporciona una instancia, se crean un experimento y una ejecución nuevos. Si se omite, se usa la instancia de TensorBoard predeterminada para el proyecto. Por ejemplo, projects/123/locations/us-central1/tensorboards/456.

Ejemplo

# Using specific tensorboard instance
sbatch --extra="tensorboard_base_output_dir=<your-cloud-storage-dir>,tensorboard_url=projects/<project-id>/locations/<location>/tensorboards/<tensorboard-instance-id>" your_script.sbatch

Cómo escribir registros desde tu trabajo de entrenamiento

Dentro de tu secuencia de comandos de entrenamiento, accede a la variable de entorno AIP_TENSORBOARD_LOG_DIR. Esta variable proporciona la ruta de acceso única de Cloud Storage en la que tu secuencia de comandos debe escribir sus registros de TensorBoard.

La ruta de acceso sigue esta estructura:

gs://<your-cloud-storage-path>/<cluster-id>-<cluster-uuid>/tensorboard/job-<job-id>/

En el siguiente ejemplo, se muestra un flujo de trabajo completo con dos componentes clave: la secuencia de comandos de envío de Slurm que configura el trabajo y la secuencia de comandos de entrenamiento de Python que lee la variable de entorno para escribir sus registros.

Secuencia de comandos de trabajo de Slurm (simple_job.sbatch):

#!/bin/bash
#SBATCH --job-name=tensorboard-simple-test
#SBATCH --output=tensorboard-simple-test-%j.out

# Activate your Python virtual environment if needed
# source /path/to/your/venv/bin/activate
python3 simple_logger.py

Secuencia de comandos de Python (simple_logger.py):

import tensorflow as tf
import os

# Get the log directory from the environment variable
log_dir = os.environ.get("AIP_TENSORBOARD_LOG_DIR")

print(f"Writing TensorBoard logs to: {log_dir}")
writer = tf.summary.create_file_writer(log_dir)

with writer.as_default():
    for step in range(10):
        # Simulate some metrics
        loss = 1.0 - (step * 0.1)
        accuracy = 0.6 + (step * 0.04)

        # Log the metrics
        tf.summary.scalar('loss', loss, step=step)
        tf.summary.scalar('accuracy', accuracy, step=step)
        writer.flush()
        print(f"Step {step}: loss={loss:.4f}, accuracy={accuracy:.4f}")

writer.close()
print(f"--- Finished writing metrics to {log_dir} ---")

Sincronización de registros en tiempo real

Para visualizar las métricas de un trabajo en ejecución, debes cerrar y volver a crear periódicamente el escritor de resumen en tu código de entrenamiento. Esto es necesario porque gcsfuse solo sincroniza los archivos de registro con Cloud Storage una vez que se cierran. Esta técnica de "vaciado" garantiza que los resultados intermedios sean visibles en la consola de TensorBoard antes de que se complete el trabajo.

Cómo ver Vertex AI TensorBoard

Una vez que se envía el trabajo, puedes supervisar su progreso en la página Experimentos de Vertex AI en la Cloud de Confiance consola.

¿Qué sigue?

Analiza las métricas de entrenamiento, visualiza los gráficos de modelos y genera perfiles de rendimiento en Vertex AI TensorBoard para depurar problemas, optimizar tu modelo y prepararlo para la integración en un flujo de trabajo completo de MLOps.