tf.distribute.DistributedDataset

Представляет набор данных, распределённый между устройствами и машинами.

tf.distribute.DistributedDataset можно рассматривать как «распределённый» набор данных. При использовании API tf.distribute для масштабирования обучения на несколько устройств или машин необходимо также распределить входные данные, в результате чего создаётся экземпляр tf.distribute.DistributedDataset вместо экземпляра tf.data.Dataset в нераспределённом случае. В TF 2.x объекты tf.distribute.DistributedDataset являются итерируемыми объектами Python.

Примечание: Экземпляры tf.distribute.DistributedDataset не относятся к типу tf.data.Dataset. Поддерживаются только два описанных ниже варианта использования: итерация и element_spec. Другие API для преобразования или анализа набора данных не поддерживаются.

Создать объект tf.distribute.DistributedDataset можно с помощью двух API: tf.distribute.Strategy.experimental_distribute_dataset(dataset)и tf.distribute.Strategy.distribute_datasets_from_function(dataset_fn). Какой из них использовать? Если у вас есть экземпляр tf.data.Dataset и вам подходят стандартное разбиение пакетов (то есть повторное разбиение входного экземпляра tf.data.Dataset с новым размером пакета, равным глобальному размеру пакета, делённому на количество синхронно работающих реплик) и автоматическое шардингование (то есть параметры tf.data.experimental.AutoShardPolicy), используйте первый API. В противном случае, если вы не используете канонический экземпляр tf.data.Dataset или хотите настроить разбиение пакетов либо шардингование, можно поместить соответствующую логику в dataset_fn и использовать второй API. Оба API автоматически выполняют предварительную загрузку данных на устройство. Дополнительные сведения и примеры см. по ссылкам на страницы API.

Объект DistributedDataset можно использовать двумя основными способами:

1. Перебирать его для получения входных данных для одного или нескольких устройств — экземпляра tf.distribute.DistributedValues. Для этого можно:

   • использовать цикл for в стиле Python:

   global_batch_size = 4
   strategy = tf.distribute.MirroredStrategy(["GPU:0", "GPU:1"])
   dataset = tf.data.Dataset.from_tensors(([1.],[1.])).repeat(4).batch(global_batch_size)
   dist_dataset = strategy.experimental_distribute_dataset(dataset)
   @tf.function
   def train_step(input):
     features, labels = input
     return labels - 0.3 * features
   for x in dist_dataset:
     # train_step trains the model using the dataset elements
     loss = strategy.run(train_step, args=(x,))
     print("Loss is", loss)
       Loss is PerReplica:{
         0: tf.Tensor(
       [[0.7]
        [0.7]], shape=(2, 1), dtype=float32),
         1: tf.Tensor(
       [[0.7]
        [0.7]], shape=(2, 1), dtype=float32)
       }
       

   Размещение цикла внутри tf.function повышает производительность. Однако break и return в настоящее время не поддерживаются, если цикл размещён внутри tf.function. Размещение цикла внутри tf.function также не поддерживается при использовании tf.distribute.experimental.MultiWorkerMirroredStrategy или tf.distribute.experimental.TPUStrategy с несколькими рабочими узлами.

   • использовать __iter__ для создания явного итератора типа tf.distribute.DistributedIterator

   global_batch_size = 4
   strategy = tf.distribute.MirroredStrategy(["GPU:0", "GPU:1"])
   train_dataset = tf.data.Dataset.from_tensors(([1.],[1.])).repeat(50).batch(global_batch_size)
   train_dist_dataset = strategy.experimental_distribute_dataset(train_dataset)
   @tf.function
   def distributed_train_step(dataset_inputs):
     def train_step(input):
       loss = tf.constant(0.1)
       return loss
     per_replica_losses = strategy.run(train_step, args=(dataset_inputs,))
     return strategy.reduce(tf.distribute.ReduceOp.SUM, per_replica_losses,axis=None)
   EPOCHS = 2
   STEPS = 3
   for epoch in range(EPOCHS):
     total_loss = 0.0
     num_batches = 0
     dist_dataset_iterator = iter(train_dist_dataset)
     for _ in range(STEPS):
       total_loss += distributed_train_step(next(dist_dataset_iterator))
       num_batches += 1
     average_train_loss = total_loss / num_batches
     template = ("Epoch {}, Loss: {:.4f}")
     print (template.format(epoch+1, average_train_loss))
       Epoch 1, Loss: 0.2000
       Epoch 2, Loss: 0.2000
       

   Для повышения производительности можно также обернуть вызов strategy.run в tf.range внутри tf.function. Это позволяет выполнять несколько шагов в tf.function. Autograph преобразует его в tf.while_loop на рабочем узле. Однако этот способ менее гибок по сравнению с выполнением одного шага внутри tf.function. Например, нельзя выполнять операции в eager-режиме или произвольный код Python внутри шагов.

2. Проверять tf.TypeSpec данных, созданных с помощью DistributedDataset.

   tf.distribute.DistributedDataset создаёт tf.distribute.DistributedValues в качестве входных данных для устройств. Если вы передаёте входные данные в tf.function и хотите указать форму и тип каждого аргумента Tensor функции, можно передать объект tf.TypeSpec в аргумент input_signature функции tf.function. Чтобы получить tf.TypeSpec входных данных, используйте свойство element_spec объекта tf.distribute.DistributedDataset или tf.distribute.DistributedIterator.

   Например:

   global_batch_size = 4
   epochs = 1
   steps_per_epoch = 1
   mirrored_strategy = tf.distribute.MirroredStrategy(["GPU:0", "GPU:1"])
   dataset = tf.data.Dataset.from_tensors(([2.])).repeat(100).batch(global_batch_size)
   dist_dataset = mirrored_strategy.experimental_distribute_dataset(dataset)
   @tf.function(input_signature=[dist_dataset.element_spec])
   def train_step(per_replica_inputs):
     def step_fn(inputs):
       return tf.square(inputs)
     return mirrored_strategy.run(step_fn, args=(per_replica_inputs,))
   for _ in range(epochs):
     iterator = iter(dist_dataset)
     for _ in range(steps_per_epoch):
       output = train_step(next(iterator))
       print(output)
     PerReplica:{
       0: tf.Tensor(
     [[4.]
      [4.]], shape=(2, 1), dtype=float32),
       1: tf.Tensor(
     [[4.]
      [4.]], shape=(2, 1), dtype=float32)
     }
     

Дополнительные примеры и замечания см. в руководстве по распределённым входным данным.

global_batch_size = 16
strategy = tf.distribute.MirroredStrategy(["GPU:0", "GPU:1"])
dataset = tf.data.Dataset.from_tensors(([1.],[2])).repeat(100).batch(global_batch_size)
dist_dataset = strategy.experimental_distribute_dataset(dataset)
dist_dataset.element_spec
(PerReplicaSpec(TensorSpec(shape=(None, 1), dtype=tf.float32, name=None),
                TensorSpec(shape=(None, 1), dtype=tf.float32, name=None)),
 PerReplicaSpec(TensorSpec(shape=(None, 1), dtype=tf.int32, name=None),
                TensorSpec(shape=(None, 1), dtype=tf.int32, name=None)))

Методы

__iter__

Посмотреть исходный код

__iter__()

Создаёт итератор для tf.distribute.DistributedDataset.

Возвращаемый итератор реализует протокол итератора Python.

Пример использования:

global_batch_size = 4
strategy = tf.distribute.MirroredStrategy(["GPU:0", "GPU:1"])
dataset = tf.data.Dataset.from_tensor_slices([1, 2, 3, 4]).repeat().batch(global_batch_size)
distributed_iterator = iter(strategy.experimental_distribute_dataset(dataset))
print(next(distributed_iterator))
PerReplica:{
  0: tf.Tensor([1 2], shape=(2,), dtype=int32),
  1: tf.Tensor([3 4], shape=(2,), dtype=int32)
}