tf.distribute.DistributedDataset

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

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

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

Существуют два API для создания объекта tf.distribute.DistributedDataset: 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. Автоматическая генерация кода преобразует его в tf.while_loop на рабочем узле. Однако это менее гибко по сравнению с выполнением одного шага внутри tf.function. Например, вы не можете выполнять действия непосредственно или произвольный код Python в рамках шагов.

2. Просмотреть tf.TypeSpec данных, генерируемых DistributedDataset.

   tf.distribute.DistributedDataset генерирует tf.distribute.DistributedValues как вход для устройств. Если вы передаете вход в tf.function и хотите указать форму и тип каждого тензора-аргумента функции, вы можете передать объект 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 Iterator.

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

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)
}