TF 数据集迭代器只前进一次而不是两次

Question

我正在使用 Tensorflow 1.4.1 并了解 Tensorflow Dataset API。在描述consuming values from an iterator的部分，有如下例子

dataset1 = tf.data.Dataset.from_tensor_slices(tf.random_uniform([4, 10]))
dataset2 = tf.data.Dataset.from_tensor_slices((tf.random_uniform([4]), tf.random_uniform([4, 100])))
dataset3 = tf.data.Dataset.zip((dataset1, dataset2))

iterator = dataset3.make_initializable_iterator()

sess.run(iterator.initializer)
next1, (next2, next3) = iterator.get_next()

...带有以下指导性引用：

请注意，评估 next1、next2 或 next3 中的任何一个都会推进 所有组件的迭代器。迭代器的典型消费者将 在一个表达式中包含所有组件。

我决定通过以下简单示例来测试这种行为。

import tensorflow as tf

dataset1 = tf.data.Dataset.range(5)
dataset2 = tf.data.Dataset.range(5)
dataset3 = tf.data.Dataset.zip((dataset1, dataset2))

iterator = dataset3.make_initializable_iterator()

with tf.Session() as sess:
    sess.run(iterator.initializer)
    next1, next2 = iterator.get_next()

    A = next1
    B = next1 + next2

    while True:
        try:
            a, b = sess.run([A,B])
            print(a,b)
        except tf.errors.OutOfRangeError:
            print('done')
            break

如您所见，我在两个表达式中评估 next1，A 和 B。根据上面的引用，如果迭代器确实为每次评估都提前了，我期待迭代器在这两个评估中都提前并看到包含

的打印输出

(0, 2)
(2, 6)

但是，我得到的是：

(0, 0)
(1, 2)
(2, 4)
(3, 6)
(4, 8)

为什么迭代器只前进一次？什么是显示我期望看到的行为的工作示例？

Answer 1

当您在 TensorFlow 图中执行改变状态的操作（如 iterator.get_next()）时，经常会出现混淆。规则很简单：

图中的每个有状态操作（不在tf.while_loop() 或tf.cond() 中）在每次Session.run() 调用时都会执行一次。

应用该规则，您的图表中只有一个 iterator.get_next() 操作，因此迭代器在每次 Session.run() 调用时只会前进一次，并且相同的元素将用于计算 A 和 B。

要获得所需的行为，您需要创建第二个iterator.get_next() op。此外，为了获得确定性行为，我们需要确保两个 iterator.get_next() 操作之间存在控制依赖关系。例如，以下程序展示了您想要的行为：

import tensorflow as tf

dataset1 = tf.data.Dataset.range(5)
dataset2 = tf.data.Dataset.range(5)
dataset3 = tf.data.Dataset.zip((dataset1, dataset2))

iterator = dataset3.make_initializable_iterator()

with tf.Session() as sess:
    sess.run(iterator.initializer)
    next1, next2 = iterator.get_next()
    A = next1

    # Get a second element from `iterator` and add a control dependency to
    # ensure that it is consumed *after* `A` is computed.
    with tf.control_dependencies([A]):
       next3, next4 = iterator.get_next()
    B = next3 + next4

    while True:
        try:
            a, b = sess.run([A,B])
            print(a,b)
        except tf.errors.OutOfRangeError:
            print('done')
            break