Сохраните дату! Google I / O возвращается 18-20 мая Зарегистрируйтесь сейчас
Эта страница переведена с помощью Cloud Translation API.
Switch to English

Используйте TPU

Посмотреть на 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, которые отличаются от локального процесса, в котором запущена пользовательская программа python. Таким образом, необходимо выполнить некоторую работу по инициализации, чтобы подключиться к удаленному кластеру и инициализировать TPU. Обратите внимание, что аргумент tpu для TPUClusterResolver - это специальный адрес только для Colab. В случае, если вы работаете в Google Compute Engine (GCE), вам следует вместо этого передать имя своего CloudTPU.

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:Initializing the TPU system: grpc://10.240.1.74:8470
INFO:tensorflow:Initializing the TPU system: grpc://10.240.1.74: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:3', 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')]

Размещение устройства вручную

После инициализации 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: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, вы можете просто передать его в API strategy.run . Ниже приведен пример того, что все ядра будут получать одинаковые входные данные (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)
}

Классификация по ТПУ

Поскольку мы изучили основные концепции, пришло время взглянуть на более конкретный пример. В этом руководстве показано, как использовать стратегию распространения tf.distribute.TPUStrategy для управления облачным TPU и обучения модели Keras.

Определите модель 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 API имеет решающее значение при использовании Cloud TPU, поскольку невозможно использовать Cloud TPU, если вы не сможете достаточно быстроtf.data.Dataset им данные. Подробную информацию о производительности набора данных см. В Руководстве по производительности входного конвейера .

Для всех, кроме простейших экспериментов (с использованием tf.data.Dataset.from_tensor_slices или других данных в графике) вам нужно будет хранить все файлы данных, считываемые набором данных, в сегментах Google Cloud Storage (GCS).

В большинстве случаев рекомендуется преобразовать ваши данные в формат TFRecord и использовать tf.data.TFRecordDataset для их чтения. См. Руководство по TFRecord и tf.Example для получения подробной информации о том, как это сделать. Однако это не является жестким требованием, и вы можете использовать другие средства чтения наборов данных ( FixedLengthRecordDataset или TextLineDataset ), если хотите.

Небольшие наборы данных можно полностью загрузить в память с помощью tf.data.Dataset.cache .

Независимо от используемого формата данных настоятельно рекомендуется использовать файлы большого размера, порядка 100 МБ. Это особенно важно в этой сетевой настройке, поскольку накладные расходы на открытие файла значительно выше.

Здесь вы должны использовать модуль 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)

  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

Обучите модель с помощью API высокого уровня Keras

Вы можете просто обучить модель с помощью API-интерфейсов Keras fit / compile. Здесь нет ничего специфичного для 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 38ms/step - loss: 0.3196 - sparse_categorical_accuracy: 0.8996 - val_loss: 0.0467 - val_sparse_categorical_accuracy: 0.9855
Epoch 2/5
300/300 [==============================] - 7s 25ms/step - loss: 0.0381 - sparse_categorical_accuracy: 0.9878 - val_loss: 0.0379 - val_sparse_categorical_accuracy: 0.9885
Epoch 3/5
300/300 [==============================] - 8s 26ms/step - loss: 0.0213 - sparse_categorical_accuracy: 0.9933 - val_loss: 0.0366 - val_sparse_categorical_accuracy: 0.9892
Epoch 4/5
300/300 [==============================] - 8s 25ms/step - loss: 0.0130 - sparse_categorical_accuracy: 0.9961 - val_loss: 0.0327 - val_sparse_categorical_accuracy: 0.9902
Epoch 5/5
300/300 [==============================] - 8s 25ms/step - loss: 0.0083 - sparse_categorical_accuracy: 0.9975 - val_loss: 0.0465 - val_sparse_categorical_accuracy: 0.9885
<tensorflow.python.keras.callbacks.History at 0x7fc0906c90b8>

Чтобы уменьшить накладные расходы python и максимизировать производительность вашего TPU, попробуйте экспериментальный аргумент 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)
WARNING:tensorflow:The argument `steps_per_execution` is no longer experimental. Pass `steps_per_execution` instead of `experimental_steps_per_execution`.
WARNING:tensorflow:The argument `steps_per_execution` is no longer experimental. Pass `steps_per_execution` instead of `experimental_steps_per_execution`.
Epoch 1/5
300/300 [==============================] - 14s 46ms/step - loss: 0.2399 - sparse_categorical_accuracy: 0.9279 - val_loss: 0.0510 - val_sparse_categorical_accuracy: 0.9838
Epoch 2/5
300/300 [==============================] - 5s 15ms/step - loss: 0.0406 - sparse_categorical_accuracy: 0.9876 - val_loss: 0.0409 - val_sparse_categorical_accuracy: 0.9882
Epoch 3/5
300/300 [==============================] - 5s 15ms/step - loss: 0.0203 - sparse_categorical_accuracy: 0.9936 - val_loss: 0.0394 - val_sparse_categorical_accuracy: 0.9879
Epoch 4/5
300/300 [==============================] - 5s 15ms/step - loss: 0.0132 - sparse_categorical_accuracy: 0.9959 - val_loss: 0.0410 - val_sparse_categorical_accuracy: 0.9879
Epoch 5/5
300/300 [==============================] - 5s 15ms/step - loss: 0.0140 - sparse_categorical_accuracy: 0.9953 - val_loss: 0.0508 - val_sparse_categorical_accuracy: 0.9863
<tensorflow.python.keras.callbacks.History at 0x7fbc743646d8>

Обучите модель с помощью настраиваемого цикла обучения.

Вы также можете создавать и обучать свои модели tf.function с tf.distribute API-интерфейсов tf.function и tf.distribute . strategy.experimental_distribute_datasets_from_function API используется для распределения набора данных с учетом функции набора данных. Обратите внимание, что размер пакета, передаваемого в набор данных, в этом случае будет соответствовать размеру пакета реплик, а не глобальному размеру пакета. Чтобы узнать больше, ознакомьтесь с учебным курсом Custom training with tf.distribute.Strategy .

Сначала создайте модель, наборы данных и tf.функции.

# 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),))
WARNING:tensorflow:From <ipython-input-1-2f075cabff81>: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-2f075cabff81>: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.1237, accuracy: 96.17%
Epoch: 1/5
Current step: 600, training loss: 0.0333, accuracy: 98.98%
Epoch: 2/5
Current step: 900, training loss: 0.0185, accuracy: 99.37%
Epoch: 3/5
Current step: 1200, training loss: 0.0135, accuracy: 99.54%
Epoch: 4/5
Current step: 1500, training loss: 0.0092, accuracy: 99.7%

Повышение производительности за счет нескольких шагов в tf.function

Производительность можно улучшить, выполнив несколько шагов в tf.function . Это достигается путем помещения вызова strategy.run в tf.range внутри tf.function , AutoGraph преобразует его в tf.while_loop на tf.while_loop TPU.

Хотя с лучшей производительностью, есть компромиссы по сравнению с одним шагом внутри tf.function . Выполнение нескольких шагов в tf.function менее гибко, вы не можете запускать что-то с нетерпением или произвольный код Python внутри шагов.

@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.009, accuracy: 99.71%

Следующие шаги