tf.distribute.DistributedIterator

Итератор по tf.distribute.DistributedDataset.

tf.distribute.DistributedIterator — основной механизм для перечисления элементов tf.distribute.DistributedDataset. Он поддерживает протокол Python Iterator, что означает, что по нему можно итерироваться с помощью цикла for или извлекая отдельные элементы явно через get_next().

Вы можете создать tf.distribute.DistributedIterator, вызвав iter на tf.distribute.DistributedDataset или создав цикл Python над tf.distribute.DistributedDataset.

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

global_batch_size = 16
strategy = tf.distribute.MirroredStrategy()
dataset = tf.data.Dataset.from_tensors(([1.],[2])).repeat(100).batch(global_batch_size)
distributed_iterator = iter(strategy.experimental_distribute_dataset(dataset))
distributed_iterator.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)))

Методы

get_next

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

get_next()

Возвращает следующий элемент из итератора для всех реплик.

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

strategy = tf.distribute.MirroredStrategy()
dataset = tf.data.Dataset.range(100).batch(2)
dist_dataset = strategy.experimental_distribute_dataset(dataset)
dist_dataset_iterator = iter(dist_dataset)
@tf.function
def one_step(input):
  return input
step_num = 5
for _ in range(step_num):
  strategy.run(one_step, args=(dist_dataset_iterator.get_next(),))
strategy.experimental_local_results(dist_dataset_iterator.get_next())
(<tf.Tensor: shape=(2,), dtype=int64, numpy=array([10, 11])>,)

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

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

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

(<tf.Tensor: shape=(1,), dtype=int64, numpy=array([10])>,
 <tf.Tensor: shape=(1,), dtype=int64, numpy=array([11])>)

get_next_as_optional

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

get_next_as_optional()

Возвращает tf.experimental.Optional, содержащий следующее значение для всех реплик.

Если tf.distribute.DistributedIterator достиг конца последовательности, возвращённый tf.experimental.Optional не будет содержать значения.

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

strategy = tf.distribute.MirroredStrategy()
global_batch_size = 2
steps_per_loop = 2
dataset = tf.data.Dataset.range(10).batch(global_batch_size)
distributed_iterator = iter(
    strategy.experimental_distribute_dataset(dataset))
def step_fn(x):
  return x
@tf.function
def train_fn(distributed_iterator):
  for _ in tf.range(steps_per_loop):
    optional_data = distributed_iterator.get_next_as_optional()
    if not optional_data.has_value():
      break
    tf.print(strategy.run(step_fn, args=(optional_data.get_value(),)))
train_fn(distributed_iterator)
# ([0 1],)
# ([2 3],)

__iter__

__iter__()