Cloud TPU의 TPUEstimator API
이 문서는 Cloud TPU에서 TPUEstimator API를 사용하는 방법을 설명합니다. TPUEstimator는 수준이 낮은 다양한 하드웨어 세부 작업을 처리하여 Cloud TPU의 모델 실행을 단순화합니다.
TPUEstimator를 사용하여 작성된 모델은 일반적으로 코드 변경 없이 CPU, GPU, 단일 TPU 기기, TPU Pod에서 작동합니다. 또한 TPUEstimator에서 일부 최적화가 자동으로 수행되므로 쉽게 최대 성능에 도달할 수 있습니다.
TPU 하드웨어에서 머신러닝 작업 부하의 일반적인 작동 방법에 대해 알아보려면 시스템 아키텍처 문서를 참조하세요.
표준 TensorFlow Estimator API
표준 TensorFlow Estimator API의 이점을 간단히 설명하면 다음과 같습니다.
Estimator.train()- 정해진 단계 수를 실행하는 동안 지정된 입력으로 모델을 학습시킵니다.Estimator.evaluate()- 테스트 세트로 모델을 평가합니다.Estimator.predict()- 학습된 모델을 사용하여 추론을 실행합니다.Estimator.export_savedmodel()- 예측을 제공할 모델을 내보냅니다.
또한 Estimator에는 체크포인트 저장 및 복원, 텐서보드 요약 생성 등의 학습 작업에 공통적으로 사용되는 기본 작업이 포함되어 있습니다.
Estimator를 사용하려면 TensorFlow 그래프의 모델 및 입력 부분에 해당되는 model_fn 및 input_fn을 작성해야 합니다.
TPUEstimator 프로그래밍 모델
TPUEstimator는 연산(model_fn)을 래핑하여 사용 가능한 모든 Cloud TPU 코어에 배포합니다. 학습률은 배치 크기에 맞게 조정해야 합니다.
input_fn함수는 원격 호스트 CPU에서 실행되는 입력 파이프라인을 모델링합니다. 프로그래머 가이드의 설명대로tf.data를 사용하여 입력 작업을 프로그래밍하세요. 호출별로 하나의 기기에 대한 전역 배치 입력을 처리합니다. 샤드 배치 크기는params['batch_size']에서 가져옵니다. 프로 팁: 최적의 성능을 위해 텐서 대신 데이터 세트를 반환하세요.model_fn함수는 연산이 TPU에 복제 및 분산되도록 모델링합니다. 연산에는 Cloud TPU에서 지원되는 작업만 포함해야 합니다. TensorFlow 작업에는 사용 가능한 작업 목록이 포함됩니다.
TPUEstimator를 사용한 학습 예
다음 코드는 TPUEstimator를 사용하여 MNIST를 학습시키는 방법을 보여줍니다.
def model_fn(features, labels, mode, params):
"""A simple CNN."""
del params # unused
input_layer = tf.reshape(features, [-1, 28, 28, 1])
conv1 = tf.layers.conv2d(
inputs=input_layer, filters=32, kernel_size=[5, 5], padding="same",
activation=tf.nn.relu)
pool1 = tf.layers.max_pooling2d(inputs=conv1, pool_size=[2, 2], strides=2)
conv2 = tf.layers.conv2d(
inputs=pool1, filters=64, kernel_size=[5, 5],
padding="same", activation=tf.nn.relu)
pool2 = tf.layers.max_pooling2d(inputs=conv2, pool_size=[2, 2], strides=2)
pool2_flat = tf.reshape(pool2, [-1, 7 * 7 * 64])
dense = tf.layers.dense(inputs=pool2_flat, units=128, activation=tf.nn.relu)
dropout = tf.layers.dropout(
inputs=dense, rate=0.4, training=mode == tf.estimator.ModeKeys.TRAIN)
logits = tf.layers.dense(inputs=dropout, units=10)
onehot_labels = tf.one_hot(indices=tf.cast(labels, tf.int32), depth=10)
loss = tf.losses.softmax_cross_entropy(
onehot_labels=onehot_labels, logits=logits)
learning_rate = tf.train.exponential_decay(
FLAGS.learning_rate, tf.train.get_global_step(), 100000, 0.96)
optimizer = tpu_optimizer.CrossShardOptimizer(
tf.train.GradientDescentOptimizer(learning_rate=learning_rate))
train_op = optimizer.minimize(loss, global_step=tf.train.get_global_step())
return tpu_estimator.TPUEstimatorSpec(mode=mode, loss=loss, train_op=train_op)
def make_input_fn(filename):
"""Returns an `input_fn` for train and eval."""
def input_fn(params):
"""An input_fn to parse 28x28 images from filename using tf.data."""
batch_size = params["batch_size"]
def parser(serialized_example):
"""Parses a single tf.Example into image and label tensors."""
features = tf.parse_single_example(
serialized_example,
features={
"image_raw": tf.FixedLenFeature([], tf.string),
"label": tf.FixedLenFeature([], tf.int64),
})
image = tf.decode_raw(features["image_raw"], tf.uint8)
image.set_shape([28 * 28])
# Normalize the values of the image from the range [0, 255] to [-0.5, 0.5]
image = tf.cast(image, tf.float32) * (1. / 255) - 0.5
label = tf.cast(features["label"], tf.int32)
return image, label
dataset = tf.contrib.data.TFRecordDataset(
filename, buffer_size=FLAGS.dataset_reader_buffer_size)
dataset = dataset.repeat()
dataset = dataset.apply(
tf.contrib.data.map_and_batch(
parser, batch_size=batch_size,
num_parallel_batches=8,
drop_remainder=True))
return dataset
return input_fn
def main(unused_argv):
tf.logging.set_verbosity(tf.logging.INFO)
run_config = tpu_config.RunConfig(
master=FLAGS.master,
model_dir=FLAGS.model_dir,
session_config=tf.ConfigProto(
allow_soft_placement=True, log_device_placement=True),
tpu_config=tpu_config.TPUConfig(FLAGS.iterations))
estimator = tpu_estimator.TPUEstimator(
model_fn=model_fn,
use_tpu=FLAGS.use_tpu,
train_batch_size=FLAGS.batch_size,
eval_batch_size=FLAGS.batch_size,
config=run_config)
estimator.train(input_fn=make_input_fn(FLAGS.train_file),
max_steps=FLAGS.train_steps)
다음 섹션은 위의 샘플에서 소개한 새로운 개념을 설명하여 Cloud TPU를 효과적으로 사용하는 데 도움이 됩니다.
TPUEstimator 개념
TPUEstimator는 텐서플로우 프로그램을 실행하는 데 그래프 내 복제 방법을 사용합니다. 그래프 내(단일 세션) 복제는 일반적으로 분산형 TensorFlow에 사용된 그래프 간(여러 세션) 복제와 다릅니다. 가장 큰 차이점은 다음과 같습니다.
TPUEstimator에서는 TensorFlow 세션 개시자가 로컬이 아닙니다. Python 프로그램은 Cloud TPU의 모든 코어에 복제되는 1개의 그래프를 생성합니다. 일반적으로 TensorFlow 세션 개시자가 첫 번째 작업자로 설정됩니다.
학습 예시가 최대한 빨리 Cloud TPU에 피드되도록 로컬 대신 원격 호스트에 입력 파이프라인을 배치합니다. 데이터 세트(
tf.data)는 필수 항목입니다.Cloud TPU 작업자는 동기식으로 작동합니다. 즉, 각 작업자가 동시에 같은 단계를 수행합니다.
텐서플로우 에스티메이터를 TPUEstimator로 변환
작은 모델을 먼저 포팅하여 해당 동작을 테스트하는 것이 좋습니다. 이렇게 하면 TPUEstimator의 기본 개념에 충실할 수 있습니다. 모델이 실행되면 점차 기능을 늘려가세요.
Cloud TPU에서 실행할 여러 샘플 모델 및 관련 안내는 가이드를 참조하세요. GitHub에서 다른 모델도 사용할 수 있습니다.
tf.estimator.Estimator 클래스의 코드를 변환하여 tf.contrib.tpu.TPUEstimator를 사용하려면 다음 항목을 변경합니다.
tf.estimator.RunConfig를tf.contrib.tpu.RunConfig로 변경합니다.TPUConfig(tf.contrib.tpu.RunConfig일부)를 설정하여iterations_per_loop을 지정합니다.iterations_per_loop은 한 번의session.run호출을 Cloud TPU에서 실행할 반복 횟수입니다(학습 루프당).
Cloud TPU는 호스트로 복귀하기 전에 학습 루프를 지정된 횟수만큼 반복 실행합니다. 모든 Cloud TPU 반복이 실행될 때까지 체크포인트나 요약이 저장되지 않습니다.
model_fn에서tf.contrib.tpu.CrossShardOptimizer를 사용하여 옵티마이저를 래핑합니다. 예를 들면 다음과 같습니다.optimizer = tf.contrib.tpu.CrossShardOptimizer( tf.train.GradientDescentOptimizer(learning_rate=learning_rate))tf.estimator.Estimator를tf.contrib.tpu.TPUEstimator로 변경합니다.
기본 RunConfig는 100단계마다 텐서보드의 요약을 저장하고 10분마다 체크포인트를 작성합니다.
FAQ
입력 파이프라인에 tf.data가 필요한 이유는 무엇인가요?
두 가지 이유가 있습니다.
TPU 연산이
worker에서 실행되는 동안 애플리케이션 코드는 클라이언트에서 실행되기 때문입니다. 입력 파이프라인 작업은 최적의 성능을 위해 작업자로 실행되어야 합니다.tf.data는 작업자로 작업을 실행합니다.TPU 실행 비용을 분할 상환하려면 모델 학습 단계를
tf.while_loop로 래핑해야 합니다. 여기서 하나의Session.run이 1회의 학습 루프 동안 여러 번 반복됩니다. 현재로서는tf.data만tf.while_loop로 래핑할 수 있습니다.
모델 학습 성능을 어떻게 평가할 수 있나요?
텐서보드에 제공된 프로파일러를 사용하여 모델 학습 성능을 프로파일링할 수 있습니다.