GradientTape дает разные градиенты в зависимости от функции потерь, декорированной tf.function или нет

Question

Я обнаружил, что вычисленные градиенты зависят от взаимодействия декораторов tf.function следующим образом.

Сначала я создаю некоторые синтетические c данные для двоичной классификации

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))

Затем я создаю очень простую модель регрессии logisti c с удалением всех наворотов для простоты

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)>])

, где правильный только второй ответ. Таким образом, когда моя внешняя функция оформлена, я получаю правильный ответ только от внутренней функции, которая также оформлена. У меня сложилось впечатление, что достаточно украсить внешний (который является тренировочным l oop во многих приложениях), но здесь мы видим его нет. Я хочу понять, почему, а также насколько глубоко нужно go, чтобы украсить используемые функции?

Добавлена ​​некоторая отладочная информация

Я добавил некоторую отладочную информацию и показываю код только для 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]]

, мы видим, что все идентично, без тензоров, которые стремятся, и все же я получаю разные градиенты!

Answer 1

Оказывается, это ошибка, и я ее поднял здесь .

Answer 2

Это несколько сложный вопрос, но у него есть объяснение. Проблема заключается в функции tf.keras.backend.categorical_crossentropy, которая имеет разное поведение в зависимости от того, работаете ли вы в режиме ожидания или в режиме графика (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, что является наиболее распространенным в Керасе , поскольку выходной уровень для категориальной классификации обычно является softmax, тогда рассматриваются две возможности. Во-первых, если заданное выходное значение поступает из операции softmax, то оно может просто использовать входные данные для этой операции и вызывать tf.nn.softmax_cross_entropy_with_logits, что предпочтительнее для вычисления фактической перекрестной энтропии с softmax значения, потому что это предотвращает "насыщенные" результаты. Однако это можно сделать только в графическом режиме, потому что тензоры активного режима не отслеживают операцию, которую он произвел, и не обращают внимания на входные данные для этой операции.

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)

Проблема в том, что, хотя это математически эквивалентные способы вычисления перекрестной энтропии, они не имеют такой же точности. Они почти одинаковы, когда логиты маленькие, но если они становятся большими, они могут сильно расходиться. Вот простой тест:

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]

Результаты здесь почти идентичны, но вы можете видеть небольшую разницу во втором значении. Это, в свою очередь, оказывает влияние (вероятно, в усилении) на вычисляемые градиенты, которые, конечно, также являются «эквивалентными» математическими выражениями, но с другими свойствами точности.