Python tf.distribute.experimental.coordinator.ClusterCoordinator.create_per_worker_dataset用法及代码示例

用法

create_per_worker_dataset(
    dataset_fn
)

通过在工作设备上调用 dataset_fn 在工作人员上创建数据集。

这将创建由dataset_fn 在workers 上生成的给定数据集，并返回一个表示这些单独数据集集合的对象。在这样的数据集集合上调用 iter 会返回一个 tf.distribute.experimental.coordinator.PerWorkerValues ，它是迭代器的集合，其中迭代器已放置在相应的工作程序上。

不支持在迭代器的 PerWorkerValues 上调用 next。迭代器旨在作为参数传递给 tf.distribute.experimental.coordinator.ClusterCoordinator.schedule 。当调度的函数即将被一个worker执行时，该函数将接收到该worker对应的个体迭代器。当迭代器是函数的输入时，可以在调度函数内的迭代器上调用 next 方法。

目前schedule 方法假设worker 都是相同的，因此假设不同worker 上的数据集是相同的，除非它们包含dataset.shuffle 操作并且未设置随机种子时它们可能会被不同地洗牌。因此，我们还建议无限重复数据集并安排有限数量的步骤，而不是依赖数据集中的OutOfRangeError。

例子：

strategy = tf.distribute.experimental.ParameterServerStrategy(
    cluster_resolver=...)
coordinator = tf.distribute.experimental.coordinator.ClusterCoordinator(
    strategy=strategy)

@tf.function
def worker_fn(iterator):
  return next(iterator)

def per_worker_dataset_fn():
  return strategy.distribute_datasets_from_function(
      lambda x:tf.data.Dataset.from_tensor_slices([3] * 3))

per_worker_dataset = coordinator.create_per_worker_dataset(
    per_worker_dataset_fn)
per_worker_iter = iter(per_worker_dataset)
remote_value = coordinator.schedule(worker_fn, args=(per_worker_iter,))
assert remote_value.fetch() == 3