[Tensorflow] 텐서 복제본 생성하기: tf.distribute.ReplicaContext, tf.distribute.get_replica_context

번역된 Tensorflow 공식 문서를 보고 정리하였습니다.

tf.distribute.get_strategy()

tf.distribute.get_strategy()는 cross-replica context에서 사용됩니다.

Typically only used in a cross-replica context

(e.g., 분산훈련)
tf.distribute.get_strategy()는 tf.distribute.Strategy의 인스턴스를 반환합니다.

tf.distribute.Strategy

tf.distribute.Strategy는 훈련을 여러 GPU 또는 여러 장비, 여러 TPU로 나누어 처리하기 위한 텐서플로 API입니다. 이 API를 사용하면 기존의 모델이나 훈련 코드를 조금만 고쳐서 분산처리를 할 수 있습니다.

Strategy 종류

동기훈련 대 비동기 훈련

  • 동기 훈련
    • 각 GPU가 입력 데이터를 나누어 갖고 동시에 훈련하여 각 단계마다 그래디언트(gradient)를 모읍니다.
    • GPU 간 통신을 통해 그래디언트의 평균을 계산하며 모델을 업데이트 합니다.
    • 일반적으로 평균을 계산하기 위해 All-reduce라는 집단 통신 작업을 사용합니다.
  • 비동기 훈련
    • 모든 워커가 독립적으로 입력 데이털르 사용해 훈련하고 각각 비동기적으로 변수들을 갱신합니다.
    • 일반적으로 비동기 훈련은 파라미터 서버 구조를 사용합니다.

하드웨어 플랫폼

사용 방법

tf.distribute.Strategy.num_replicas_in_sync를 통해 복제본의 수를 얻을 수 있습니다.

케라스와 함께 tf.distribute.Strategy 사용하기

  1. 적절한 tf.distribute.Strategy 인스턴스를 만듭니다.
  2. 케라스 모델의 생성과 컴파일을 strategy.scope 안으로 옮겨줍니다.

tf.distribute.StrategySequential, 함수형 API, 클래스 상속 등 모든 방식으로 만든 케라스 모델을 다 지원합니다.

데이터셋이나 넘파이를 사용하는 두 경우 모두 입력 배치가 동일한 크기로 나누어져서 여러 개로 복제된 작업에 전달됩니다. 예를 들어, MirroredStrategy를 2개의 GPU에서 사용한다면, 크기가 10개인 배치(batch)가 두 개의 GPU로 배분됩니다. 즉, 각 GPU는 한 단계마다 5개의 입력을 받게 됩니다. 따라서 GPU가 추가될수록 각 에포크(epoch) 당 훈련 시간은 줄어들게 됩니다. 일반적으로는 가속기를 더 추가할 때마다 배치 사이즈도 더 키웁니다. 추가한 컴퓨팅 자원을 더 효과적으로 사용하기 위해서입니다. 모델에 따라서는 학습률(learning rate)을 재조정해야 할 수도 있을 것입니다.

사용자 정의 훈련 루프와 함게 tf.distribute.Strategy 사용하기

  1. tf.distribute.Strategy.scope 메서드를 호출하여 범위(scope) 안에서 모델과 옵티마이저를 만듭니다. 이는 모델이나 오티마이저로 만들어진 변수가 미러링 되도록 만듭니다.
with mirrored_strategy.scope():
  model = tf.keras.Sequential([tf.keras.layers.Dense(1, input_shape=(1,))])
  optimizer = tf.keras.optimizers.SGD()
  1. 입력 데이터셋을 만든 다음, tf.distribute.Strategy.experimental_distribute_dataset 메서드를 호출하여 전략에 맞게 데이터셋을 분배합니다.
with mirrored_strategy.scope():
  dataset = tf.data.Dataset.from_tensors(([1.], [1.])).repeat(1000).batch(
      global_batch_size)
  dist_dataset = mirrored_strategy.experimental_distribute_dataset(dataset)

tf.distribute.get_replica_context

tf.distribute.get_replica_context로 ReplicaContext의 인스턴스를 얻을 수 있습니다.

Reference

[1] https://www.tensorflow.org/guide/distributed_training
[2] https://brunch.co.kr/@chris-song/96