使用 Vertex AI TensorBoard 直观呈现作业

如果您对 Gemini Enterprise Agent Platform Managed Training 感兴趣，请与您的销售代表联系以获取访问权限。

借助托管式训练，您可以使用 Vertex AI TensorBoard 近乎实时地直观呈现训练日志。只需将工作负载配置为将日志保存到 Cloud Storage 存储桶，日志就会自动流式传输到 TensorBoard 界面以供分析。

前提条件

在开始之前，请确保您具备以下条件：

  • 正在运行的托管式训练集群。
  • 用于存储 TensorBoard 日志的 Cloud Storage 存储桶。此存储桶必须与您的 TensorBoard 实例位于同一区域。如需查看设置说明，请参阅创建 Cloud Storage 存储桶。
  • Vertex AI TensorBoard 实例。如需查看创建说明， 请参阅创建 Vertex AI TensorBoard 实例。
  • 正确的 IAM 权限。如需允许 Cloud Storage FUSE 从存储桶读取数据和向存储桶写入数据，集群虚拟机使用的服务账号需要具有 Storage Object User (roles/storage.objectUser) 角色。

启用 Tensorboard 上传

如需为作业配置 TensorBoard 集成，请在 Slurm 作业提交中使用 --extra 标志传递以下参数：

  • tensorboard_base_output_dir：指定要将日志上传到的 Cloud Storage 路径。例如 gs://my-bucket/my-logs。

  • tensorboard_url：指定 Vertex AI TensorBoard 实例、实验或运行网址。如果仅提供实例，系统会创建新的实验和运行。如果省略，则系统会使用项目的默认 TensorBoard 实例。例如 projects/123/locations/us-central1/tensorboards/456。

示例

# 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

从训练作业写入日志

在训练脚本中，访问 AIP_TENSORBOARD_LOG_DIR 环境变量。此变量提供唯一的 Cloud Storage 路径，脚本应将 TensorBoard 日志写入到该路径。

路径遵循以下结构：

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

以下示例展示了一个完整的工作流，其中包含两个关键组件：用于配置作业的 Slurm 提交脚本，以及用于读取环境变量以写入日志的 Python 训练脚本。

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

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

实时日志同步

如需直观呈现正在运行的作业中的指标，您必须定期关闭并重新创建训练代码中的摘要写入器。这是必要的，因为 gcsfuse 仅在日志文件关闭后才会将其同步到 Cloud Storage。这种“刷新”技术可确保在作业完成之前，中间结果在 TensorBoard 控制台中可见。

查看 Vertex AI TensorBoard

提交作业后，您可以在 Google Cloud 控制台中前往 Vertex AI Experiments 页面，监控作业的进度。

后续步骤

在 Vertex AI TensorBoard 中分析训练指标、可视化模型图表并进行性能剖析，以调试问题、优化模型，并为将模型集成到完整的 MLOps 工作流中做好准备。