如何为 Unified Cloud AI Platform 上的分布式训练创建 config.yaml 文件

问题描述

我希望使用 Google Cloud 的新服务统一 AI 平台来训练模型。为此，我使用了如下所示的 config.yaml：

workerPoolSpecs:
  workerPoolSpec:
    machinespec:
      machineType: n1-highmem-16
      acceleratorType: NVIDIA_TESLA_P100
      acceleratorCount: 2
    replicaCount: 1
    pythonPackageSpec:
      executorImageUri: us-docker.pkg.dev/cloud-aiplatform/training/tf-gpu.2-4:latest
      packageUris: gs://path/to/bucket/unified_ai_platform/src_dist/trainer-0.1.tar.gz
      pythonmodule: trainer.task
  workerPoolSpec:
    machinespec:
      machineType: n1-highmem-16
      acceleratorType: NVIDIA_TESLA_P100
      acceleratorCount: 2
    replicaCount: 2
    pythonPackageSpec:
      executorImageUri: us-docker.pkg.dev/cloud-aiplatform/training/tf-gpu.2-4:latest
      packageUris: gs://path/to/bucket/unified_ai_platform/src_dist/trainer-0.1.tar.gz
      pythonmodule: trainer.task

但是，对于分布式训练，我无法理解如何在此文件中传递多个 workerPoolSpec。提供的示例 yaml file 并未考虑我可以提供多个 workerPoolSpec 的情况。

该示例的文档还说“您可以指定多个工作池规范，以创建具有多个工作池的自定义作业”。

在这方面的任何帮助将不胜感激。

解决方法

回答我自己的问题。 config.yaml 文件应如下所示：

workerPoolSpecs:
  - machineSpec:
      machineType: n1-standard-16
      acceleratorType: NVIDIA_TESLA_P100
      acceleratorCount: 2
    replicaCount: 1
    containerSpec:
      imageUri: gcr.io/path/to/container:v2
      args: 
        - --model-dir=gs://path/to/model
        - --tfrecord-dir=gs://path/to/training/data/
        - --epochs=2
  - machineSpec:
      machineType: n1-standard-16
      acceleratorType: NVIDIA_TESLA_P100
      acceleratorCount: 2
    replicaCount: 2
    containerSpec:
      imageUri: gcr.io/path/to/container:v2
      args: 
        - --model-dir=gs://path/to/models
        - --tfrecord-dir=gs://path/to/training/data/
        - --epochs=2