python - 如果条件输出改变，用@tf.function 装饰

Question

我正在尝试评估我的变量 a 是否为空(即 size == 0)。然而，当使用 @tf.function 装饰代码时，if 语句错误地评估为 True，而在移除装饰器时它评估为 False。 tf.size(a) 在这两种情况下似乎都正确评估为 0。如何解决这个问题？

import tensorflow as tf
a=tf.Variable([[]])
@tf.function
def test(a):
    print_op = tf.print(tf.size(a))
    print(tf.size(a))
    if tf.math.not_equal(tf.size(a),0):
        print('fail')
    with tf.control_dependencies([print_op]):
        return None
test(a)

Answer 1

这有点让人头疼，但是，一旦我们了解 tf.function 正在将 python 操作和控制流映射到 tf 图，而裸函数只是急切地执行，我们就可以挑选它，它更有意义。

我已经调整了您的示例以说明发生了什么。考虑下面的 test1 和 test2:

@tf.function
def test1(a):
    print_op = tf.print(tf.size(a))
    print("python print size: {}".format(tf.size(a)))
    if tf.math.not_equal(tf.size(a),0):
        print('fail')
    with tf.control_dependencies([print_op]):
        return None

def test2(a):
    print_op = tf.print(tf.size(a))
    print("python print size: {}".format(tf.size(a)))
    if tf.math.not_equal(tf.size(a),0):
        print('fail')
    with tf.control_dependencies([print_op]):
        return None

除了 @tf.function 装饰器之外，它们彼此相同。

现在执行 test2(tf.Variable([[]])) 给我们:

0
python print size: 0

这是我假设您期望的行为。而 test1(tf.Variable([[]])) 给出:

python print size: Tensor("Size_1:0", shape=(), dtype=int32)
fail
0

关于此输出，您可能会发现一些令人惊讶的事情(除了 fail):

• print() 语句打印出一个(尚未评估的)张量而不是零
• print() 和 tf.print() 的顺序颠倒了

这是因为通过添加 @tf.function 我们不再有 python 函数，而是使用 autograph 从函数代码映射的 tf 图。这意味着，在评估 if 条件时，我们还没有执行 tf.math.not_equal(tf.size(a),0) 并且只是有一个对象(Tensor 对象的一个​​实例)，在 python 中，它是真实的:

class MyClass:
  pass
my_obj = MyClass()
if (my_obj):
  print ("my_obj evaluates to true") ## outputs "my_obj evaluates to true"

这意味着我们在计算 tf.math.not_equal(tf.size(a) 之前到达 test1 中的 print('fail') 语句,0)。

那么解决办法是什么？

好吧，如果我们在 if block 中删除对 python-only print() 函数的调用，并将其替换为亲笔签名的 tf. print() 语句然后 autograph 将无缝地将我们的 if ... else ... 逻辑转换为图形友好的 tf.cond 语句，确保一切都发生在正确顺序:

def test3(a):    print_op = tf.print(tf.size(a))    print("python print size: {}".format(tf.size(a)))    if tf.math.not_equal(tf.size(a),0):        tf.print('fail')    with tf.control_dependencies([print_op]):        return None

test3(tf.Variable([[]]))

0
python print size: 0