TPU Pod에서 학습

개요

Cloud TPU의 장점 중 하나는 크기가 다른 TPU 유형 간에 학습을 확장할 수 있다는 점입니다. 모든 TPU 유형은 동일한 데이터 병렬 아키텍처를 사용합니다. 유일한 변경사항은 병렬 처리가 8개 코어에서 2,048개 코어로 증가한다는 점입니다.

더 많은 수의 TPU를 최대한 활용하려면 여러 학습 작업 매개변수를 조정해야 합니다. 이 문서에서는 몇 가지 일반적인 문제와 모델에서 변경해야 하는 사항을 설명합니다.

배치 크기 조정 및 학습 단계

더 큰 TPU 유형에서 선형적 확장을 달성하려면 코어당 배치 크기를 동일하게 유지하세요.

예를 들어 v2-8에서 배치 크기 1,024를 사용하는 경우 v2-32에서는 배치 크기 4,096(4*1,024)을 사용합니다. 이렇게 하면 TPU 하드웨어를 완전히 활용합니다. 더 작은 배치 크기를 사용할 수 있지만 그렇게 하면 학습이 선형적으로 확장되지 않습니다.

많은 모델에는 1단계가 데이터의 단일 배치에 해당하는 train_steps 플래그가 포함되어 있습니다. 배치 크기를 늘릴 때 총 학습 예시 수가 동일하게 유지되도록 학습 단계 수를 줄입니다.

예를 들어 원래 구성이 1,000 단계의 배치 크기가 1,024인 경우 배치 크기를 4,096으로 늘리고 단계 수를 250(1000/4)으로 줄일 수 있습니다. 모델에서 epochs 플래그를 사용하면 단계 수를 확장할 필요가 없습니다.

배치 크기가 크면 모델의 수렴 동작이 변경될 수 있으므로 일부 초매개변수도 조정할 수 있습니다.

대규모 배치 옵티마이저 사용

더 큰 배치 크기를 사용하도록 확장하는 경우 큰 배치 크기를 지원하는 옵티마이저를 사용해야 할 수도 있습니다. 예를 들어 참조 Resnet 모델은 더 큰 TPU 유형에 LARSOptimizer를 사용합니다.

TPU Pod와 동일한 리전에서 리전별 Google Cloud Storage 버킷 사용

일반적으로 TPU 학습 시 권장사항은 항상 같은 리전의 리소스를 사용하는 것입니다. Google Cloud Storage로부터의 데이터 전송 속도가 더 높은 경향이 있으므로 리소스 리전은 특히 TPU Pod를 사용할 때 중요합니다. 학습 데이터세트 및 체크포인트에 TPU와 동일한 리전의 리전 Google Cloud Storage 버킷을 사용해야 합니다.

TPUClusterResolver에 TPU 이름 및 리전 제공

단일 기기 학습의 경우 TPU 이름 또는 IP 주소(예: grpc://1.2.3.4:8470)를 지정할 수 있습니다.

TPU Pod의 경우 TensorFlow가 학습 배포에 사용 가능한 모든 호스트의 IP 주소를 찾을 수 있도록 TPU 이름을 사용해야 합니다.

TPU Pod의 IP 주소를 지정하면 다음 오류가 발생합니다.

ValueError: TPUConfig.num_shards is not set correctly ....

모델 및 TPU를 다른 영역 또는 리전에서 실행해야 하는 경우 TPU 영역과 함께 TPUClusterResolver를 제공해야 합니다.

tputpu_zone 플래그를 참조 모델에 전달하거나 TPUClusterResolver의 값을 직접 설정할 수 있습니다. 예:

cluster = tf.contrib.cluster_resolver.TPUClusterResolver(
  tpu='tpuname',
  zone='europe-west4-a',
  project='myproject')

TPUEstimator를 사용할 때 발생할 수 있는 관련 문제는 클러스터 리졸버가 tf.contrib.tpu.RunConfig로 전달되지 않는 것입니다. 이 문제는 RunConfig가 master 매개변수로 초기화되는 경우에 발생할 수 있습니다. 대신 cluster 매개변수를 사용하여 초기화하세요.

config = tf.contrib.tpu.RunConfig(cluster=cluster)

TPUEstimator를 사용하지 않고 TPU Pod 및 tf.Session을 사용하는 경우 다음과 같이 세션 구성을 구성해야 합니다.

with tf.Session(cluster.get_master(), config=config.session_config) as sess:

TPU Pod에서의 평가

TPU Pod에서는 평가가 지원되지 않거나 비용 효율적이지 않습니다. Google에서는 대부분의 모델에서 지원되는 mode 플래그를 사용하여 단일 기기 TPU에서 평가를 실행하는 것을 권장합니다.

TPU Pod 또는 단일 기기 TPU에서의 학습을 위해서는 플래그를 train으로 설정합니다. 평가를 위해서는 단일 기기 TPU의 경우에만 eval을 설정합니다. 단일 기기 TPU에서는 train_and_eval 모드를 사용할 수 있지만 TPU Pod에서는 사용할 수 없습니다.