tensorflow - 将 tensorflow Defun 视为闭包

Question

我在 tensorflow 中使用 Defun 装饰器时遇到问题。也就是说，Defun 不能关闭任何外部创建的 TF 操作。下面是一个独立的例子，展示了我想做的事情。请注意，张量 x 属于对 custom_op 的调用内外的不同图。 Defun 代码创建一个临时图，将图转换为函数原型(prototype)，然后将其合并到原始图中。代码在第一步崩溃，因为我们关闭的张量不在新的临时图中。有没有解决的办法？能够关闭事物将非常有帮助。

    import tensorflow as tf
    from tensorflow.python.framework import function

    w = tf.Variable(1.0)
    function_factory = lambda x: x*w

    @function.Defun(x=tf.float32)

    def custom_op(x):
        print('graph for x inside custom_op: ', x.graph)
        return function_factory(x)

    x = tf.constant(2.0)

    print('graph for x outside custom_op: ', x.graph)
    y = custom_op(x)

    with tf.Session() as sess:
        sess.run(tf.initialize_all_variables())
        sess.run(y)

Answer 1

不，Defun装饰器不会捕获所有内容。您需要传入w明确地，如下所示:

import tensorflow as tf
from tensorflow.python.framework import function

w = tf.Variable(1.0)

@function.Defun(tf.float32, tf.float32)
def custom_op(x, w):
    print('graph for x inside custom_op: ', x.graph)
    return x * w

x = tf.constant(2.0)
print('graph for x outside custom_op: ', x.graph)
y = custom_op(x, tf.identity(w))

with tf.Session() as sess:
    sess.run(tf.initialize_all_variables())
    sess.run(y)

(如果需求很高，我们可能会添加更完整的捕获支持。)