使用 Vertex AI TensorBoard 視覺化工作

如要使用 Gemini Enterprise Agent Platform Managed Training,請與銷售代表聯絡以取得存取權。

使用 Managed Training 時,您可以透過 Vertex AI TensorBoard,近乎即時地將訓練記錄視覺化。只要將工作負載設定為將記錄儲存至 Cloud Storage bucket,系統就會自動將記錄串流至 TensorBoard 介面進行分析。

必要條件

開始之前,請確認您具備以下項目:

  • 正在執行的 Managed Training 叢集。
  • 用來儲存 TensorBoard 記錄的 Cloud Storage bucket。這個 bucket 必須與 TensorBoard 執行個體位於相同區域。如需設定操作說明,請參閱「建立 Cloud Storage bucket」。
  • Vertex AI TensorBoard 執行個體。如需建立操作說明,請參閱「建立 Vertex AI TensorBoard 執行個體」。
  • 正確的 IAM 權限。如要允許 Cloud Storage FUSE 從儲存空間 bucket 讀取及寫入資料,叢集 VM 使用的服務帳戶必須具備「Storage 物件使用者」角色 (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

提交工作後,請前往 Cloud de Confiance 控制台的 Vertex AI Experiments 頁面,監控工作進度。

後續步驟

在 Vertex AI TensorBoard 中分析訓練指標、將模型圖表視覺化,以及剖析效能,藉此偵錯問題、最佳化模型,並準備將模型整合至完整的機器學習運作工作流程。