如果条件输出改变，用@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),0).

之前到达 test1 中的 print('fail') 语句

那么解决方法是什么？

好吧，如果我们在 if 块中删除对 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

如果条件输出改变，用@tf.function修饰

Decorating with @tf.function changes if condition output

python

tensorflow

tensorflow2.0