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 для преобразования или проверки набора данных.

Существуют два 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)
    }
    

Placing the loop inside a <a href="../../tf/function"><code>tf.function</code></a> will give a performance boost.
However `break` and `return` are currently not supported if the loop is
placed inside a <a href="../../tf/function"><code>tf.function</code></a>. We also don't support placing the loop
inside a <a href="../../tf/function"><code>tf.function</code></a> when using
<a href="../../tf/distribute/experimental/MultiWorkerMirroredStrategy"><code>tf.distribute.experimental.MultiWorkerMirroredStrategy</code></a> or
<a href="../../tf/distribute/experimental/TPUStrategy"><code>tf.distribute.experimental.TPUStrategy</code></a> with multiple workers.

• использовать __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. Например, вы не можете выполнять операции жадно или произвольный код Python внутри шагов.

1. Проверить 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)
}