아르바이트를 하다보니 GPU를 2개씩 사용할 수 있게 되었다. 그런데 분산처리에 대한 코드를 작성하지 않으면 GPU를 하나만 사용할 수 있다. 이게 너무 아까워서 분산처리를 공부하려고 한다. 공식문서를 읽어보자.

개요

tf.distribute.Strategy 는 다수의 GPU, 다수의 TPU를 사용하여 분산처리를 할 수 있도록 도와주는 API이다. 이 API를 사용하면 아주 적은 코드만 수정하고도 이미 구현해놓은 코드들을 분산처리할 수 있다.

이 녀석의 목표는 다음과 같다.

• 연구원, 개발자 등 모두가 쉽게 사용할 수 있도록 한다.
• 즉시 사용가능하고 좋은 성능을 낸다.
• 분산처리 방식을 쉽게 바꾼다.

이 녀석은 Keras 와도 사용할 수 있고 TensorFlow 로 작성한 커스텀 루프와도 사용할 수 있다.

TF 2.x 부터는 모델을 eagerly하게 실행할 수도 있고 tf.function, tf.distribute.Strategy 를 사용해서 그래프로 실행할 수도 있다. 다만 tf.function 에서 가장 잘 돌아간다.
Eager mode 는 디버깅 목적으로만 추천하고 TPUStrategy 를 지원하지 않는다. 지금 읽는 문서에서는 이 녀석을 사용해 학습하는 것을 주로 이야기하지만 다양한 플랫폼에서 평가, 예측을 배포하는데에도 사용할 수 있다고 한다.

우리가 분산처리 방식을 변경하면 TF의 변수, 레이어, 모델, 옵티마이저, 메트릭, 서머리, 체크포인트 들이 그것을 인식하도록 구현을 해놨다고 한다. 만세!! 그래서 코드의 변경이 최소화될 수 있다.

더 알고 싶으면 이 동영상을 보세요

Types of Strategies

tf.distribute.Strategy 는 몇 가지 케이스를 커버할 수 있도록 의도되었다. 여러가지 케이스 중 몇가지는 이미 지원하고 있고 몇가지는 미래에 지원할 예정이다. 지금 지원하는 케이스는 동기/비동기 학습, 하드웨어 플랫폼(TPU, GPU)이다.
현재 지원중인 케이스들을 기반으로는 총 6가지의 Strategy를 사용할 수 있다. 표를 살펴보자.

Training API	MirroredStrategy	TPUStrategy	MultiWorkerMirroredStrategy	CentralStorageStrategy	ParameterServerStrategy
Keras API	Supported	Supported	Experimental support	Experimental support	Supported planned post 2.3
Custom training loop	Supported	Supported	Experimental support	Experimental support	Supported planned post 2.3
Estimator API	Limited Support	Not supported	Limited Support	Limited Support	Limited Support

MirroredStrategy

이 녀석은 하나의 머신에 여러 대의 GPU가 장착되어 있는 환경에서 동기 분산처리를 지원한다. 하나의 GPU 디바이스 당 하나의 replica를 생성한다(replica가 뭐지) 모델에 있는 각 변수들은 모든 replica에 미러링된다. 이 변수들은 MirroredVariable 이라고 부른다. 이 변수들은 똑같은 업데이트를 적용해서 서로 동기화 상태를 유지한다.
효율적인 all-reduce 알고리즘을 사용해서 변수들을 업데이트한다. all-reduce는 모든 장치에서 텐서를 합산하여 집계하고 각 장치에서 사용할 수 있도록한다. 아주 효율적이고 동기화 오버 헤드를 크게 줄일 수 있는 융합 알고리즘이다. 기본적으로 NVIDIA NCCL 알고리즘을 사용하지만 다른 알고리즘을 사용해도 된다.

간단한 사용법

mirrored_strategy = tf.distribute.MirroredStrategy(devices=["/gpu:0", "/gpu:1"])

위와 같이 선언 후 모델을 scope 안에서 선언하면 된다. (아래에서 다시 다룸)

override all-reduce algorithms

TF에서 제공하는 알고리즘은 다음과 같다.

1. tf.distribute.HierarchicalCopyAllReduce
2. tf.distribute.ReductionToOneDevice
3. tf.distribute.NcclAllReduce (default)

사용법은 다음과 같다.

mirrored_strategy = tf.distribute.MirroredStrategy(
	cross_device_ops=tf.distribute.HierarchicalCopyAllReduce())

TPUStrategy

TPU는 구글에서 설계한 특수한 ASICs인데 머신러닝 속도가 엄청나게 빨라진다고 한다. MirroredStrategy 와 동작은 똑같다고 한다. TPU를 사용할 떄가 오면 다시 돌아오자.

MultiWorkerMirroredStrategy

이 녀석은 tf.distribute.experimental.MultiWorkerMirroredStrategy 에 위치해있고 MirroredStrategy 와 매우 흡사하다. 이 녀석은 잠재적으로 다중 GPU를 사용하는 다중 워커에 동기 분산처리를 구현한다. 이 녀석도 MirroredStrategy 처럼 모든 변수를 모든 워커와 디바이스에 복사한다.
요런 느낌....

keras에서 사용하기

다음 두가지 과정을 거치면 사용이 가능하다.

1. tf.distribute.Strategy 인스턴스 생성하기
2. 작성한 Keras 모델을 strategy.scope 로 옮기기

Sequential, Functional, Subclassed 모두 사용 가능하다

dense layer 하나로 구성된 모델을 구성하여 예제를 살펴보자

mirrored_strategy = tf.distribute.MirroredStrategy()

with mirrored_strategy.scope():
  model = tf.keras.Sequential([tf.keras.layers.Dense(1, input_shape=(1,))])

model.compile(loss='mse', optimizer='sgd')

scope 안에 넣은 후에도 일반적인 모델, 옵티마이저, 메트릭 등을 생성하고 사용할 수 있다.
분산 처리를 완료하면 각 배치를 replica들에게 균등하게 나눠준다. 예를들어서 2 개의 GPU를 사용하고 배치의 크기가 10이면 각 GPU는 5개의 인풋을 받는다. 효율적인 배치 사이즈를 위해 다음과 같은 예제코드를 사용한다.

# Compute global batch size using number of replicas.
BATCH_SIZE_PER_REPLICA = 5
global_batch_size = (BATCH_SIZE_PER_REPLICA *
                     mirrored_strategy.num_replicas_in_sync)
dataset = tf.data.Dataset.from_tensors(([1.], [1.])).repeat(100)
dataset = dataset.batch(global_batch_size)

LEARNING_RATES_BY_BATCH_SIZE = {5: 0.1, 10: 0.15}
learning_rate = LEARNING_RATES_BY_BATCH_SIZE[global_batch_size]

mirrored_strategy.num_replicas_in_sync 를 통해 replica의 개수를 가져오고 BATCH_SIZE_PER_REPLICA 를 통해 GPU들이 받을 배치 사이즈를 작성한다.

참고문서

공식문서 - Distributed training with TensorFlow
공식문서 - Distributed training with Keras
공식문서 - Multi-worker training with Keras