GradientTape 根据损失函数是否被 tf.function 修饰给出不同的梯度

Question

我发现计算的梯度取决于 tf.function 装饰器的相互作用，如下所示。

首先，我为二元分类创建一些合成数据

tf.random.set_seed(42)
np.random.seed(42)
x=tf.random.normal((2,1))
y=tf.constant(np.random.choice([0,1],2))

然后我定义两个仅在 tf.function 装饰器中不同的损失函数

weights=tf.constant([1.,.1])[tf.newaxis,...]

def customloss1(y_true,y_pred,sample_weight=None):
    y_true_one_hot=tf.one_hot(tf.cast(y_true,tf.uint8),2)
    y_true_scale=tf.multiply(weights,y_true_one_hot)
    return tf.reduce_mean(tf.keras.losses.categorical_crossentropy(y_true_scale,y_pred))

@tf.function
def customloss2(y_true,y_pred,sample_weight=None):
    y_true_one_hot=tf.one_hot(tf.cast(y_true,tf.uint8),2)
    y_true_scale=tf.multiply(weights,y_true_one_hot)
    return tf.reduce_mean(tf.keras.losses.categorical_crossentropy(y_true_scale,y_pred))

然后我制作了一个非常简单的逻辑回归模型，删除了所有花哨的东西以保持简单

tf.random.set_seed(42)
np.random.seed(42)
model=tf.keras.Sequential([
    tf.keras.layers.Dense(2,use_bias=False,activation='softmax',input_shape=[1,])
])

最后定义两个函数来计算上述损失函数的梯度，一个被tf.function修饰，另一个不被tf.function修饰

def get_gradients1(x,y):
    with tf.GradientTape() as tape1:
        p1=model(x)
        l1=customloss1(y,p1)
    with tf.GradientTape() as tape2:
        p2=model(x)
        l2=customloss2(y,p2)

    gradients1=tape1.gradient(l1,model.trainable_variables)
    gradients2=tape2.gradient(l2,model.trainable_variables)

    return gradients1, gradients2

@tf.function
def get_gradients2(x,y):
    with tf.GradientTape() as tape1:
        p1=model(x)
        l1=customloss1(y,p1)
    with tf.GradientTape() as tape2:
        p2=model(x)
        l2=customloss2(y,p2)

    gradients1=tape1.gradient(l1,model.trainable_variables)
    gradients2=tape2.gradient(l2,model.trainable_variables)

    return gradients1, gradients2

现在当我跑步时

get_gradients1(x,y)

我明白了

([<tf.Tensor: shape=(1, 2), dtype=float32, numpy=array([[ 0.11473544, -0.11473544]], dtype=float32)>],
 [<tf.Tensor: shape=(1, 2), dtype=float32, numpy=array([[ 0.11473544, -0.11473544]], dtype=float32)>])

并且梯度与预期相等。然而当我跑步时

get_gradients2(x,y)

我明白了

([<tf.Tensor: shape=(1, 2), dtype=float32, numpy=array([[ 0.02213785, -0.5065186 ]], dtype=float32)>],
 [<tf.Tensor: shape=(1, 2), dtype=float32, numpy=array([[ 0.11473544, -0.11473544]], dtype=float32)>])

其中只有第二个答案是正确的。因此，当我的外部函数被装饰时，我只能从也被装饰的内部函数中得到正确的答案。我的印象是装饰外部（这是许多应用程序中的训练循环）就足够了，但在这里我们看到它不是。我想了解为什么以及需要多深入才能装饰正在使用的功能？

添加了一些调试信息

我添加了一些调试信息，并且仅显示了 customloss2 的代码（另一个是相同的）

@tf.function
def customloss2(y_true,y_pred,sample_weight=None):
    y_true_one_hot=tf.one_hot(tf.cast(y_true,tf.uint8),2)
    y_true_scale=tf.multiply(weights,y_true_one_hot)
    tf.print('customloss2',type(y_true_scale),type(y_pred))
    tf.print('y_true_scale','\n',y_true_scale)
    tf.print('y_pred','\n',y_pred)
    return tf.reduce_mean(tf.keras.losses.categorical_crossentropy(y_true_scale,y_pred))

在运行 get_gradients1 时我得到

customloss1 <type 'EagerTensor'> <type 'EagerTensor'>
y_true_scale 
 [[1 0]
 [0 0.1]]
y_pred 
 [[0.510775387 0.489224613]
 [0.529191136 0.470808864]]
customloss2 <class 'tensorflow.python.framework.ops.Tensor'> <class 'tensorflow.python.framework.ops.Tensor'>
y_true_scale 
 [[1 0]
 [0 0.1]]
y_pred 
 [[0.510775387 0.489224613]
 [0.529191136 0.470808864]]

我们看到 customloss1 的张量是 Eager，但 customloss2 的张量是张量，但我们得到相同的梯度值。

另一方面，当我在 get_gradients2 上运行它时

customloss1 <class 'tensorflow.python.framework.ops.Tensor'> <class 'tensorflow.python.framework.ops.Tensor'>
y_true_scale 
 [[1 0]
 [0 0.1]]
y_pred 
 [[0.510775387 0.489224613]
 [0.529191136 0.470808864]]
customloss2 <class 'tensorflow.python.framework.ops.Tensor'> <class 'tensorflow.python.framework.ops.Tensor'>
y_true_scale 
 [[1 0]
 [0 0.1]]
y_pred 
 [[0.510775387 0.489224613]
 [0.529191136 0.470808864]]

我们看到一切都是相同的，没有张量是 Eager，但我得到了不同的梯度！

Answer 1

事实证明这是一个错误，我已经在这里提出了。