tensorflow - cond 可以支持有副作用的 TF 操作吗？

Question

tf.cond 的(源代码)文档不清楚评估谓词时执行的函数是否会产生副作用。我做了一些测试，但得到了相互矛盾的结果。例如，下面的代码不起作用:

import tensorflow as tf
from tensorflow.python.ops import control_flow_ops

pred = tf.placeholder(tf.bool, [])
count = tf.Variable(0)
adder = count.assign_add(1)
subtractor = count.assign_sub(2)

my_op = control_flow_ops.cond(pred, lambda: adder, lambda: subtractor)

sess = tf.InteractiveSession()
tf.initialize_all_variables().run()

my_op.eval(feed_dict={pred: True})
count.eval() # returns -1

my_op.eval(feed_dict={pred: False})
count.eval() # returns -2

即无论谓词的计算结果是什么，两个函数都会运行，因此最终结果是减 1。另一方面，此代码片段确实有效，唯一的区别是我向图中添加了新操作每次调用 my_op 时:

pred = tf.placeholder(tf.bool, [])
count = tf.Variable(0)

my_op = control_flow_ops.cond(pred, lambda:count.assign_add(1), lambda:count.assign_sub(2))

sess = tf.InteractiveSession()
tf.initialize_all_variables().run()

my_op.eval(feed_dict={pred: False})
count.eval() # returns -2

my_op.eval(feed_dict={pred: True})
count.eval() # returns -1

不知道为什么每次创建新操作都有效，而其他情况则不然，但我显然不想添加节点，因为图表最终会变得太大。

Answer 1

您的第二个版本 - 其中 assign_add() 和 assign_sub() 操作在传递给 cond() 的 lambda 内创建 - 是执行此操作的正确方法。幸运的是，在调用 cond() 期间，这两个 lambda 表达式仅计算一次，因此您的图形不会无限制地增长。

cond() 的主要作用如下:

1. 创建一个 Switch 节点，该节点根据 pred 的值将其输入仅转发到两个输出之一。我们将输出称为 pred_true 和 pred_false。 (它们与 pred 具有相同的值，但这并不重要，因为它永远不会被直接评估。)

2. 构建与 if_true lambda 相对应的子图，其中所有节点都对 pred_true 具有控制依赖性。

3. 构建与 if_false lambda 相对应的子图，其中所有节点都对 pred_false 具有控制依赖性。

4. 将两个 lambda 的返回值列表压缩在一起，并为每个表达式创建一个 Merge 节点。 Merge 节点接受两个输入(预计只会生成其中一个)，并将其转发到其输出。

5. 返回作为 Merge 节点输出的张量。

这意味着您可以运行第二个版本，并且无论您运行多少步，图表都保持固定大小即可。

您的第一个版本不起作用的原因是，当捕获 Tensor 时(例如示例中的 adder 或 subtractor) ，添加了一个额外的 Switch 节点来强制执行张量的值仅转发到实际执行的分支的逻辑。这是 TensorFlow 如何在其执行模型中结合前馈数据流和控制流的产物。结果是捕获的张量(在本例中为 assign_add 和 assign_sub 的结果)将始终被评估，即使它们没有被使用，并且您将看看他们的副作用。这是我们需要更好地记录的内容，并且 as Michael says ，我们将使其在未来更加有用。