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

Question

我希望使用 Google Cloud 的新服务 - 统一人工智能平台来训练模型。为此，我使用了如下所示的 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 的情况。

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

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

Answer 1

回答我自己的问题。 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