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.experimental_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 = 2
strategy = tf.distribute.MirroredStrategy()
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 tf.Tensor(
    [[0.7]
     [0.7]], shape=(2, 1), dtype=float32)
    Loss is 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()
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: {}")
  print (template.format(epoch+1, average_train_loss))
    Epoch 1, Loss: 0.10000000894069672
    Epoch 2, Loss: 0.10000000894069672
    

Для повышения производительности вы также можете обернуть вызов 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 = 2
epochs = 1
steps_per_epoch = 1
mirrored_strategy = tf.distribute.MirroredStrategy()
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)
  tf.Tensor(
  [[4.]
   [4.]], shape=(2, 1), dtype=float32)
  

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

global_batch_size = 16
strategy = tf.distribute.MirroredStrategy()
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
(TensorSpec(shape=(None, 1), dtype=tf.float32, name=None),
TensorSpec(shape=(None, 1), dtype=tf.int32, name=None))

strategy = tf.distribute.MirroredStrategy(['/gpu:0', '/gpu:1'])

(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()
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))
tf.Tensor([1 2 3 4], shape=(4,), dtype=int32)

Приведённый выше пример соответствует случаю с одним устройством. Если у вас есть два устройства, например,

strategy = tf.distribute.MirroredStrategy(['/gpu:0', '/gpu:1'])

Тогда последняя строка выведет:

PerReplica:{
  0: tf.Tensor([1 2], shape=(2,), dtype=int32),
  1: tf.Tensor([3 4], shape=(2,), dtype=int32)
}