От прототипа к производству: распределенное обучение Vertex AI

1. Обзор

В этой лабораторной работе вы будете использовать Vertex AI для запуска задания распределенного обучения Vertex AI Training с использованием TensorFlow.

Эта лабораторная работа является частью серии видеороликов «От прототипа к производству» . Обязательно выполните предыдущие лабораторные работы, прежде чем приступать к выполнению этой. Чтобы узнать больше, вы можете посмотреть соответствующую серию видеороликов:

.

Что вы узнаете

Вы узнаете, как:

• Запускайте распределенное обучение на одной машине с несколькими графическими процессорами.
• Запускайте распределенное обучение на нескольких машинах

Общая стоимость запуска этой лаборатории в Google Cloud составляет около 2 долларов США .

2. Знакомство с Vertex AI

В этой лаборатории используются новейшие продукты искусственного интеллекта, доступные в Google Cloud. Vertex AI интегрирует предложения машинного обучения в Google Cloud в единый процесс разработки. Раньше модели, обученные с помощью AutoML, и пользовательские модели были доступны через отдельные сервисы. Новое предложение объединяет оба API в одном API, а также другие новые продукты. Вы также можете перенести существующие проекты на Vertex AI.

Vertex AI включает в себя множество различных продуктов для поддержки комплексных рабочих процессов машинного обучения. Эта лабораторная работа будет сосредоточена на продуктах, перечисленных ниже: Обучение и рабочие места.

3. Обзор распределенного обучения

Если у вас один графический процессор, TensorFlow будет использовать этот ускоритель для ускорения обучения модели без каких-либо дополнительных усилий с вашей стороны. Однако, если вы хотите получить дополнительный импульс от использования нескольких графических процессоров, вам нужно будет использовать tf.distribute , который является модулем TensorFlow для выполнения вычислений на нескольких устройствах.

В первом разделе этой лабораторной работы используется tf.distribute.MirroredStrategy , который вы можете добавить в свои обучающие приложения, внеся лишь несколько изменений в код. Эта стратегия создает копию модели на каждом графическом процессоре вашего компьютера. Последующие обновления градиента будут происходить синхронно. Это означает, что каждый графический процессор вычисляет прямые и обратные проходы через модель на разных фрагментах входных данных. Вычисленные градиенты каждого из этих срезов затем агрегируются по всем графическим процессорам и усредняются в процессе, известном как all-reduce . Параметры модели обновляются с использованием этих усредненных градиентов.

В необязательном разделе в конце лабораторной работы используется tf.distribute.MultiWorkerMirroredStrategy , который похож на MirroredStrategy за исключением того, что он работает на нескольких машинах. Каждая из этих машин также может иметь несколько графических процессоров. Как и MirroredStrategy , MultiWorkerMirroredStrategy — это стратегия синхронного параллелизма данных, которую можно использовать лишь с небольшими изменениями кода. Основное отличие при переходе от синхронного параллелизма данных на одной машине ко многим заключается в том, что градиенты в конце каждого шага теперь необходимо синхронизировать на всех графических процессорах на машине и на всех машинах в кластере.

Для выполнения этой лабораторной работы вам не обязательно знать подробности, но если вы хотите узнать больше о том, как работает распределенное обучение в TensorFlow, посмотрите видео ниже:

4. Настройте свою среду

Выполните действия, описанные в разделе «Обучение пользовательских моделей с помощью лаборатории Vertex AI», чтобы настроить среду.

5. Обучение на одной машине, с несколькими графическими процессорами

Вы отправите задание по распределенному обучению в Vertex AI, поместив код приложения для обучения в контейнер Docker и отправив этот контейнер в реестр Google Artifact . Используя этот подход, вы можете обучить модель, построенную с помощью любого фреймворка.

Для начала в меню запуска записной книжки Workbench, созданной вами в предыдущих лабораторных работах, откройте окно терминала.

Шаг 1. Напишите обучающий код

Создайте новый каталог с именем flowers-multi-gpu и перейдите в него:

mkdir flowers-multi-gpu
cd flowers-multi-gpu

Выполните следующую команду, чтобы создать каталог для обучающего кода и файла Python, куда вы добавите приведенный ниже код.

mkdir trainer
touch trainer/task.py

Теперь в каталоге flowers-multi-gpu/ у вас должно быть следующее:

+ trainer/
    + task.py

Затем откройте только что созданный файл task.py и скопируйте приведенный ниже код.

Вам нужно будет заменить {your-gcs-bucket} в BUCKET_ROOT на корзину Cloud Storage, в которой вы сохранили набор данных цветов в лабораторной работе 1 .

import tensorflow as tf
import numpy as np
import os

## Replace {your-gcs-bucket} !!
BUCKET_ROOT='/gcs/{your-gcs-bucket}'

# Define variables
NUM_CLASSES = 5
EPOCHS=10
BATCH_SIZE = 32

IMG_HEIGHT = 180
IMG_WIDTH = 180

DATA_DIR = f'{BUCKET_ROOT}/flower_photos'

def create_datasets(data_dir, batch_size):
  '''Creates train and validation datasets.'''

  train_dataset = tf.keras.utils.image_dataset_from_directory(
    data_dir,
    validation_split=0.2,
    subset="training",
    seed=123,
    image_size=(IMG_HEIGHT, IMG_WIDTH),
    batch_size=batch_size)

  validation_dataset = tf.keras.utils.image_dataset_from_directory(
    data_dir,
    validation_split=0.2,
    subset="validation",
    seed=123,
    image_size=(IMG_HEIGHT, IMG_WIDTH),
    batch_size=batch_size)

  train_dataset = train_dataset.cache().shuffle(1000).prefetch(buffer_size=tf.data.AUTOTUNE)
  validation_dataset = validation_dataset.cache().prefetch(buffer_size=tf.data.AUTOTUNE)

  return train_dataset, validation_dataset


def create_model():
  '''Creates model.'''

  model = tf.keras.Sequential([
    tf.keras.layers.Resizing(IMG_HEIGHT, IMG_WIDTH),
    tf.keras.layers.Rescaling(1./255, input_shape=(IMG_HEIGHT, IMG_WIDTH, 3)),
    tf.keras.layers.Conv2D(16, 3, padding='same', activation='relu'),
    tf.keras.layers.MaxPooling2D(),
    tf.keras.layers.Conv2D(32, 3, padding='same', activation='relu'),
    tf.keras.layers.MaxPooling2D(),
    tf.keras.layers.Conv2D(64, 3, padding='same', activation='relu'),
    tf.keras.layers.MaxPooling2D(),
    tf.keras.layers.Flatten(),
    tf.keras.layers.Dense(128, activation='relu'),
    tf.keras.layers.Dense(NUM_CLASSES, activation='softmax')
  ])
  return model

def main():  

  # Create distribution strategy
  strategy = tf.distribute.MirroredStrategy()

  # Get data
  GLOBAL_BATCH_SIZE = BATCH_SIZE * strategy.num_replicas_in_sync
  train_dataset, validation_dataset = create_datasets(DATA_DIR, BATCH_SIZE)

  # Wrap model creation and compilation within scope of strategy
  with strategy.scope():
    model = create_model()
    model.compile(optimizer=tf.keras.optimizers.Adam(),
                  loss=tf.keras.losses.SparseCategoricalCrossentropy(),
                  metrics=['accuracy'])

  history = model.fit(
    train_dataset,
    validation_data=validation_dataset,
    epochs=EPOCHS
  )

  model.save(f'{BUCKET_ROOT}/model_output')


if __name__ == "__main__":
    main()

Прежде чем создавать контейнер, давайте более подробно рассмотрим код. Есть несколько компонентов, специфичных для использования распределенного обучения.

• В функции main() создается объект MirroredStrategy . Далее вы переносите создание переменных модели в область действия стратегии. Этот шаг сообщает TensorFlow, какие переменные следует зеркально отображать на графических процессорах.
• Размер пакета увеличивается с помощью num_replicas_in_sync . Масштабирование размера пакета — лучшая практика при использовании стратегий параллелизма синхронных данных в TensorFlow. Вы можете узнать больше здесь.

Шаг 2. Создайте файл Dockerfile.

Чтобы контейнеризировать ваш код, вам необходимо создать Dockerfile. В Dockerfile вы включите все команды, необходимые для запуска образа. Он установит все необходимые библиотеки и настроит точку входа для обучающего кода.

В терминале создайте пустой Dockerfile в корне каталога цветов:

touch Dockerfile

Теперь в каталоге flowers-multi-gpu/ у вас должно быть следующее:

+ Dockerfile
+ trainer/
    + task.py

Откройте Dockerfile и скопируйте в него следующее:

FROM gcr.io/deeplearning-platform-release/tf2-gpu.2-8

WORKDIR /

# Copies the trainer code to the docker image.
COPY trainer /trainer

# Sets up the entry point to invoke the trainer.
ENTRYPOINT ["python", "-m", "trainer.task"]

Шаг 3. Создайте контейнер.

В своем терминале выполните следующую команду, чтобы определить переменную env для вашего проекта, обязательно заменяя your-cloud-project идентификатором вашего проекта:

Вы можете получить идентификатор своего проекта, запустив gcloud config list --format 'value(core.project)' в своем терминале.

PROJECT_ID='your-cloud-project'

Создайте репозиторий в реестре артефактов. Мы будем использовать репозиторий, который мы создали в первой лабораторной работе.

REPO_NAME='flower-app'

Определите переменную с URI вашего образа контейнера в реестре артефактов:

IMAGE_URI=us-central1-docker.pkg.dev/$PROJECT_ID/$REPO_NAME/flower_image_distributed:single_machine

Настроить докер

gcloud auth configure-docker \
    us-central1-docker.pkg.dev

Затем создайте контейнер, запустив следующую команду из корня вашего каталога flowers-multi-gpu :

docker build ./ -t $IMAGE_URI

Наконец, отправьте его в реестр артефактов:

docker push $IMAGE_URI

Теперь, когда контейнер отправлен в реестр артефактов, вы готовы приступить к работе по обучению.

Шаг 4. Запустите задание с помощью SDK.

В этом разделе вы увидите, как настроить и запустить задание распределенного обучения с помощью Vertex AI Python SDK.

В панели запуска создайте блокнот TensorFlow 2.

Импортируйте Vertex AI SDK.

from google.cloud import aiplatform

Затем определите CustomContainerTrainingJob .

Вам нужно будет заменить {PROJECT_ID} container_uri и {YOUR_BUCKET} в staging_bucket .

job = aiplatform.CustomContainerTrainingJob(display_name='flowers-multi-gpu',
                                            container_uri='us-central1-docker.pkg.dev/{PROJECT_ID}/flower-app/flower_image_distributed:single_machine',
                                            staging_bucket='gs://{YOUR_BUCKET}')

После определения задания его можно запустить. Вы установите количество ускорителей равным 2. Если бы мы использовали только 1 графический процессор, это не считалось бы распределенным обучением. Распределенное обучение на одной машине — это когда вы используете 2 и более ускорителей.

my_custom_job.run(replica_count=1,
                  machine_type='n1-standard-4',
                  accelerator_type='NVIDIA_TESLA_V100',
                  accelerator_count=2)

В консоли вы сможете видеть ход выполнения вашей работы.

6. [Дополнительно] Обучение нескольких сотрудников

Теперь, когда вы попробовали распределенное обучение на одной машине с несколькими графическими процессорами, вы можете поднять свои навыки распределенного обучения на новый уровень, обучаясь на нескольких машинах. Чтобы снизить затраты, мы не будем добавлять графические процессоры к этим машинам, но вы можете поэкспериментировать, добавив графические процессоры, если хотите.

Откройте новое окно терминала в экземпляре вашего блокнота:

Шаг 1. Напишите обучающий код

Создайте новый каталог с именем flowers-multi-machine и перейдите в него:

mkdir flowers-multi-machine
cd flowers-multi-machine

Выполните следующую команду, чтобы создать каталог для обучающего кода и файла Python, куда вы добавите приведенный ниже код.

mkdir trainer
touch trainer/task.py

Теперь в каталоге flowers-multi-machine/ у вас должно быть следующее:

+ trainer/
    + task.py

Затем откройте только что созданный файл task.py и скопируйте приведенный ниже код.

Вам нужно будет заменить {your-gcs-bucket} в BUCKET_ROOT на корзину Cloud Storage, в которой вы сохранили набор данных цветов в лабораторной работе 1 .

import tensorflow as tf
import numpy as np
import os

## Replace {your-gcs-bucket} !!
BUCKET_ROOT='/gcs/{your-gcs-bucket}'

# Define variables
NUM_CLASSES = 5
EPOCHS=10
BATCH_SIZE = 32

IMG_HEIGHT = 180
IMG_WIDTH = 180

DATA_DIR = f'{BUCKET_ROOT}/flower_photos'
SAVE_MODEL_DIR = f'{BUCKET_ROOT}/multi-machine-output'

def create_datasets(data_dir, batch_size):
  '''Creates train and validation datasets.'''

  train_dataset = tf.keras.utils.image_dataset_from_directory(
    data_dir,
    validation_split=0.2,
    subset="training",
    seed=123,
    image_size=(IMG_HEIGHT, IMG_WIDTH),
    batch_size=batch_size)

  validation_dataset = tf.keras.utils.image_dataset_from_directory(
    data_dir,
    validation_split=0.2,
    subset="validation",
    seed=123,
    image_size=(IMG_HEIGHT, IMG_WIDTH),
    batch_size=batch_size)

  train_dataset = train_dataset.cache().shuffle(1000).prefetch(buffer_size=tf.data.AUTOTUNE)
  validation_dataset = validation_dataset.cache().prefetch(buffer_size=tf.data.AUTOTUNE)

  return train_dataset, validation_dataset


def create_model():
  '''Creates model.'''

  model = tf.keras.Sequential([
    tf.keras.layers.Resizing(IMG_HEIGHT, IMG_WIDTH),
    tf.keras.layers.Rescaling(1./255, input_shape=(IMG_HEIGHT, IMG_WIDTH, 3)),
    tf.keras.layers.Conv2D(16, 3, padding='same', activation='relu'),
    tf.keras.layers.MaxPooling2D(),
    tf.keras.layers.Conv2D(32, 3, padding='same', activation='relu'),
    tf.keras.layers.MaxPooling2D(),
    tf.keras.layers.Conv2D(64, 3, padding='same', activation='relu'),
    tf.keras.layers.MaxPooling2D(),
    tf.keras.layers.Flatten(),
    tf.keras.layers.Dense(128, activation='relu'),
    tf.keras.layers.Dense(NUM_CLASSES, activation='softmax')
  ])
  return model

def _is_chief(task_type, task_id):
  '''Helper function. Determines if machine is chief.'''

  return task_type == 'chief'


def _get_temp_dir(dirpath, task_id):
  '''Helper function. Gets temporary directory for saving model.'''

  base_dirpath = 'workertemp_' + str(task_id)
  temp_dir = os.path.join(dirpath, base_dirpath)
  tf.io.gfile.makedirs(temp_dir)
  return temp_dir


def write_filepath(filepath, task_type, task_id):
  '''Helper function. Gets filepath to save model.'''

  dirpath = os.path.dirname(filepath)
  base = os.path.basename(filepath)
  if not _is_chief(task_type, task_id):
    dirpath = _get_temp_dir(dirpath, task_id)
  return os.path.join(dirpath, base)

def main():
  # Create distribution strategy
  strategy = tf.distribute.MultiWorkerMirroredStrategy()

  # Get data
  GLOBAL_BATCH_SIZE = BATCH_SIZE * strategy.num_replicas_in_sync
  train_dataset, validation_dataset = create_datasets(DATA_DIR, BATCH_SIZE)

  # Wrap variable creation within strategy scope
  with strategy.scope():
    model = create_model()
    model.compile(optimizer=tf.keras.optimizers.Adam(),
                  loss=tf.keras.losses.SparseCategoricalCrossentropy(),
                  metrics=['accuracy'])

  history = model.fit(
    train_dataset,
    validation_data=validation_dataset,
    epochs=EPOCHS
  )

  # Determine type and task of the machine from
  # the strategy cluster resolver
  task_type, task_id = (strategy.cluster_resolver.task_type,
                        strategy.cluster_resolver.task_id)

  # Based on the type and task, write to the desired model path
  write_model_path = write_filepath(SAVE_MODEL_DIR, task_type, task_id)
  model.save(write_model_path)

if __name__ == "__main__":
    main()

Прежде чем создавать контейнер, давайте более подробно рассмотрим код. В коде есть несколько компонентов, необходимых вашему обучающему приложению для работы с MultiWorkerMirroredStrategy .

• В функции main() создается объект MultiWorkerMirroredStrategy . Далее вы переносите создание переменных модели в область действия стратегии. Этот важный шаг сообщает TensorFlow, какие переменные следует зеркально отображать в репликах.
• Размер пакета увеличивается с помощью num_replicas_in_sync . Масштабирование размера пакета — лучшая практика при использовании стратегий параллелизма синхронных данных в TensorFlow.
• Сохранение вашей модели немного сложнее в случае с несколькими рабочими процессами, поскольку место назначения должно быть разным для каждого из рабочих процессов. Главный работник сохранит модель в желаемом каталоге модели, а другие работники сохранят модель во временных каталогах. Важно, чтобы эти временные каталоги были уникальными, чтобы предотвратить запись нескольких рабочих процессов в одно и то же место. Сбережение может включать в себя коллективные операции, а это означает, что экономить должны все работники, а не только руководитель. Функции _is_chief() , _get_temp_dir() , write_filepath() , а также функция main() включают стандартный код, который помогает сохранить модель.

Обратите внимание: если вы использовали MultiWorkerMirroredStrategy в другой среде, возможно, вы настроили переменную среды TF_CONFIG . Vertex AI автоматически устанавливает TF_CONFIG , поэтому вам не нужно определять эту переменную на каждом компьютере в вашем кластере.

Шаг 2. Создайте файл Dockerfile.

Чтобы контейнеризировать ваш код, вам необходимо создать Dockerfile. В Dockerfile вы включите все команды, необходимые для запуска образа. Он установит все необходимые библиотеки и настроит точку входа для обучающего кода.

В терминале создайте пустой Dockerfile в корне каталога цветов:

touch Dockerfile

Теперь в каталоге flowers-multi-machine/ у вас должно быть следующее:

+ Dockerfile
+ trainer/
    + task.py

Откройте Dockerfile и скопируйте в него следующее:

FROM gcr.io/deeplearning-platform-release/tf2-gpu.2-8

WORKDIR /

# Copies the trainer code to the docker image.
COPY trainer /trainer

# Sets up the entry point to invoke the trainer.
ENTRYPOINT ["python", "-m", "trainer.task"]

Шаг 3. Создайте контейнер.

В своем терминале выполните следующую команду, чтобы определить переменную env для вашего проекта, обязательно заменяя your-cloud-project идентификатором вашего проекта:

Вы можете получить идентификатор своего проекта, запустив gcloud config list --format 'value(core.project)' в своем терминале.

PROJECT_ID='your-cloud-project'

Создайте репозиторий в реестре артефактов. Мы будем использовать репозиторий, который мы создали в первой лабораторной работе.

REPO_NAME='flower-app'

Определите переменную с URI вашего образа контейнера в реестре Google Artifact:

IMAGE_URI=us-central1-docker.pkg.dev/$PROJECT_ID/$REPO_NAME/flower_image_distributed:multi_machine

Настроить докер

gcloud auth configure-docker \
    us-central1-docker.pkg.dev

Затем создайте контейнер, запустив следующую команду из корня вашего каталога flowers-multi-machine :

docker build ./ -t $IMAGE_URI

Наконец, отправьте его в реестр артефактов:

docker push $IMAGE_URI

Теперь, когда контейнер отправлен в реестр артефактов, вы готовы приступить к работе по обучению.

Шаг 4. Запустите задание с помощью SDK.

В этом разделе вы увидите, как настроить и запустить задание распределенного обучения с помощью Vertex AI Python SDK.

В панели запуска создайте блокнот TensorFlow 2.

Импортируйте Vertex AI SDK.

from google.cloud import aiplatform

Затем определите worker_pool_specs .

Vertex AI предоставляет 4 рабочих пула для выполнения различных типов машинных задач.

Рабочий пул 0 настраивает основного, главного, планировщика или «главного». В MultiWorkerMirroredStrategy все машины обозначены как рабочие, то есть физические машины, на которых выполняются реплицированные вычисления. Помимо того, что каждая машина является рабочей, должен быть один рабочий, который берет на себя дополнительную работу, например сохранение контрольных точек и запись файлов сводки в TensorBoard. Эта машина известна как главная. Всегда существует только один главный работник, поэтому количество рабочих для пула рабочих 0 всегда будет равно 1.

В пуле рабочих 1 вы настраиваете дополнительных рабочих для вашего кластера.

Первый словарь в списке worker_pool_specs представляет рабочий пул 0, а второй словарь представляет рабочий пул 1. В этом примере две конфигурации идентичны. Однако, если вы хотите обучать на трех машинах, вы должны добавить дополнительных рабочих в пул рабочих 1, установив для параметра replica_count значение 2. Если вы хотите добавить графические процессоры, вам нужно будет добавить аргументы accelerator_type и accelerator_count в machine_spec для обоих. рабочие пулы. Обратите внимание: если вы хотите использовать графические процессоры с MultiWorkerMirroredStrategy , каждая машина в кластере должна иметь одинаковое количество графических процессоров. В противном случае работа провалится.

Вам нужно будет заменить {PROJECT_ID} в image_uri .

# The spec of the worker pools including machine type and Docker image
# Be sure to replace PROJECT_ID in the "image_uri" with your project.

worker_pool_specs=[
     {
        "replica_count": 1,
        "machine_spec": {
          "machine_type": "n1-standard-4",
        },
        "container_spec": {"image_uri": "us-central1-docker.pkg.dev/{PROJECT_ID}/flower-app/flower_image_distributed:multi_machine"}
      },
      {
        "replica_count": 1,
        "machine_spec": {
          "machine_type": "n1-standard-4",
        },
        "container_spec": {"image_uri": "us-central1-docker.pkg.dev/{PROJECT_ID}/flower-app/flower_image_distributed:multi_machine"}
      }
          ]

Затем создайте и запустите CustomJob , заменив {YOUR_BUCKET} в staging_bucket на корзину в вашем проекте для промежуточного хранения.

my_custom_job = aiplatform.CustomJob(display_name='flowers-multi-worker',
                                     worker_pool_specs=worker_pool_specs,
                                     staging_bucket='gs://{YOUR_BUCKET}')

my_custom_job.run()

В консоли вы сможете видеть ход выполнения вашей работы.

🎉Поздравляем! 🎉

Вы узнали, как использовать Vertex AI, чтобы:

• Запускайте распределенные обучающие задания с помощью TensorFlow

Чтобы узнать больше о различных частях Vertex, ознакомьтесь с документацией .

7. Очистка

Поскольку мы настроили блокнот на тайм-аут после 60 минут простоя, нам не нужно беспокоиться об отключении экземпляра. Если вы хотите вручную завершить работу экземпляра, нажмите кнопку «Стоп» в разделе консоли Vertex AI Workbench. Если вы хотите полностью удалить блокнот, нажмите кнопку «Удалить».

Чтобы удалить сегмент хранилища, в меню навигации облачной консоли перейдите к разделу «Хранилище», выберите сегмент и нажмите «Удалить»: