python - 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 修饰，另一个不被它修饰

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 的张量是 Tensor，但我们得到了相同的梯度值。

另一方面，当我在 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

这是一个有点复杂的问题，但它有一个解释。问题出在函数 tf.keras.backend.categorical_crossentropy 内，它具有不同的行为，具体取决于您是在 Eager 还是图形 ( tf.function ) 模式下运行。

该函数考虑了三种可能的情况。第一个是你通过from_logits=True ，在这种情况下它只是调用 tf.nn.softmax_cross_entropy_with_logits :

if from_logits:
  return nn.softmax_cross_entropy_with_logits_v2(
      labels=target, logits=output, axis=axis)

如果你给 from_logits=False ，这是Keras中最常见的，由于分类分类的输出层一般是softmax，那么它考虑两种可能。第一个是，如果给定的输出值来自 softmax 操作，那么它可以只使用该操作的输入并调用 tf.nn.softmax_cross_entropy_with_logits ，它是用 softmax 值计算实际交叉熵的首选，因为它可以防止“饱和”结果。然而，这只能在图形模式下完成，因为 Eager 模式张量不会跟踪它产生的操作，更不用说该操作的输入。

if not isinstance(output, (ops.EagerTensor, variables_module.Variable)):
  output = _backtrack_identity(output)
  if output.op.type == 'Softmax':
    # When softmax activation function is used for output operation, we
    # use logits from the softmax function directly to compute loss in order
    # to prevent collapsing zero when training.
    # See b/117284466
    assert len(output.op.inputs) == 1
    output = output.op.inputs[0]
    return nn.softmax_cross_entropy_with_logits_v2(
        labels=target, logits=output, axis=axis)

最后一种情况是当您提供 from_logits=False 时并且您处于急切模式或给定的输出张量不直接来自 softmax 操作，在这种情况下，唯一的选择是从 softmax 值计算交叉熵。

# scale preds so that the class probas of each sample sum to 1
output = output / math_ops.reduce_sum(output, axis, True)
# Compute cross entropy from probabilities.
epsilon_ = _constant_to_tensor(epsilon(), output.dtype.base_dtype)
output = clip_ops.clip_by_value(output, epsilon_, 1. - epsilon_)
return -math_ops.reduce_sum(target * math_ops.log(output), axis)

问题是，尽管这些是计算交叉熵的数学等效方法，但它们的精度并不相同。当 logits 很小时，它们几乎相同，但如果它们变大，它们就会有很大差异。这是一个简单的测试:

import tensorflow as tf

@tf.function
def test_keras_xent(y, p, from_logits=False, mask_op=False):
    # p is always logits
    if not from_logits:
        # Compute softmax if not using logits
        p = tf.nn.softmax(p)
    if mask_op:
        # A dummy addition prevents Keras from detecting that
        # the value comes from a softmax operation
        p = p + tf.constant(0, p.dtype)
    return tf.keras.backend.categorical_crossentropy(y, p, from_logits=from_logits)

# Test
tf.random.set_seed(0)
y = tf.constant([1., 0., 0., 0.])

# Logits in [0, 1)
p = tf.random.uniform([4], minval=0, maxval=1)
tf.print(test_keras_xent(y, p, from_logits=True))
# 1.50469065
tf.print(test_keras_xent(y, p, from_logits=False, mask_op=False))
# 1.50469065
tf.print(test_keras_xent(y, p, from_logits=False, mask_op=True))
# 1.50469065

# Logits in [0, 10)
p = tf.random.uniform([4], minval=0, maxval=10)
tf.print(test_keras_xent(y, p, from_logits=True))
# 3.47569656
tf.print(test_keras_xent(y, p, from_logits=False, mask_op=False))
# 3.47569656
tf.print(test_keras_xent(y, p, from_logits=False, mask_op=True))
# 3.47569656

# Logits in [0, 100)
p = tf.random.uniform([4], minval=0, maxval=100)
tf.print(test_keras_xent(y, p, from_logits=True))
# 68.0106506
tf.print(test_keras_xent(y, p, from_logits=False, mask_op=False))
# 68.0106506
tf.print(test_keras_xent(y, p, from_logits=False, mask_op=True))
# 16.1180954

以你的例子为例:

import tensorflow as tf

tf.random.set_seed(42)
x = tf.random.normal((2, 1))
y = tf.constant(np.random.choice([0, 1], 2))
y1h = tf.one_hot(y, 2, dtype=x.dtype)
model = tf.keras.Sequential([
    # Linear activation because we want the logits for testing
    tf.keras.layers.Dense(2, use_bias=False, activation='linear', input_shape=[1,])
])
p = model(x)
tf.print(test_keras_xent(y1h, p, from_logits=True))
# [0.603375256 0.964639068]
tf.print(test_keras_xent(y1h, p, from_logits=False, mask_op=False))
# [0.603375256 0.964639068]
tf.print(test_keras_xent(y1h, p, from_logits=False, mask_op=True))
# [0.603375256 0.964638948]

此处的结果几乎相同，但您可以看到第二个值存在微小差异。这反过来又会对计算的梯度产生影响(可能是放大的)，这当然也是“等效的”数学表达式，但具有不同的精度属性。