 TensorFlow.org에서보기 |  Google Colab에서 실행 |  GitHub에서 소스보기 |  노트북 다운로드 |
Cloud TPU에 대한 실험적 지원은 현재 Keras 및 Google Colab에 제공됩니다. 이 Colab 노트북을 실행하기 전에 노트북 설정 (런타임> 런타임 유형 변경> 하드웨어 가속기> TPU)을 확인하여 하드웨어 가속기가 TPU인지 확인하십시오.
설정
import tensorflow as tf
import os
import tensorflow_datasets as tfds
TPU 초기화
TPU는 일반적으로 사용자 파이썬 프로그램을 실행하는 로컬 프로세스와 다른 Cloud TPU 작업자에 있습니다. 따라서 원격 클러스터에 연결하고 TPU를 초기화하려면 일부 초기화 작업을 수행해야합니다. TPUClusterResolver 대한 tpu 인수는 TPUClusterResolver 에 대한 특별한 주소입니다. GCE (Google Compute Engine)에서 실행중인 경우 대신 CloudTPU 이름을 전달해야합니다.
resolver = tf.distribute.cluster_resolver.TPUClusterResolver(tpu='grpc://' + os.environ['COLAB_TPU_ADDR'])
tf.config.experimental_connect_to_cluster(resolver)
# This is the TPU initialization code that has to be at the beginning.
tf.tpu.experimental.initialize_tpu_system(resolver)
print("All devices: ", tf.config.list_logical_devices('TPU'))
INFO:tensorflow:Initializing the TPU system: grpc://10.240.1.2:8470 INFO:tensorflow:Initializing the TPU system: grpc://10.240.1.2:8470 INFO:tensorflow:Clearing out eager caches INFO:tensorflow:Clearing out eager caches INFO:tensorflow:Finished initializing TPU system. INFO:tensorflow:Finished initializing TPU system. All devices: [LogicalDevice(name='/job:worker/replica:0/task:0/device:TPU:7', device_type='TPU'), LogicalDevice(name='/job:worker/replica:0/task:0/device:TPU:6', device_type='TPU'), LogicalDevice(name='/job:worker/replica:0/task:0/device:TPU:5', device_type='TPU'), LogicalDevice(name='/job:worker/replica:0/task:0/device:TPU:4', device_type='TPU'), LogicalDevice(name='/job:worker/replica:0/task:0/device:TPU:0', device_type='TPU'), LogicalDevice(name='/job:worker/replica:0/task:0/device:TPU:1', device_type='TPU'), LogicalDevice(name='/job:worker/replica:0/task:0/device:TPU:2', device_type='TPU'), LogicalDevice(name='/job:worker/replica:0/task:0/device:TPU:3', device_type='TPU')]
수동 장치 배치
TPU가 초기화 된 후 수동 장치 배치를 사용하여 단일 TPU 장치에 계산을 배치 할 수 있습니다.
a = tf.constant([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]])
b = tf.constant([[1.0, 2.0], [3.0, 4.0], [5.0, 6.0]])
with tf.device('/TPU:0'):
c = tf.matmul(a, b)
print("c device: ", c.device)
print(c)
c device: /job:worker/replica:0/task:0/device:TPU:0 tf.Tensor( [[22. 28.] [49. 64.]], shape=(2, 2), dtype=float32)
유통 전략
대부분의 경우 사용자는 데이터를 병렬로 여러 TPU에서 모델을 실행하려고합니다. 배포 전략은 CPU, GPU 또는 TPU에서 모델을 구동하는 데 사용할 수있는 추상화입니다. 배포 전략을 바꾸면 모델이 지정된 장치에서 실행됩니다. 자세한 내용은 배포 전략 안내서 를 참조하십시오.
먼저 TPUStrategy 객체를 만듭니다.
strategy = tf.distribute.TPUStrategy(resolver)
INFO:tensorflow:Found TPU system: INFO:tensorflow:Found TPU system: INFO:tensorflow:*** Num TPU Cores: 8 INFO:tensorflow:*** Num TPU Cores: 8 INFO:tensorflow:*** Num TPU Workers: 1 INFO:tensorflow:*** Num TPU Workers: 1 INFO:tensorflow:*** Num TPU Cores Per Worker: 8 INFO:tensorflow:*** Num TPU Cores Per Worker: 8 INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:localhost/replica:0/task:0/device:CPU:0, CPU, 0, 0) INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:localhost/replica:0/task:0/device:CPU:0, CPU, 0, 0) INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:localhost/replica:0/task:0/device:XLA_CPU:0, XLA_CPU, 0, 0) INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:localhost/replica:0/task:0/device:XLA_CPU:0, XLA_CPU, 0, 0) INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:CPU:0, CPU, 0, 0) INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:CPU:0, CPU, 0, 0) INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:0, TPU, 0, 0) INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:0, TPU, 0, 0) INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:1, TPU, 0, 0) INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:1, TPU, 0, 0) INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:2, TPU, 0, 0) INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:2, TPU, 0, 0) INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:3, TPU, 0, 0) INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:3, TPU, 0, 0) INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:4, TPU, 0, 0) INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:4, TPU, 0, 0) INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:5, TPU, 0, 0) INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:5, TPU, 0, 0) INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:6, TPU, 0, 0) INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:6, TPU, 0, 0) INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:7, TPU, 0, 0) INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU:7, TPU, 0, 0) INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU_SYSTEM:0, TPU_SYSTEM, 0, 0) INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:TPU_SYSTEM:0, TPU_SYSTEM, 0, 0) INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:XLA_CPU:0, XLA_CPU, 0, 0) INFO:tensorflow:*** Available Device: _DeviceAttributes(/job:worker/replica:0/task:0/device:XLA_CPU:0, XLA_CPU, 0, 0)
모든 TPU 코어에서 실행될 수 있도록 계산을 복제하기 위해 간단히 strategy.run API로 전달할 수 있습니다. 아래는 모든 코어가 동일한 입력 (a, b) 얻고 각 코어에서 독립적으로 matmul을 수행하는 예입니다. 출력은 모든 복제본의 값이됩니다.
@tf.function
def matmul_fn(x, y):
z = tf.matmul(x, y)
return z
z = strategy.run(matmul_fn, args=(a, b))
print(z)
PerReplica:{
0: tf.Tensor(
[[22. 28.]
[49. 64.]], shape=(2, 2), dtype=float32),
1: tf.Tensor(
[[22. 28.]
[49. 64.]], shape=(2, 2), dtype=float32),
2: tf.Tensor(
[[22. 28.]
[49. 64.]], shape=(2, 2), dtype=float32),
3: tf.Tensor(
[[22. 28.]
[49. 64.]], shape=(2, 2), dtype=float32),
4: tf.Tensor(
[[22. 28.]
[49. 64.]], shape=(2, 2), dtype=float32),
5: tf.Tensor(
[[22. 28.]
[49. 64.]], shape=(2, 2), dtype=float32),
6: tf.Tensor(
[[22. 28.]
[49. 64.]], shape=(2, 2), dtype=float32),
7: tf.Tensor(
[[22. 28.]
[49. 64.]], shape=(2, 2), dtype=float32)
}
TPU 분류
기본 개념을 배웠으므로보다 구체적인 예를 살펴볼 차례입니다. 이 가이드는 배포 전략 tf.distribute.experimental.TPUStrategy 를 사용하여 Cloud TPU를 구동하고 tf.distribute.experimental.TPUStrategy 모델을 학습하는 방법을 보여줍니다.
Keras 모델 정의
아래는 Keras를 사용한 MNIST 모델의 정의이며 CPU 또는 GPU에서 사용하는 것과 변경되지 않았습니다. Keras 모델 작성은 strategy.scope 안에 있어야하므로 각 TPU 디바이스에서 변수를 작성할 수 있습니다. 코드의 다른 부분은 전략 범위 내에있을 필요는 없습니다.
def create_model():
return tf.keras.Sequential(
[tf.keras.layers.Conv2D(256, 3, activation='relu', input_shape=(28, 28, 1)),
tf.keras.layers.Conv2D(256, 3, activation='relu'),
tf.keras.layers.Flatten(),
tf.keras.layers.Dense(256, activation='relu'),
tf.keras.layers.Dense(128, activation='relu'),
tf.keras.layers.Dense(10)])
입력 데이터 세트
Cloud TPU를 사용하는 경우 tf.data.Dataset API를 효율적으로 사용하는 것이 중요합니다. 데이터를 충분히 빠르게 공급할 수 없다면 Cloud TPU를 사용할 수 없습니다. 데이터 세트 성능에 대한 자세한 내용은 입력 파이프 라인 성능 안내서 를 참조하십시오.
가장 간단한 실험 ( tf.data.Dataset.from_tensor_slices 또는 기타 in-graph 데이터 사용)을 제외한 모든 데이터 파일에서 읽은 모든 데이터 파일을 Google Cloud Storage (GCS) 버킷에 저장해야합니다.
대부분의 사용 사례의 경우 데이터를 TFRecord 형식으로 변환하고 tf.data.TFRecordDataset 을 사용하여 tf.data.TFRecordDataset 를 읽는 것이 좋습니다. 이 작업을 수행하는 방법에 대한 자세한 내용은 TFRecord 및 tf.Example 학습서 를 참조하십시오. 그러나 이것은 어려운 요구 사항이 아니며 원하는 경우 다른 데이터 세트 리더 ( FixedLengthRecordDataset 또는 TextLineDataset )를 사용할 수 있습니다.
작은 데이터 세트는 tf.data.Dataset.cache 사용하여 메모리에 완전히로드 할 수 있습니다.
사용 된 데이터 형식에 관계없이 100MB 정도의 큰 파일을 사용하는 것이 좋습니다. 파일을 여는 오버 헤드가 상당히 높기 때문에이 네트워크 설정에서 특히 중요합니다.
여기서 tensorflow_datasets 모듈을 사용하여 MNIST 교육 데이터의 사본을 tensorflow_datasets 합니다. try_gcs 는 퍼블릭 GCS 버킷에서 사용 가능한 사본을 사용하도록 지정되었습니다. 이를 지정하지 않으면 TPU가 다운로드 된 데이터에 액세스 할 수 없습니다.
def get_dataset(batch_size, is_training=True):
split = 'train' if is_training else 'test'
dataset, info = tfds.load(name='mnist', split=split, with_info=True,
as_supervised=True, try_gcs=True)
def scale(image, label):
image = tf.cast(image, tf.float32)
image /= 255.0
return image, label
dataset = dataset.map(scale)
# Only shuffle and repeat the dataset in training. The advantage to have a
# infinite dataset for training is to avoid the potential last partial batch
# in each epoch, so users don't need to think about scaling the gradients
# based on the actual batch size.
if is_training:
dataset = dataset.shuffle(10000)
dataset = dataset.repeat()
dataset = dataset.batch(batch_size)
return dataset
Keras 고급 API를 사용하여 모델 학습
Keras fit / compile API를 사용하여 모델을 간단히 학습 할 수 있습니다. 여기에 TPU에만 해당되는 것은 없으며 여러 GPU가 있고 TPUStrategy 대신 MirroredStrategy 사용하는 경우 아래 동일한 코드를 작성하십시오. 자세한 내용 은 Keras를 사용한 분산 교육 튜토리얼을 확인하십시오.
with strategy.scope():
model = create_model()
model.compile(optimizer='adam',
loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True),
metrics=['sparse_categorical_accuracy'])
batch_size = 200
steps_per_epoch = 60000 // batch_size
validation_steps = 10000 // batch_size
train_dataset = get_dataset(batch_size, is_training=True)
test_dataset = get_dataset(batch_size, is_training=False)
model.fit(train_dataset,
epochs=5,
steps_per_epoch=steps_per_epoch,
validation_data=test_dataset,
validation_steps=validation_steps)
Epoch 1/5 WARNING:tensorflow:From /tmpfs/src/tf_docs_env/lib/python3.6/site-packages/tensorflow/python/data/ops/multi_device_iterator_ops.py:601: get_next_as_optional (from tensorflow.python.data.ops.iterator_ops) is deprecated and will be removed in a future version. Instructions for updating: Use `tf.data.Iterator.get_next_as_optional()` instead. Warning:tensorflow:From /tmpfs/src/tf_docs_env/lib/python3.6/site-packages/tensorflow/python/data/ops/multi_device_iterator_ops.py:601: get_next_as_optional (from tensorflow.python.data.ops.iterator_ops) is deprecated and will be removed in a future version. Instructions for updating: Use `tf.data.Iterator.get_next_as_optional()` instead. 1/300 [..............................] - ETA: 20:25 - loss: 2.3062 - sparse_categorical_accuracy: 0.0500WARNING:tensorflow:Callbacks method `on_train_batch_end` is slow compared to the batch time (batch time: 0.0015s vs `on_train_batch_end` time: 0.0237s). Check your callbacks. Warning:tensorflow:Callbacks method `on_train_batch_end` is slow compared to the batch time (batch time: 0.0015s vs `on_train_batch_end` time: 0.0237s). Check your callbacks. 300/300 [==============================] - ETA: 0s - loss: 0.1450 - sparse_categorical_accuracy: 0.9561WARNING:tensorflow:Callbacks method `on_test_batch_end` is slow compared to the batch time (batch time: 0.0013s vs `on_test_batch_end` time: 0.0097s). Check your callbacks. Warning:tensorflow:Callbacks method `on_test_batch_end` is slow compared to the batch time (batch time: 0.0013s vs `on_test_batch_end` time: 0.0097s). Check your callbacks. 300/300 [==============================] - 15s 50ms/step - loss: 0.1450 - sparse_categorical_accuracy: 0.9561 - val_loss: 0.0456 - val_sparse_categorical_accuracy: 0.9852 Epoch 2/5 300/300 [==============================] - 8s 28ms/step - loss: 0.0357 - sparse_categorical_accuracy: 0.9890 - val_loss: 0.0372 - val_sparse_categorical_accuracy: 0.9884 Epoch 3/5 300/300 [==============================] - 9s 28ms/step - loss: 0.0193 - sparse_categorical_accuracy: 0.9940 - val_loss: 0.0557 - val_sparse_categorical_accuracy: 0.9835 Epoch 4/5 300/300 [==============================] - 9s 29ms/step - loss: 0.0141 - sparse_categorical_accuracy: 0.9954 - val_loss: 0.0405 - val_sparse_categorical_accuracy: 0.9883 Epoch 5/5 300/300 [==============================] - 9s 29ms/step - loss: 0.0092 - sparse_categorical_accuracy: 0.9967 - val_loss: 0.0428 - val_sparse_categorical_accuracy: 0.9887 <tensorflow.python.keras.callbacks.History at 0x7f480c793240>
파이썬 오버 헤드를 줄이고 TPU의 성능을 최대화하려면 Model.compile 대한 실험 experimental_steps_per_execution 인수를 Model.compile . 여기에서 처리량이 약 50 % 증가합니다.
with strategy.scope():
model = create_model()
model.compile(optimizer='adam',
# Anything between 2 and `steps_per_epoch` could help here.
experimental_steps_per_execution = 50,
loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True),
metrics=['sparse_categorical_accuracy'])
model.fit(train_dataset,
epochs=5,
steps_per_epoch=steps_per_epoch,
validation_data=test_dataset,
validation_steps=validation_steps)
Epoch 1/5 300/300 [==============================] - 16s 52ms/step - loss: 0.1337 - sparse_categorical_accuracy: 0.9583 - val_loss: 0.0521 - val_sparse_categorical_accuracy: 0.9840 Epoch 2/5 300/300 [==============================] - 5s 16ms/step - loss: 0.0331 - sparse_categorical_accuracy: 0.9898 - val_loss: 0.0360 - val_sparse_categorical_accuracy: 0.9884 Epoch 3/5 300/300 [==============================] - 5s 16ms/step - loss: 0.0188 - sparse_categorical_accuracy: 0.9939 - val_loss: 0.0405 - val_sparse_categorical_accuracy: 0.9887 Epoch 4/5 300/300 [==============================] - 5s 16ms/step - loss: 0.0117 - sparse_categorical_accuracy: 0.9961 - val_loss: 0.0808 - val_sparse_categorical_accuracy: 0.9820 Epoch 5/5 300/300 [==============================] - 5s 16ms/step - loss: 0.0119 - sparse_categorical_accuracy: 0.9962 - val_loss: 0.0488 - val_sparse_categorical_accuracy: 0.9862 <tensorflow.python.keras.callbacks.History at 0x7f47ac791240>
사용자 지정 훈련 루프를 사용하여 모델을 훈련시킵니다.
tf.function 및 tf.distribute API를 직접 사용하여 모델을 작성하고 학습 할 수도 있습니다. strategy.experimental_distribute_datasets_from_function API는 데이터 세트 함수가 주어진 데이터 세트를 분배하는 데 사용됩니다. 이 경우 데이터 집합에 전달 된 배치 크기는 전역 배치 크기가 아닌 복제본 배치 크기입니다. 자세한 내용 은 tf.distribute.Strategy를 사용한 사용자 지정 교육 자습서를 확인하십시오.
먼저 모델, 데이터 세트 및 tf.functions를 작성하십시오.
# Create the model, optimizer and metrics inside strategy scope, so that the
# variables can be mirrored on each device.
with strategy.scope():
model = create_model()
optimizer = tf.keras.optimizers.Adam()
training_loss = tf.keras.metrics.Mean('training_loss', dtype=tf.float32)
training_accuracy = tf.keras.metrics.SparseCategoricalAccuracy(
'training_accuracy', dtype=tf.float32)
# Calculate per replica batch size, and distribute the datasets on each TPU
# worker.
per_replica_batch_size = batch_size // strategy.num_replicas_in_sync
train_dataset = strategy.experimental_distribute_datasets_from_function(
lambda _: get_dataset(per_replica_batch_size, is_training=True))
@tf.function
def train_step(iterator):
"""The step function for one training step"""
def step_fn(inputs):
"""The computation to run on each TPU device."""
images, labels = inputs
with tf.GradientTape() as tape:
logits = model(images, training=True)
loss = tf.keras.losses.sparse_categorical_crossentropy(
labels, logits, from_logits=True)
loss = tf.nn.compute_average_loss(loss, global_batch_size=batch_size)
grads = tape.gradient(loss, model.trainable_variables)
optimizer.apply_gradients(list(zip(grads, model.trainable_variables)))
training_loss.update_state(loss * strategy.num_replicas_in_sync)
training_accuracy.update_state(labels, logits)
strategy.run(step_fn, args=(next(iterator),))
그런 다음 훈련 루프를 실행하십시오.
steps_per_eval = 10000 // batch_size
train_iterator = iter(train_dataset)
for epoch in range(5):
print('Epoch: {}/5'.format(epoch))
for step in range(steps_per_epoch):
train_step(train_iterator)
print('Current step: {}, training loss: {}, accuracy: {}%'.format(
optimizer.iterations.numpy(),
round(float(training_loss.result()), 4),
round(float(training_accuracy.result()) * 100, 2)))
training_loss.reset_states()
training_accuracy.reset_states()
Epoch: 0/5 Current step: 300, training loss: 0.1335, accuracy: 95.8% Epoch: 1/5 Current step: 600, training loss: 0.0344, accuracy: 98.93% Epoch: 2/5 Current step: 900, training loss: 0.0196, accuracy: 99.35% Epoch: 3/5 Current step: 1200, training loss: 0.0119, accuracy: 99.61% Epoch: 4/5 Current step: 1500, training loss: 0.0102, accuracy: 99.65%
tf.function 내에서 여러 단계로 성능 향상
tf.function 내에서 여러 단계를 실행하여 성능을 향상시킬 수 있습니다. 이것은 tf.range 안에 tf.function 사용하여 strategy.run 호출을 tf.range 으로써 달성됩니다. AutoGraph는이를 TPU 워커에서 tf.while_loop 로 변환합니다.
더 나은 성능이지만 tf.function 내부의 단일 단계와 비교할 때 장단점이 있습니다. tf.function 에서 여러 단계를 실행하면 유연성이 떨어 tf.function 단계 내에서 열성적이거나 임의의 파이썬 코드를 실행할 수 없습니다.
@tf.function
def train_multiple_steps(iterator, steps):
"""The step function for one training step"""
def step_fn(inputs):
"""The computation to run on each TPU device."""
images, labels = inputs
with tf.GradientTape() as tape:
logits = model(images, training=True)
loss = tf.keras.losses.sparse_categorical_crossentropy(
labels, logits, from_logits=True)
loss = tf.nn.compute_average_loss(loss, global_batch_size=batch_size)
grads = tape.gradient(loss, model.trainable_variables)
optimizer.apply_gradients(list(zip(grads, model.trainable_variables)))
training_loss.update_state(loss * strategy.num_replicas_in_sync)
training_accuracy.update_state(labels, logits)
for _ in tf.range(steps):
strategy.run(step_fn, args=(next(iterator),))
# Convert `steps_per_epoch` to `tf.Tensor` so the `tf.function` won't get
# retraced if the value changes.
train_multiple_steps(train_iterator, tf.convert_to_tensor(steps_per_epoch))
print('Current step: {}, training loss: {}, accuracy: {}%'.format(
optimizer.iterations.numpy(),
round(float(training_loss.result()), 4),
round(float(training_accuracy.result()) * 100, 2)))
Current step: 1800, training loss: 0.0087, accuracy: 99.72%
다음 단계
- Google Cloud TPU 설명서-Google Cloud TPU를 설정하고 실행합니다.
- TensorFlow 를 사용한 분산 교육 -분산 전략을 사용하는 방법과 모범 사례를 보여주는 많은 예제에 대한 링크.
- TensorFlow 공식 모델 -Cloud TPU와 호환되는 최신 TensorFlow 2.x 모델의 예입니다.
- Google Cloud TPU 성능 가이드 -애플리케이션에 대한 Cloud TPU 구성 매개 변수를 조정하여 Cloud TPU 성능을 더욱 향상시킵니다.