tf.distribute.Распределенный набор данных

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

Объект 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)
}