Как создать файл config.yaml для распределенного обучения на платформе Unified Cloud AI Platform

Question

Я хочу обучить модель с помощью нового сервиса Google Cloud - Unified AI Platform. Для этого я использую 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

Однако для распределенного обучения я не могу понять, как передать несколько s в этом файле. В приведенном примере файла yaml не рассматривается случай, когда я могу предоставить несколько 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