将 Vertex AI TensorBoard 与自定义训练搭配使用

使用自定义训练来训练模型时，您可以设置训练作业以自动将 Vertex AI TensorBoard 日志上传到 Vertex AI TensorBoard。

您可以使用此集成来近乎实时地监控训练，因为 Vertex AI TensorBoard 会在 Vertex AI TensorBoard 日志写入 Cloud Storage 时流式传输这些日志。

对训练脚本的更改

您的训练脚本必须配置为将 TensorBoard 日志写入 Cloud Storage 存储桶，即 Vertex AI Training 服务通过预定义的环境变量 AIP_TENSORBOARD_LOG_DIR 自动提供的位置。

这通常可以通过将 os.environ['AIP_TENSORBOARD_LOG_DIR'] 作为日志目录提供给开源 TensorBoard 日志写入 API 来实现。 AIP_TENSORBOARD_LOG_DIR 的位置通常使用 staging_bucket 变量进行设置。

如需在 TensorFlow 2.x 中配置训练脚本，请创建 TensorBoard 回调并将 log_dir 变量设置为 os.environ['AIP_TENSORBOARD_LOG_DIR']。然后，TensorBoard 回调会添加到 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],
  )

详细了解 Vertex AI 如何在自定义训练环境中设置环境变量。

创建自定义训练作业

以下示例展示了如何创建您自己的自定义训练作业。

如需查看有关如何创建自定义训练作业的详细示例，请参阅 Hello 自定义训练。如需了解构建自定义训练容器的步骤，请参阅创建用于训练的自定义容器映像。

如需创建自定义训练作业，请使用 Vertex AI SDK for Python 或 Vertex AI SDK for REST。

Python

def create_training_pipeline_custom_job_sample(
    project: str,
    location: str,
    staging_bucket: str,
    display_name: str,
    script_path: str,
    container_uri: str,
    model_serving_container_image_uri: str,
    dataset_id: Optional[str] = None,
    model_display_name: Optional[str] = None,
    args: Optional[List[Union[str, float, int]]] = None,
    replica_count: int = 0,
    machine_type: str = "n1-standard-4",
    accelerator_type: str = "ACCELERATOR_TYPE_UNSPECIFIED",
    accelerator_count: int = 0,
    training_fraction_split: float = 0.8,
    validation_fraction_split: float = 0.1,
    test_fraction_split: float = 0.1,
    sync: bool = True,
    tensorboard_resource_name: Optional[str] = None,
    service_account: Optional[str] = None,
):
    aiplatform.init(project=project, location=location, staging_bucket=staging_bucket)

    job = aiplatform.CustomTrainingJob(
        display_name=display_name,
        script_path=script_path,
        container_uri=container_uri,
        model_serving_container_image_uri=model_serving_container_image_uri,
    )

    # This example uses an ImageDataset, but you can use another type
    dataset = aiplatform.ImageDataset(dataset_id) if dataset_id else None

    model = job.run(
        dataset=dataset,
        model_display_name=model_display_name,
        args=args,
        replica_count=replica_count,
        machine_type=machine_type,
        accelerator_type=accelerator_type,
        accelerator_count=accelerator_count,
        training_fraction_split=training_fraction_split,
        validation_fraction_split=validation_fraction_split,
        test_fraction_split=test_fraction_split,
        sync=sync,
        tensorboard=tensorboard_resource_name,
        service_account=service_account,
    )

    model.wait()

    print(model.display_name)
    print(model.resource_name)
    print(model.uri)
    return model

project：您的项目 ID。您可以在 Google Cloud 控制台的欢迎页面中找到这些 ID。
location：要在其中运行 CustomJob 的位置。这应是提供的 TensorBoard 实例所在的位置。
staging_bucket：要在 API 调用期间暂存工件（包括 TensorBoard 日志）的 Cloud Storage 存储桶。
display_name：自定义训练作业的显示名称。
script_path：相对于本地文件系统上的工作目录的路径，指向作为训练代码入口点的脚本。
container_uri：训练容器映像的 URI 可以是 Vertex AI 预构建训练容器或自定义容器。
model_serving_container_image_uri：适合为训练脚本生成的模型提供服务的模型服务容器的 URI。
dataset_id：用于训练的数据集的 ID 编号。
model_display_name：经过训练的模型的显示名称。
args：要传递给 Python 脚本的命令行参数。
replica_count：要使用的工作器副本的数量。在大多数情况下，对于第一个工作器池，请将此值设置为 1。
machine_type：要使用的虚拟机类型。如需查看受支持的虚拟机列表，请参阅机器类型。
accelerator_type：要挂接到资源池中每个虚拟机的 GPU 的类型。如需查看受支持的 GPU 列表，请参阅 GPU。
accelerator_count：要挂接到资源池中每个虚拟机的 GPU 的数量。默认值为 1。
training_fraction_split：用于训练模型的数据集比例。
validation_fraction_split：用于验证模型的数据集比例。
test_fraction_split：用于评估模型的数据集比例。
sync：是否同步执行此方法。
tensorboard_resource_name：CustomJob 将 TensorBoard 日志上传到的 Vertex TensorBoard 实例的资源名称。
service_account：使用 TensorBoard 运行时为必需。请参阅创建具有所需权限的服务账号。

REST

在使用任何请求数据之前，请先进行以下替换：

LOCATION_ID：要在其中运行 CustomJob 的位置，例如 us-central1。这应是提供的 TensorBoard 实例所在的位置。
PROJECT_ID：您的项目 ID。
TENSORBOARD_INSTANCE_NAME：（必需）用来存储 Vertex AI TensorBoard 日志的现有 Vertex AI TensorBoard 实例的全名：
projects/PROJECT_ID/locations/LOCATION_ID/tensorboards/TENSORBOARD_INSTANCE_ID
注意：如果 TensorBoard 实例不是现有实例，则创建自定义作业会抛出一个 404 错误。
GCS_BUCKET_NAME："${PROJECT_ID}-tensorboard-logs-${LOCATION}"
USER_SA_EMAIL：（必需）在先前步骤中创建的服务账号，或您自己的服务账号。"USER_SA_NAME@${PROJECT_ID}.iam.gserviceaccount.com"
TRAINING_CONTAINER：TRAINING_CONTAINER。
INVOCATION_TIMESTAMP："$(date +'%Y%m%d-%H%M%S')"
JOB_NAME："tensorboard-example-job-${INVOCATION_TIMESTAMP}"
BASE_OUTPUT_DIR：（必需）要将训练的所有输出写入到的 Google Cloud 路径。"gs://$GCS_BUCKET_NAME/$JOB_NAME"

HTTP 方法和网址：

POST https://LOCATION_ID-aiplatform.googleapis.com/v1/projects/PROJECT_ID/locations/LOCATION_ID/customJobs

请求 JSON 正文：

{
"displayName": JOB_NAME,
"jobSpec":{
"workerPoolSpecs":[
  {
    "replicaCount": "1",
     "machineSpec": {
        "machineType": "n1-standard-8",
      },
      "containerSpec": {
        "imageUri": TRAINING_CONTAINER,
      }
    }
  ],
  
  "base_output_directory": {
  "output_uri_prefix": BASE_OUTPUT_DIR,
   },
  "serviceAccount": USER_SA_EMAIL,
  "tensorboard": TENSORBOARD_INSTANCE_NAME,
  }
}

如需发送您的请求，请展开以下选项之一：

curl（Linux、macOS 或 Cloud Shell）

注意：以下命令假定您已使用您的用户账号通过运行 gcloud init 或 gcloud auth login 登录 gcloud CLI，或者使用了 Cloud Shell，这会使您自动登录 gcloud CLI。您可以运行 gcloud auth list 来检查当前活跃的账号。

将请求正文保存在名为 request.json 的文件中。在终端中运行以下命令，在当前目录中创建或覆盖此文件：

cat > request.json << 'EOF'
{
"displayName": JOB_NAME,
"jobSpec":{
"workerPoolSpecs":[
  {
    "replicaCount": "1",
     "machineSpec": {
        "machineType": "n1-standard-8",
      },
      "containerSpec": {
        "imageUri": TRAINING_CONTAINER,
      }
    }
  ],
  
  "base_output_directory": {
  "output_uri_prefix": BASE_OUTPUT_DIR,
   },
  "serviceAccount": USER_SA_EMAIL,
  "tensorboard": TENSORBOARD_INSTANCE_NAME,
  }
}
EOF

然后，执行以下命令以发送 REST 请求：

curl -X POST \
     -H "Authorization: Bearer $(gcloud auth print-access-token)" \
     -H "Content-Type: application/json; charset=utf-8" \
     -d @request.json \
     "https://LOCATION_ID-aiplatform.googleapis.com/v1/projects/PROJECT_ID/locations/LOCATION_ID/customJobs"

PowerShell (Windows)

注意：以下命令假定您已使用您的用户账号通过运行 gcloud init 或 gcloud auth login 登录 gcloud CLI。您可以运行 gcloud auth list 来检查当前活跃的账号。

将请求正文保存在名为 request.json 的文件中。在终端中运行以下命令，在当前目录中创建或覆盖此文件：

@'
{
"displayName": JOB_NAME,
"jobSpec":{
"workerPoolSpecs":[
  {
    "replicaCount": "1",
     "machineSpec": {
        "machineType": "n1-standard-8",
      },
      "containerSpec": {
        "imageUri": TRAINING_CONTAINER,
      }
    }
  ],
  
  "base_output_directory": {
  "output_uri_prefix": BASE_OUTPUT_DIR,
   },
  "serviceAccount": USER_SA_EMAIL,
  "tensorboard": TENSORBOARD_INSTANCE_NAME,
  }
}
'@  | Out-File -FilePath request.json -Encoding utf8

然后，执行以下命令以发送 REST 请求：

$cred = gcloud auth print-access-token
$headers = @{ "Authorization" = "Bearer $cred" }

Invoke-WebRequest `
    -Method POST `
    -Headers $headers `
    -ContentType: "application/json; charset=utf-8" `
    -InFile request.json `
    -Uri "https://LOCATION_ID-aiplatform.googleapis.com/v1/projects/PROJECT_ID/locations/LOCATION_ID/customJobs" | Select-Object -Expand Content

您应该收到类似以下内容的 JSON 响应：

{
  "name": "projects/PROJECT_ID/locations/LOCATION_ID/customJobs/CUSTOM_JOB_ID",
  "displayName": "DISPLAY_NAME",
  "jobSpec": {
    "workerPoolSpecs": [
      {
        "machineSpec": {
          "machineType": "n1-standard-8"
        },
        "replicaCount": "1",
        "diskSpec": {
          "bootDiskType": "pd-ssd",
          "bootDiskSizeGb": 100
        },
        "containerSpec": {
          "imageUri": "IMAGE_URI"
        }
      }
    ],
    "serviceAccount": "SERVICE_ACCOUNT",
    "baseOutputDirectory": {
      "outputUriPrefix": "OUTPUT_URI_PREFIX"
    },
    "tensorboard": "projects//locations/LOCATION_ID/tensorboards/tensorboard-id"
  },
  "state": "JOB_STATE_PENDING",
  "createTime": "CREATE-TIME",
  "updateTime": "UPDATE-TIME"
}

后续步骤

了解查看 Vertex AI TensorBoard。
了解如何使用 Cloud Profiler 优化自定义训练作业的性能。