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

1. Обзор

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

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

.

Чему вы научитесь

Вы научитесь:

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

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

2. Введение в Vertex AI

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

Vertex AI предлагает множество различных продуктов для поддержки комплексных рабочих процессов машинного обучения. В этой лабораторной работе мы сосредоточимся на продуктах, перечисленных ниже: Training и Workbench.

Обзор продукции Vertex

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 Registry . Используя этот подход, вы можете обучать модель, созданную с помощью любого фреймворка.

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

Откройте терминал в ноутбуке.

Шаг 1: Напишите код для обучения.

Создайте новую директорию с именем flowers-multi-gpu и перейдите в неё с помощью команды cd:

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 в корневом каталоге вашего проекта flowers:

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: Создайте контейнер

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

PROJECT_ID='your-cloud-project'

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

REPO_NAME='flower-app'

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

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

Настройка Docker 

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

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

docker build ./ -t $IMAGE_URI

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

docker push $IMAGE_URI

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

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

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

В панели запуска создайте блокнот 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)

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

multigpu_job

6. [Необязательно] Обучение для нескольких работников

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

Откройте новое окно терминала в вашем ноутбуке:

Откройте терминал в ноутбуке.

Шаг 1: Напишите код для обучения.

Создайте новую директорию с именем flowers-multi-machine и перейдите в неё с помощью команды cd: 

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() содержат шаблонный код, помогающий сохранить модель.

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

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

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

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: Создайте контейнер

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

PROJECT_ID='your-cloud-project'

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

REPO_NAME='flower-app'

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

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

Настройка Docker 

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

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

docker build ./ -t $IMAGE_URI

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

docker push $IMAGE_URI

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

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

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

В панели запуска создайте блокнот 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. В этом примере обе конфигурации идентичны. Однако, если вы хотите обучать модель на 3 машинах, вам нужно добавить дополнительные рабочие процессы в пул рабочих процессов 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 в консоли. Если вы хотите полностью удалить ноутбук, нажмите кнопку «Удалить».

Остановить экземпляр

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

Удалить хранилище