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()ops在传递给lambda的lambd中创建cond()- 是正确的方法.幸运的是,在调用期间,两个lambdas中的每一个仅被评估一次cond(),因此您的图形将不会无限制地增长.

基本上cond()是以下内容:

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

2. 构建对应于if_truelambda 的子图,其中所有节点都具有控制依赖性pred_true.

3. 构建对应于if_falselambda 的子图,其中所有节点都具有控制依赖性pred_false.

4. 将两个lambda中的返回值列表压缩在一起,并Merge为每个lambda创建一个节点.一个Merge节点需要两个输入,其中只有一个输出,并将其转发到其输出.

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

这意味着您可以运行第二个版本,并且无论您运行多少步骤,都可以确保图形保持固定大小.

您的第一个版本不起作用的原因是,当Tensor捕获a时(例如adder或subtractor在您的示例中),Switch添加了一个额外的节点以强制执行一个逻辑,即张量的值仅转发到实际执行的分支.这是TensorFlow如何在其执行模型中组合前馈数据流和控制流的工件.结果是捕获的张量(在这种情况下是assign_add和的结果assign_sub)将始终被评估,即使它们没有被使用,你会看到它们的副作用.这是我们需要更好地记录的内容,正如迈克尔所说,我们将来会更加有用.

Answer 2

第二种情况有效,因为你在cond中添加了ops:这会导致它们有条件地执行.

第一种情况类似于:

adder = (count += 1)
subtractor = (count -= 2)
if (cond) { adder } else { subtractor }

由于加法器和减法器在条件之外,因此它们总是被执行.

第二种情况更像是说

if (cond) { adder = (count += 1) } else { subtractor = (count -= 2) }

在这种情况下,你做了你所期望的.

我们意识到副作用和(有些)懒惰评估之间的相互作用令人困惑,我们有一个中期目标,即使事情更加统一.但是现在要理解的重要一点是我们不进行真正的惰性评估:条件获取对在任一分支中使用的条件之外定义的每个数量的依赖性.