tf.custom_gradient

Декоратор для определения функции с пользовательским градиентом.

tf.custom_gradient(
    f
)

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

Например, рассмотрим следующую функцию, которая часто встречается при вычислении кросс-энтропии и логарифмических правдоподобий:

def log1pexp(x):
  return tf.math.log(1 + tf.exp(x))

Из-за численной нестабильности градиент этой функции, вычисленный при x=100, является NaN. Например:

x = tf.constant(100.)
y = log1pexp(x)
dy = tf.gradients(y, x) # Will be NaN when evaluated.

Выражение градиента можно аналитически упростить, чтобы обеспечить численную стабильность:

@tf.custom_gradient
def log1pexp(x):
  e = tf.exp(x)
  def grad(dy):
    return dy * (1 - 1 / (1 + e))
  return tf.math.log(1 + e), grad

С этим определением градиент в точке x=100 будет правильно вычислен как 1.0.

См. также tf.RegisterGradient, который регистрирует функцию градиента для примитивной операции TensorFlow. tf.custom_gradient, с другой стороны, позволяет точно контролировать вычисление градиента последовательности операций.

Обратите внимание, что если декорированная функция использует Variable, то окружающая область переменных должна использовать ResourceVariable.