我对以下示例中观察到的行为感到困惑:
import tensorflow as tf
@tf.function
def f(a):
c = a * 2
b = tf.reduce_sum(c ** 2 + 2 * c)
return b, c
def fplain(a):
c = a * 2
b = tf.reduce_sum(c ** 2 + 2 * c)
return b, c
a = tf.Variable([[0., 1.], [1., 0.]])
with tf.GradientTape() as tape:
b, c = f(a)
print('tf.function gradient: ', tape.gradient([b], [c]))
# outputs: tf.function gradient: [None]
with tf.GradientTape() as tape:
b, c = fplain(a)
print('plain gradient: ', tape.gradient([b], [c]))
# outputs: plain gradient: [<tf.Tensor: shape=(2, 2), dtype=float32, numpy=
# array([[2., 6.],
# [6., 2.]], dtype=float32)>]
较低的行为是我所期望的。我如何理解 @tf.function 案例?
非常感谢您!
(请注意,这个问题不同于:使用 tf.function 时缺少梯度,因为这里所有的计算都在函数内部。)