년 10 월 19 일 ML 심포지엄에서 여성에 실질적으로 우리와 함께 지금 등록

TPU 사용

TensorFlow.org에서 보기 Google Colab에서 실행 GitHub에서 소스 보기 노트북 다운로드

런타임> 변경 실행시의 형태> 하드웨어 가속기> TPU :이 Colab 노트북을 실행하기 전에 하드웨어 가속기는 노트북 설정을 확인하여 TPU가 있는지 확인합니다.

설정

import tensorflow as tf

import os
import tensorflow_datasets as tfds
/home/kbuilder/.local/lib/python3.6/site-packages/requests/__init__.py:104: RequestsDependencyWarning: urllib3 (1.26.6) or chardet (2.3.0)/charset_normalizer (2.0.4) doesn't match a supported version!
  RequestsDependencyWarning)

TPU 초기화

TPU는 일반적으로 사용자의 Python 프로그램을 실행하는 로컬 프로세스와 다른 Cloud TPU 작업자입니다. 따라서 원격 클러스터에 연결하고 TPU를 초기화하려면 일부 초기화 작업을 수행해야 합니다. 합니다 tpu 에 인수 tf.distribute.cluster_resolver.TPUClusterResolver 단지 Colab위한 특별한 주소입니다. Google Compute Engine(GCE)에서 코드를 실행하는 경우 대신 Cloud TPU의 이름을 전달해야 합니다.

resolver = tf.distribute.cluster_resolver.TPUClusterResolver(tpu='')
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:Clearing out eager caches
INFO:tensorflow:Clearing out eager caches
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:Finished initializing TPU system.
INFO:tensorflow:Finished initializing TPU system.
All devices:  [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'), LogicalDevice(name='/job:worker/replica:0/task:0/device:TPU:4', 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:6', device_type='TPU'), LogicalDevice(name='/job:worker/replica:0/task:0/device:TPU:7', 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에서 모델을 실행합니다. 여러 TPU(또는 다른 가속기)에 모델을 배포하기 위해 TensorFlow는 여러 배포 전략을 제공합니다. 배포 전략을 교체할 수 있으며 모델은 지정된(TPU) 장치에서 실행됩니다. 체크 유통 전략 가이드 자세한 내용은.

이를 증명하기 위해, 생성 tf.distribute.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: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) 각각 독립적 코어의 행렬 곱셈을 수행하는 단계. 출력은 모든 복제본의 값이 됩니다.

@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.TPUStrategy 클라우드 TPU에 Keras 모델을 훈련 -to.

Keras 모델 정의

(A)의 정의로 시작 Sequential Keras를 사용하여 MNIST 데이터 세트에 영상 분류에 대한 Keras 모델. 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)])

데이터세트 로드

의 효율적인 사용 tf.data.Dataset 당신이 신속하게 그 데이터를 공급 할 수있는 경우를 제외하고는 클라우드 TPU를를 사용하는 것은 불가능 한, 클라우드 TPU를 사용하는 경우 API는 매우 중요하다. 당신은에서 데이터 세트의 성능에 대해 자세히 배울 수있는 입력 파이프 라인 성능 가이드 .

모든하지만 (사용하여 간단한 실험을 위해 tf.data.Dataset.from_tensor_slices 또는 기타의 그래프 데이터), 당신은 모든 데이터 파일은 Google 클라우드 스토리지에서 데이터 집합 (GCS) 버킷 읽을 저장해야합니다.

대부분의 사용 사례를 들어,으로 데이터를 변환하는 것이 좋습니다 TFRecord 형식과 사용 tf.data.TFRecordDataset 읽을 수 있습니다. 체크 TFRecord 및 tf.Example 튜토리얼 이 작업을 수행하는 방법에 대한 자세한 내용을. 그것은 힘든 요구 사항이 아니므로이 같은 다른 데이터 세트의 리더, 사용할 수 있습니다 tf.data.FixedLengthRecordDataset 또는 tf.data.TextLineDataset .

다음을 사용하여 메모리에 전체 작은 데이터 세트를로드 할 수 있습니다 tf.data.Dataset.cache .

사용하는 데이터 형식에 관계없이 100MB 정도의 대용량 파일을 사용하는 것이 좋습니다. 이것은 파일을 여는 오버헤드가 훨씬 높기 때문에 이 네트워크 설정에서 특히 중요합니다.

아래의 코드와 같이, 당신은 사용해야 tensorflow_datasets MNIST 교육 및 테스트 데이터의 사본을 얻을 모듈. 참고 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)

  # Normalize the input data.
  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 of having an
  # infinite dataset for training is to avoid the potential last partial batch
  # in each epoch, so that you 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를. 당신이 mutliple GPU를하고 사용하는 것처럼이 단계 - 당신이 코드를 작성 아무것도 TPU 별이 MirroredStrategy 대신의 TPUStrategy . 당신은 더 배울 수 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
300/300 [==============================] - 19s 35ms/step - loss: 0.1352 - sparse_categorical_accuracy: 0.9583 - val_loss: 0.0404 - val_sparse_categorical_accuracy: 0.9863
Epoch 2/5
300/300 [==============================] - 6s 21ms/step - loss: 0.0346 - sparse_categorical_accuracy: 0.9893 - val_loss: 0.0376 - val_sparse_categorical_accuracy: 0.9882
Epoch 3/5
300/300 [==============================] - 7s 22ms/step - loss: 0.0199 - sparse_categorical_accuracy: 0.9934 - val_loss: 0.0432 - val_sparse_categorical_accuracy: 0.9872
Epoch 4/5
300/300 [==============================] - 6s 21ms/step - loss: 0.0132 - sparse_categorical_accuracy: 0.9957 - val_loss: 0.0475 - val_sparse_categorical_accuracy: 0.9870
Epoch 5/5
300/300 [==============================] - 6s 21ms/step - loss: 0.0100 - sparse_categorical_accuracy: 0.9970 - val_loss: 0.0443 - val_sparse_categorical_accuracy: 0.9877
<keras.callbacks.History at 0x7faf8829f278>

파이썬 오버 헤드를 줄이고 TPU의 성능을 극대화하기 위해 argument- 전달 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.
                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 [==============================] - 13s 43ms/step - loss: 0.1405 - sparse_categorical_accuracy: 0.9571 - val_loss: 0.0438 - val_sparse_categorical_accuracy: 0.9860
Epoch 2/5
300/300 [==============================] - 3s 10ms/step - loss: 0.0337 - sparse_categorical_accuracy: 0.9890 - val_loss: 0.0503 - val_sparse_categorical_accuracy: 0.9847
Epoch 3/5
300/300 [==============================] - 3s 10ms/step - loss: 0.0170 - sparse_categorical_accuracy: 0.9947 - val_loss: 0.0448 - val_sparse_categorical_accuracy: 0.9884
Epoch 4/5
300/300 [==============================] - 3s 10ms/step - loss: 0.0137 - sparse_categorical_accuracy: 0.9956 - val_loss: 0.0362 - val_sparse_categorical_accuracy: 0.9903
Epoch 5/5
300/300 [==============================] - 3s 10ms/step - loss: 0.0090 - sparse_categorical_accuracy: 0.9971 - val_loss: 0.0458 - val_sparse_categorical_accuracy: 0.9880
<keras.callbacks.History at 0x7fae1054f908>

사용자 지정 훈련 루프를 사용하여 모델 훈련

또한 생성하고 사용하여 모델을 학습 할 수 tf.functiontf.distribute 직접 API를. 당신은 사용할 수 있습니다 strategy.experimental_distribute_datasets_from_function 데이터 집합 기능 주어진 데이터 집합을 배포하는 API를. 아래 예에서 데이터세트에 전달된 배치 크기는 전역 배치 크기가 아닌 복제본당 배치 크기입니다. 더 많은 내용은 체크 아웃 tf.distribute.Strategy의와 정의 훈련 자습서를.

먼저 모델, 데이터세트 및 tf.functions를 생성합니다.

# Create the model, optimizer and metrics inside the 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),))
WARNING:tensorflow:From <ipython-input-1-5625c2a14441>:15: StrategyBase.experimental_distribute_datasets_from_function (from tensorflow.python.distribute.distribute_lib) is deprecated and will be removed in a future version.
Instructions for updating:
rename to distribute_datasets_from_function
WARNING:tensorflow:From <ipython-input-1-5625c2a14441>:15: StrategyBase.experimental_distribute_datasets_from_function (from tensorflow.python.distribute.distribute_lib) is deprecated and will be removed in a future version.
Instructions for updating:
rename to distribute_datasets_from_function

그런 다음 훈련 루프를 실행합니다.

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.1352, accuracy: 95.85%
Epoch: 1/5
Current step: 600, training loss: 0.0355, accuracy: 98.86%
Epoch: 2/5
Current step: 900, training loss: 0.0194, accuracy: 99.37%
Epoch: 3/5
Current step: 1200, training loss: 0.0127, accuracy: 99.6%
Epoch: 4/5
Current step: 1500, training loss: 0.0111, accuracy: 99.65%

내부의 여러 단계와 성능 향상 tf.function

당신은 내에서 여러 단계를 실행하여 성능을 향상시킬 수 있습니다 tf.function . 이것은 포장에 의해 달성된다 strategy.run A를 호출 tf.range 내부 tf.function 하고, 사인은로 변환됩니다 tf.while_loop TPU의 작업자.

개선 된 성능에도 불구하고, 내부의 단일 단계로 실행하는 비교 방법에 장단점이있다 tf.function . A의 여러 단계를 실행 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.0081, accuracy: 99.73%

다음 단계