处理张量流Defun作为闭包

2023年10月30日 245次阅读

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

    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)

最佳答案 不,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)

(如果需求很高,我们可能会提供更全面的捕获支持.)

点赞