프로토타입에서 프로덕션으로: Vertex AI에서 분산 학습

1. 개요

이 실습에서는 TensorFlow를 사용하는 Vertex AI Training에서 Vertex AI를 사용하여 분산 학습 작업을 실행합니다.

이 실습은 프로토타입에서 프로덕션으로 동영상 시리즈 중 하나이며, 이전 실습을 완료한 후에 이 실습을 진행해야 합니다. 함께 제공되는 동영상 시리즈를 시청하여 자세히 알아볼 수 있습니다.

.

학습 내용

다음 작업을 수행하는 방법을 배우게 됩니다.

  • 여러 GPU가 있는 단일 머신에서 분산 학습 실행
  • 여러 머신에서 분산 학습 실행

Google Cloud에서 이 실습을 진행하는 데 드는 총 비용은 약 $2입니다.

2. Vertex AI 소개

이 실습에서는 Google Cloud에서 제공되는 최신 AI 제품을 사용합니다. Vertex AI는 Google Cloud 전반의 ML 제품을 원활한 개발 환경으로 통합합니다. 이전에는 AutoML로 학습된 모델과 커스텀 모델은 별도의 서비스를 통해 액세스할 수 있었습니다. 새 서비스는 다른 새로운 제품과 함께 두 가지 모두를 단일 API로 결합합니다. 기존 프로젝트를 Vertex AI로 이전할 수도 있습니다.

Vertex AI에는 엔드 투 엔드 ML 워크플로를 지원하는 다양한 제품이 포함되어 있습니다. 이 실습에서는 아래에 강조 표시된 학습Workbench 제품에 중점을 둡니다.

Vertex 제품 개요

3. 분산 학습 개요

단일 GPU가 있는 경우 TensorFlow가 이 액셀러레이터를 사용하여 추가 작업 없이 모델 학습 속도를 높입니다. 그러나 여러 GPU를 사용하여 성능을 더욱 높이려면 여러 기기에서 연산을 실행하는 TensorFlow 모듈인 tf.distribute를 사용해야 합니다.

이 실습의 첫 번째 섹션에서는 코드를 몇 번만 변경하여 학습 애플리케이션에 추가할 수 있는 tf.distribute.MirroredStrategy를 사용합니다. 이 전략은 시스템의 각 GPU에 모델 사본을 만듭니다. 후속 경사 업데이트는 동기 방식으로 진행됩니다. 즉, 각 GPU가 입력 데이터의 여러 슬라이스에서 모델을 정방향 및 역방향으로 전달합니다. 그런 다음 각 슬라이스의 계산된 경사는 모든 GPU에서 집계되며 all-reduce라고 하는 프로세스로 평균화됩니다. 모델 매개변수는 이러한 평균 경사를 사용하여 업데이트됩니다.

실습 끝 부분에 있는 선택사항 섹션은 여러 머신에서 작동한다는 점을 제외하면 MirroredStrategy와 유사한 tf.distribute.MultiWorkerMirroredStrategy를 사용합니다. 이러한 각 머신에는 GPU가 여러 개 있을 수도 있습니다. MirroredStrategy와 유사한 MultiWorkerMirroredStrategy는 몇 가지 코드 변경만으로 사용할 수 있는 동기식 데이터 병렬화 전략입니다. 한 머신의 동기식 데이터 병렬 처리에서 여러 머신의 동기식 데이터 병렬 처리로 이동할 때 주요 차이점은 각 단계가 끝날 때 이제 경사가 머신의 모든 GPU와 클러스터의 모든 머신에서 동기화되어야 한다는 것입니다.

이 실습을 완료하기 위해 세부정보를 알아야 할 필요는 없지만 TensorFlow에서 분산 학습이 작동하는 방식을 자세히 알아보려면 아래 동영상을 확인하세요.

4. 환경 설정

Vertex AI를 사용한 커스텀 모델 학습 실습의 단계를 완료하여 환경을 설정합니다.

5. 단일 머신, 다중 GPU 학습

학습 애플리케이션 코드를 Docker 컨테이너에 넣고 이 컨테이너를 Google Artifact Registry로 푸시하여 분산 학습 작업을 Vertex AI에 제출합니다. 이 방식을 사용하면 모든 프레임워크로 빌드된 모델을 학습시킬 수 있습니다.

시작하려면 이전 실습에서 만든 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 파일을 열고 아래의 코드를 복사합니다.

BUCKET_ROOT{your-gcs-bucket}실습 1에서 flowers 데이터 세트를 저장한 Cloud Storage 버킷으로 바꿔야 합니다.

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 객체가 생성됩니다. 다음은 전략 범위 내에 모델 변수 생성을 래핑합니다. 이 단계는 모든 GPU에서 미러링해야 하는 변수를 TensorFlow에 알려줍니다.
  • 배치 크기는 num_replicas_in_sync에 의해 확장됩니다. TensorFlow에서 동기 데이터 동시 로드 전략을 사용할 때는 배치 크기를 확장하는 것이 좋습니다. 자세한 내용은 여기를 참조하세요.

2단계: Dockerfile 만들기

코드를 컨테이너화하려면 Dockerfile을 만들어야 합니다. Dockerfile에는 이미지를 실행하는 데 필요한 모든 명령어가 포함됩니다. 필요한 모든 라이브러리를 설치하고 학습 코드의 진입점을 설정합니다.

터미널에서 flowers 디렉터리의 루트에 빈 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를 프로젝트의 ID로 바꿉니다.

PROJECT_ID='your-cloud-project'

Artifact Registry에서 저장소를 만듭니다. 첫 번째 실습에서 만든 저장소를 사용합니다.

REPO_NAME='flower-app'

Artifact Registry에서 컨테이너 이미지의 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

마지막으로 Artifact Registry로 푸시합니다.

docker push $IMAGE_URI

컨테이너가 Artifact Registry로 푸시되었으므로 이제 학습 작업을 시작할 준비가 되었습니다.

4단계: SDK로 작업 실행

이 섹션에서는 Vertex AI Python SDK를 사용하여 분산 학습 작업을 구성하고 실행하는 방법을 알아봅니다.

런처에서 TensorFlow 2 노트북을 만듭니다.

new_notebook

Vertex AI SDK를 가져옵니다.

from google.cloud import aiplatform

그런 다음 CustomContainerTrainingJob을 정의합니다.

container_uri{PROJECT_ID}staging_bucket{YOUR_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로 설정합니다. GPU가 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. [선택사항] 다중 작업자 학습

이제 GPU가 여러 개 있는 단일 머신에서 분산 학습을 시도했으므로 여러 머신에서 학습을 통해 분산 학습 기술을 한 단계 높일 수 있습니다. 비용 절감을 위해 이러한 머신에 GPU를 추가하지 않겠지만 원하는 경우 GPU를 추가하여 실험할 수 있습니다.

노트북 인스턴스에서 새 터미널 창을 엽니다.

노트북에서 터미널 열기

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 파일을 열고 아래의 코드를 복사합니다.

BUCKET_ROOT{your-gcs-bucket}실습 1에서 flowers 데이터 세트를 저장한 Cloud Storage 버킷으로 바꿔야 합니다.

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에는 이미지를 실행하는 데 필요한 모든 명령어가 포함됩니다. 필요한 모든 라이브러리를 설치하고 학습 코드의 진입점을 설정합니다.

터미널에서 flowers 디렉터리의 루트에 빈 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를 프로젝트의 ID로 바꿉니다.

PROJECT_ID='your-cloud-project'

Artifact Registry에서 저장소를 만듭니다. 첫 번째 실습에서 만든 저장소를 사용합니다.

REPO_NAME='flower-app'

Google Artifact Registry에서 컨테이너 이미지의 URI로 변수를 정의합니다.

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

마지막으로 Artifact Registry로 푸시합니다.

docker push $IMAGE_URI

컨테이너가 Artifact Registry로 푸시되었으므로 이제 학습 작업을 시작할 준비가 되었습니다.

4단계: SDK로 작업 실행

이 섹션에서는 Vertex AI Python SDK를 사용하여 분산 학습 작업을 구성하고 실행하는 방법을 알아봅니다.

런처에서 TensorFlow 2 노트북을 만듭니다.

new_notebook

Vertex AI SDK를 가져옵니다.

from google.cloud import aiplatform

그런 다음 worker_pool_specs를 정의합니다.

Vertex AI는 다양한 유형의 머신 작업을 처리하기 위해 4개의 작업자 풀을 제공합니다.

작업자 풀 0은 기본, 최고, 스케줄러 또는 '마스터'를 구성합니다. MultiWorkerMirroredStrategy에서 모든 머신은 복제된 계산이 실행되는 물리적 머신인 작업자로 지정됩니다. 작업자가 되는 각 머신 외에도 체크포인트 저장 및 텐서보드에 요약 파일 쓰기와 같은 일부 추가 작업을 수행하는 작업자가 하나 있어야 합니다. 이 머신을 최고 작업자라고 합니다. 최고 작업자는 하나뿐이므로 작업자 풀 0의 작업자 수는 항상 1개입니다.

작업자 풀 1에는 클러스터의 추가 작업자를 구성합니다.

worker_pool_specs 목록의 첫 번째 사전은 작업자 풀 0을 나타내고 두 번째 사전은 작업자 풀 1을 나타냅니다. 이 샘플에서 두 구성은 동일합니다. 그러나 3대의 머신에서 학습시키려면 replica_count를 2로 설정하여 작업자 풀 1에 작업자를 추가합니다. GPU를 추가하려면 두 작업자 풀의 machine_spec에 인수 accelerator_typeaccelerator_count를 추가해야 합니다. MultiWorkerMirroredStrategy에서 GPU를 사용하려면 클러스터의 각 머신에 동일한 수의 GPU가 있어야 합니다. 그렇지 않으면 작업이 실패합니다.

image_uri에서 {PROJECT_ID}를 바꿔야 합니다.

# 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을 만들고 실행하여 staging_bucket{YOUR_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()

콘솔에서 작업 진행 상황을 볼 수 있습니다.

multi_worker_job

🎉 수고하셨습니다. 🎉

Vertex AI를 사용하여 다음을 수행하는 방법을 배웠습니다.

  • TensorFlow로 분산 학습 작업 실행

Vertex의 다른 부분에 대해 자세히 알아보려면 문서를 확인하세요.

7. 정리

노트북이 유휴 시간 60분 후에 타임아웃되도록 구성했으므로 인스턴스를 종료할 필요가 없습니다. 인스턴스를 수동으로 종료하려면 콘솔의 Vertex AI Workbench 섹션에서 중지 버튼을 클릭합니다. 노트북을 완전히 삭제하려면 '삭제' 버튼을 클릭합니다.

인스턴스 중지

스토리지 버킷을 삭제하려면 Cloud 콘솔의 탐색 메뉴를 사용하여 스토리지로 이동하고 버킷을 선택하고 '삭제'를 클릭합니다.

스토리지 삭제