Question

最近我在调查Tensorflow中的数据集API，并且有一个方法dataset.shard()用于分布式计算。

这是Tensorflow的文档中所述的内容：

Creates a Dataset that includes only 1/num_shards of this dataset.

d = tf.data.TFRecordDataset(FLAGS.input_file)
d = d.shard(FLAGS.num_workers, FLAGS.worker_index)
d = d.repeat(FLAGS.num_epochs)
d = d.shuffle(FLAGS.shuffle_buffer_size)
d = d.map(parser_fn, num_parallel_calls=FLAGS.num_map_threads)

据说此方法返回原始数据集的一部分。如果我有两个工人，我应该这样做：

d_0 = d.shard(FLAGS.num_workers, worker_0)
d_1 = d.shard(FLAGS.num_workers, worker_1)
......
iterator_0 = d_0.make_initializable_iterator()
iterator_1 = d_1.make_initializable_iterator()

for worker_id in workers:
    with tf.device(worker_id):
        if worker_id == 0:
            data = iterator_0.get_next()
        else:
            data = iterator_1.get_next()
        ......

因为文档没有指定如何进行后续调用，所以我在这里有点困惑。

谢谢！

Answer 1

您应首先查看tutorial on Distributed TensorFlow，以便更好地了解其工作原理。

您有多个工作人员，每个工作人员运行相同的代码但差别很小：每个工作人员都有不同的FLAGS.worker_index。

当您使用tf.data.Dataset.shard时，您将提供此工作人员索引，并且数据将在工作人员之间平均分配。

这是一个有3名工人的例子。

dataset = tf.data.Dataset.range(6)
dataset = dataset.shard(FLAGS.num_workers, FLAGS.worker_index)


iterator = dataset.make_one_shot_iterator()
res = iterator.get_next()

# Suppose you have 3 workers in total
with tf.Session() as sess:
    for i in range(2):
        print(sess.run(res))

我们将得到输出：

0, 3
1, 4
2, 5

如何在tensorflow中使用dataset.shard？

1 个答案: