Python tf.distribute.DistributedIterator.get_next用法及代码示例

用法

get_next()

从迭代器返回所有副本的下一个输入。

示例使用：

strategy = tf.distribute.MirroredStrategy(["GPU:0", "GPU:1"])
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=(1,), dtype=int64, numpy=array([10])>,
 <tf.Tensor:shape=(1,), dtype=int64, numpy=array([11])>)