Индивидуальное обучение с tf.distribute.Strategy

Посмотреть на TensorFlow.org Запускаем в Google Colab Посмотреть исходный код на GitHub Скачать блокнот

В данном руководстве показано , как использовать tf.distribute.Strategy с пользовательскими учебными циклами. Мы обучим простую модель CNN на наборе данных Fashion MNIST. Набор данных Fashion MNIST содержит 60000 изображений поездов размером 28 x 28 и 10000 тестовых изображений размером 28 x 28.

Мы используем настраиваемые циклы обучения для обучения нашей модели, потому что они дают нам гибкость и больший контроль над обучением. Кроме того, легче отлаживать модель и цикл обучения.

# Import TensorFlow
import tensorflow as tf

# Helper libraries
import numpy as np
import os

print(tf.__version__)
2.5.0

Загрузите набор данных Fashion MNIST

fashion_mnist = tf.keras.datasets.fashion_mnist

(train_images, train_labels), (test_images, test_labels) = fashion_mnist.load_data()

# Adding a dimension to the array -> new shape == (28, 28, 1)
# We are doing this because the first layer in our model is a convolutional
# layer and it requires a 4D input (batch_size, height, width, channels).
# batch_size dimension will be added later on.
train_images = train_images[..., None]
test_images = test_images[..., None]

# Getting the images in [0, 1] range.
train_images = train_images / np.float32(255)
test_images = test_images / np.float32(255)

Создайте стратегию распределения переменных и графика

Как tf.distribute.MirroredStrategy стратегия работы?

  • Все переменные и граф модели воспроизводятся на репликах.
  • Входные данные равномерно распределяются по репликам.
  • Каждая реплика вычисляет потери и градиенты для полученных входных данных.
  • Градиенты синхронизируются на всех репликах путем их суммирования.
  • После синхронизации выполняется такое же обновление копий переменных на каждой реплике.
# If the list of devices is not specified in the
# `tf.distribute.MirroredStrategy` constructor, it will be auto-detected.
strategy = tf.distribute.MirroredStrategy()
WARNING:tensorflow:Collective ops is not configured at program startup. Some performance features may not be enabled.
INFO:tensorflow:Using MirroredStrategy with devices ('/job:localhost/replica:0/task:0/device:GPU:0',)
print ('Number of devices: {}'.format(strategy.num_replicas_in_sync))
Number of devices: 1

Настройка входного конвейера

Экспортируйте график и переменные в формат SavedModel, не зависящий от платформы. После сохранения модели вы можете загрузить ее с прицелом или без него.

BUFFER_SIZE = len(train_images)

BATCH_SIZE_PER_REPLICA = 64
GLOBAL_BATCH_SIZE = BATCH_SIZE_PER_REPLICA * strategy.num_replicas_in_sync

EPOCHS = 10

Создайте наборы данных и распространите их:

train_dataset = tf.data.Dataset.from_tensor_slices((train_images, train_labels)).shuffle(BUFFER_SIZE).batch(GLOBAL_BATCH_SIZE) 
test_dataset = tf.data.Dataset.from_tensor_slices((test_images, test_labels)).batch(GLOBAL_BATCH_SIZE) 

train_dist_dataset = strategy.experimental_distribute_dataset(train_dataset)
test_dist_dataset = strategy.experimental_distribute_dataset(test_dataset)

Создать модель

Создание модели с помощью tf.keras.Sequential . Для этого также можно использовать API создания подклассов моделей.

def create_model():
  model = tf.keras.Sequential([
      tf.keras.layers.Conv2D(32, 3, activation='relu'),
      tf.keras.layers.MaxPooling2D(),
      tf.keras.layers.Conv2D(64, 3, activation='relu'),
      tf.keras.layers.MaxPooling2D(),
      tf.keras.layers.Flatten(),
      tf.keras.layers.Dense(64, activation='relu'),
      tf.keras.layers.Dense(10)
    ])

  return model
# Create a checkpoint directory to store the checkpoints.
checkpoint_dir = './training_checkpoints'
checkpoint_prefix = os.path.join(checkpoint_dir, "ckpt")

Определите функцию потерь

Обычно на одной машине с 1 ГП / ЦП потери делятся на количество примеров в пакете входных данных.

Итак, как следует потеря рассчитывается при использовании tf.distribute.Strategy ?

  • Например, предположим, что у вас есть 4 графических процессора и размер пакета 64. Один пакет ввода распределяется по репликам (4 графических процессора), каждая реплика получает ввод размером 16.

  • Модель на каждой реплике выполняет прямой проход с соответствующими входными данными и вычисляет потери. Теперь, вместо деления потерь на количество примеров в соответствующем вводе (BATCH_SIZE_PER_REPLICA = 16), потери следует разделить на GLOBAL_BATCH_SIZE (64).

Зачем это делать?

  • Это должно быть сделано , потому что после того, как градиенты вычисляются на каждую реплике, они синхронизируются через реплики путем суммирования их.

Как это сделать в TensorFlow?

  • Если вы пишете на заказ тренировочный цикл, как в этом учебнике, вы должны суммировать за пример потери и разделить сумму на GLOBAL_BATCH_SIZE: scale_loss = tf.reduce_sum(loss) * (1. / GLOBAL_BATCH_SIZE) или вы можете использовать tf.nn.compute_average_loss , который принимает за потери , например, необязательные веса выборки, и GLOBAL_BATCH_SIZE в качестве аргументов и возвращает масштабированную потерю.

  • Если вы используете в своей модели потери регуляризации, вам необходимо масштабировать значение потерь по количеству реплик. Вы можете сделать это с помощью tf.nn.scale_regularization_loss функции.

  • Использование tf.reduce_mean не рекомендуется. При этом потери делятся на фактический размер пакета реплик, который может меняться шаг за шагом.

  • Это уменьшение и масштабирование осуществляется автоматически в keras model.compile и model.fit

  • При использовании tf.keras.losses классов (как в приведенном ниже примере), то снижение потерь должно быть явно указано , чтобы быть одним из NONE или SUM . AUTO и SUM_OVER_BATCH_SIZE запрещены при использовании tf.distribute.Strategy . AUTO отвергается , так как пользователь должен явно думать о том, что сокращение они хотят , чтобы убедиться , что это правильно в распределенном случае. SUM_OVER_BATCH_SIZE отвергается , потому что в настоящее время он будет делить только на реплику размера партии и оставить Разделительный по количеству реплик для пользователя, который может быть легко пропустить. Поэтому вместо этого мы просим пользователя сделать сокращение самостоятельно.

  • Если labels многомерна, а затем усреднить per_example_loss по количеству элементов в каждом образце. Например, если форма predictions является (batch_size, H, W, n_classes) и labels является (batch_size, H, W) , вам необходимо обновить per_example_loss как: per_example_loss /= tf.cast(tf.reduce_prod(tf.shape(labels)[1:]), tf.float32)

with strategy.scope():
  # Set reduction to `none` so we can do the reduction afterwards and divide by
  # global batch size.
  loss_object = tf.keras.losses.SparseCategoricalCrossentropy(
      from_logits=True,
      reduction=tf.keras.losses.Reduction.NONE)
  def compute_loss(labels, predictions):
    per_example_loss = loss_object(labels, predictions)
    return tf.nn.compute_average_loss(per_example_loss, global_batch_size=GLOBAL_BATCH_SIZE)

Определите метрики для отслеживания потерь и точности

Эти метрики отслеживают потери в тестах, а также точность обучения и тестирования. Вы можете использовать .result() , чтобы получить накопленную статистику в любое время.

with strategy.scope():
  test_loss = tf.keras.metrics.Mean(name='test_loss')

  train_accuracy = tf.keras.metrics.SparseCategoricalAccuracy(
      name='train_accuracy')
  test_accuracy = tf.keras.metrics.SparseCategoricalAccuracy(
      name='test_accuracy')
INFO:tensorflow:Reduce to /job:localhost/replica:0/task:0/device:CPU:0 then broadcast to ('/job:localhost/replica:0/task:0/device:CPU:0',).
INFO:tensorflow:Reduce to /job:localhost/replica:0/task:0/device:CPU:0 then broadcast to ('/job:localhost/replica:0/task:0/device:CPU:0',).
INFO:tensorflow:Reduce to /job:localhost/replica:0/task:0/device:CPU:0 then broadcast to ('/job:localhost/replica:0/task:0/device:CPU:0',).
INFO:tensorflow:Reduce to /job:localhost/replica:0/task:0/device:CPU:0 then broadcast to ('/job:localhost/replica:0/task:0/device:CPU:0',).
INFO:tensorflow:Reduce to /job:localhost/replica:0/task:0/device:CPU:0 then broadcast to ('/job:localhost/replica:0/task:0/device:CPU:0',).
INFO:tensorflow:Reduce to /job:localhost/replica:0/task:0/device:CPU:0 then broadcast to ('/job:localhost/replica:0/task:0/device:CPU:0',).
INFO:tensorflow:Reduce to /job:localhost/replica:0/task:0/device:CPU:0 then broadcast to ('/job:localhost/replica:0/task:0/device:CPU:0',).
INFO:tensorflow:Reduce to /job:localhost/replica:0/task:0/device:CPU:0 then broadcast to ('/job:localhost/replica:0/task:0/device:CPU:0',).
INFO:tensorflow:Reduce to /job:localhost/replica:0/task:0/device:CPU:0 then broadcast to ('/job:localhost/replica:0/task:0/device:CPU:0',).
INFO:tensorflow:Reduce to /job:localhost/replica:0/task:0/device:CPU:0 then broadcast to ('/job:localhost/replica:0/task:0/device:CPU:0',).

Цикл обучения

# model, optimizer, and checkpoint must be created under `strategy.scope`.
with strategy.scope():
  model = create_model()

  optimizer = tf.keras.optimizers.Adam()

  checkpoint = tf.train.Checkpoint(optimizer=optimizer, model=model)
def train_step(inputs):
  images, labels = inputs

  with tf.GradientTape() as tape:
    predictions = model(images, training=True)
    loss = compute_loss(labels, predictions)

  gradients = tape.gradient(loss, model.trainable_variables)
  optimizer.apply_gradients(zip(gradients, model.trainable_variables))

  train_accuracy.update_state(labels, predictions)
  return loss 

def test_step(inputs):
  images, labels = inputs

  predictions = model(images, training=False)
  t_loss = loss_object(labels, predictions)

  test_loss.update_state(t_loss)
  test_accuracy.update_state(labels, predictions)
# `run` replicates the provided computation and runs it
# with the distributed input.
@tf.function
def distributed_train_step(dataset_inputs):
  per_replica_losses = strategy.run(train_step, args=(dataset_inputs,))
  return strategy.reduce(tf.distribute.ReduceOp.SUM, per_replica_losses,
                         axis=None)

@tf.function
def distributed_test_step(dataset_inputs):
  return strategy.run(test_step, args=(dataset_inputs,))

for epoch in range(EPOCHS):
  # TRAIN LOOP
  total_loss = 0.0
  num_batches = 0
  for x in train_dist_dataset:
    total_loss += distributed_train_step(x)
    num_batches += 1
  train_loss = total_loss / num_batches

  # TEST LOOP
  for x in test_dist_dataset:
    distributed_test_step(x)

  if epoch % 2 == 0:
    checkpoint.save(checkpoint_prefix)

  template = ("Epoch {}, Loss: {}, Accuracy: {}, Test Loss: {}, "
              "Test Accuracy: {}")
  print (template.format(epoch+1, train_loss,
                         train_accuracy.result()*100, test_loss.result(),
                         test_accuracy.result()*100))

  test_loss.reset_states()
  train_accuracy.reset_states()
  test_accuracy.reset_states()
Epoch 1, Loss: 0.5044084787368774, Accuracy: 81.87333679199219, Test Loss: 0.3816865086555481, Test Accuracy: 86.5999984741211
Epoch 2, Loss: 0.3375805616378784, Accuracy: 87.8566665649414, Test Loss: 0.3369813859462738, Test Accuracy: 87.76000213623047
Epoch 3, Loss: 0.2896445095539093, Accuracy: 89.50499725341797, Test Loss: 0.299490362405777, Test Accuracy: 89.22000122070312
Epoch 4, Loss: 0.259074866771698, Accuracy: 90.58833312988281, Test Loss: 0.2881558835506439, Test Accuracy: 89.33000183105469
Epoch 5, Loss: 0.2341146171092987, Accuracy: 91.38999938964844, Test Loss: 0.2916182577610016, Test Accuracy: 89.61000061035156
Epoch 6, Loss: 0.21513047814369202, Accuracy: 92.02333068847656, Test Loss: 0.2755740284919739, Test Accuracy: 89.85000610351562
Epoch 7, Loss: 0.1952667236328125, Accuracy: 92.88333129882812, Test Loss: 0.27464523911476135, Test Accuracy: 90.36000061035156
Epoch 8, Loss: 0.17831537127494812, Accuracy: 93.3566665649414, Test Loss: 0.26432710886001587, Test Accuracy: 90.19000244140625
Epoch 9, Loss: 0.16429665684700012, Accuracy: 93.85333251953125, Test Loss: 0.2659859359264374, Test Accuracy: 91.0999984741211
Epoch 10, Loss: 0.1503313183784485, Accuracy: 94.42166900634766, Test Loss: 0.2602477967739105, Test Accuracy: 91.06999969482422

На что следует обратить внимание в приведенном выше примере:

  • Мы итерация над train_dist_dataset и test_dist_dataset используя for x in ... конструкте.
  • Масштабируется потеря является возвращаемым значением distributed_train_step . Это значение агрегируется по репликам , используя tf.distribute.Strategy.reduce вызов , а затем через партий путем суммирования значения , возвращаемого из tf.distribute.Strategy.reduce вызовов.
  • tf.keras.Metrics должны быть обновлены внутри train_step и test_step , который получает выполняемую tf.distribute.Strategy.run . * tf.distribute.Strategy.run возвращает результаты каждой локальной реплики в стратегии, и есть несколько способов , чтобы потреблять этот результат. Вы можете сделать tf.distribute.Strategy.reduce , чтобы получить агрегированные значения. Вы также можете сделать tf.distribute.Strategy.experimental_local_results , чтобы получить список значений , содержащиеся в результате, по одному на локальную реплику.

Восстановите последнюю контрольную точку и протестируйте

Модель контрольной точки с tf.distribute.Strategy может быть восстановлена или без стратегии.

eval_accuracy = tf.keras.metrics.SparseCategoricalAccuracy(
      name='eval_accuracy')

new_model = create_model()
new_optimizer = tf.keras.optimizers.Adam()

test_dataset = tf.data.Dataset.from_tensor_slices((test_images, test_labels)).batch(GLOBAL_BATCH_SIZE)
@tf.function
def eval_step(images, labels):
  predictions = new_model(images, training=False)
  eval_accuracy(labels, predictions)
checkpoint = tf.train.Checkpoint(optimizer=new_optimizer, model=new_model)
checkpoint.restore(tf.train.latest_checkpoint(checkpoint_dir))

for images, labels in test_dataset:
  eval_step(images, labels)

print ('Accuracy after restoring the saved model without strategy: {}'.format(
    eval_accuracy.result()*100))
Accuracy after restoring the saved model without strategy: 91.0999984741211

Альтернативные способы перебора набора данных

Использование итераторов

Если вы хотите перебрать заданное число шагов , а не через весь набор данных , вы можете создать итератор с помощью iter вызова и вызова в явной форме next итератора. Вы можете перебирать набор данных как внутри, так и вне функции tf. Вот небольшой фрагмент, демонстрирующий итерацию набора данных за пределами функции tf. с использованием итератора.

for _ in range(EPOCHS):
  total_loss = 0.0
  num_batches = 0
  train_iter = iter(train_dist_dataset)

  for _ in range(10):
    total_loss += distributed_train_step(next(train_iter))
    num_batches += 1
  average_train_loss = total_loss / num_batches

  template = ("Epoch {}, Loss: {}, Accuracy: {}")
  print (template.format(epoch+1, average_train_loss, train_accuracy.result()*100))
  train_accuracy.reset_states()
Epoch 10, Loss: 0.14126229286193848, Accuracy: 95.0
Epoch 10, Loss: 0.1343936026096344, Accuracy: 95.0
Epoch 10, Loss: 0.12443388998508453, Accuracy: 94.84375
Epoch 10, Loss: 0.1607474684715271, Accuracy: 94.21875
Epoch 10, Loss: 0.10524413734674454, Accuracy: 96.71875
Epoch 10, Loss: 0.11492376029491425, Accuracy: 96.71875
Epoch 10, Loss: 0.16041627526283264, Accuracy: 94.21875
Epoch 10, Loss: 0.13022005558013916, Accuracy: 94.6875
Epoch 10, Loss: 0.17113295197486877, Accuracy: 93.28125
Epoch 10, Loss: 0.12315043061971664, Accuracy: 95.625

Итерация внутри tf.функции

Вы также итерация по всему входному можете train_dist_dataset внутри tf.function , используя for x in ... конструкте или путем создания итератора , как мы делали выше. Приведенный ниже пример демонстрирует оберточной одной эпохи обучения в tf.function и итерации по train_dist_dataset внутри функции.

@tf.function
def distributed_train_epoch(dataset):
  total_loss = 0.0
  num_batches = 0
  for x in dataset:
    per_replica_losses = strategy.run(train_step, args=(x,))
    total_loss += strategy.reduce(
      tf.distribute.ReduceOp.SUM, per_replica_losses, axis=None)
    num_batches += 1
  return total_loss / tf.cast(num_batches, dtype=tf.float32)

for epoch in range(EPOCHS):
  train_loss = distributed_train_epoch(train_dist_dataset)

  template = ("Epoch {}, Loss: {}, Accuracy: {}")
  print (template.format(epoch+1, train_loss, train_accuracy.result()*100))

  train_accuracy.reset_states()
Epoch 1, Loss: 0.13766956329345703, Accuracy: 94.89666748046875
Epoch 2, Loss: 0.12510614097118378, Accuracy: 95.35166931152344
Epoch 3, Loss: 0.11464647948741913, Accuracy: 95.70333099365234
Epoch 4, Loss: 0.10295023769140244, Accuracy: 96.12000274658203
Epoch 5, Loss: 0.09352775663137436, Accuracy: 96.49666595458984
Epoch 6, Loss: 0.08494547754526138, Accuracy: 96.87166595458984
Epoch 7, Loss: 0.07917638123035431, Accuracy: 97.09166717529297
Epoch 8, Loss: 0.07128290832042694, Accuracy: 97.37833404541016
Epoch 9, Loss: 0.06662175804376602, Accuracy: 97.47999572753906
Epoch 10, Loss: 0.06016768515110016, Accuracy: 97.82833099365234

Отслеживание потерь при обучении по репликам

Мы не рекомендуем использовать tf.metrics.Mean отслеживать потери обучения через различные реплики, из-за масштабирование вычисления потерь, которая осуществляется.

Например, если вы выполняете учебное задание со следующими характеристиками:

  • Две реплики
  • На каждой реплике обрабатываются два образца.
  • Результирующие значения потерь: [2, 3] и [4, 5] на каждой реплике.
  • Глобальный размер партии = 4

При масштабировании потерь вы вычисляете значение потерь на выборку для каждой реплики, добавляя значения потерь и затем деля на глобальный размер пакета. В этом случае: (2 + 3) / 4 = 1.25 и (4 + 5) / 4 = 2.25 .

Если вы используете tf.metrics.Mean для отслеживания потерь через две реплики, результат отличается. В этом примере, вы в конечном итоге с total 3,50 и count из 2, что приводит к total / count = 1,75 , когда result() а result() называется на метрике. Потери вычисляются с tf.keras.Metrics масштабируется дополнительным фактором , который равен количеству реплик в синхронизации.

Руководство и примеры

Вот несколько примеров использования стратегии распространения с настраиваемыми циклами обучения:

  1. Распределенное руководство по обучению
  2. DenseNet пример использования MirroredStrategy .
  3. БЕРТ пример обучение с использованием MirroredStrategy и TPUStrategy . Этот пример особенно полезен для понимания того, как выполнять загрузку из контрольной точки и создавать периодические контрольные точки во время распределенного обучения и т. Д.
  4. NCF пример обучен с помощью MirroredStrategy , который может быть активирован с помощью keras_use_ctl флага.
  5. NMT пример обучение с использованием MirroredStrategy .

Другие примеры , приведенные в руководстве стратегии распространения .

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

  • Попробуйте новый tf.distribute.Strategy API на ваших моделях.
  • Посетите раздел Performance в руководстве , чтобы узнать больше о других стратегиях и инструментах , которые можно использовать для оптимизации производительности ваших моделей TensorFlow.