Python tf.distribute.DistributedDataset.__iter__用法及代码示例

用法

__iter__()

为 tf.distribute.DistributedDataset 创建一个迭代器。

返回的迭代器实现 Python 迭代器协议。

示例用法：

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)
}