如何在tensorflow中使用dataset.shard？

Question

最近，我正在研究Tensorflow中的dataset 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

您应该先看看分布式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上的0, 3
• 工作人员1的1, 4
• 工作人员2的2, 5