tf.custom_gradient

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

tf.custom_gradient(
    f=None
)

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

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

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.

Вложенность пользовательских градиентов может привести к неинтуитивным результатам. Поведение по умолчанию не соответствует производным n-го порядка. Например

@tf.custom_gradient
def op(x):
  y = op1(x)
  @tf.custom_gradient
  def grad_fn(dy):
    gdy = op2(x, y, dy)
    def grad_grad_fn(ddy):  # Not the 2nd order gradient of op w.r.t. x.
      return op3(x, y, dy, ddy)
    return gdy, grad_grad_fn
  return y, grad_fn

Функция grad_grad_fn будет вычислять градиент первого порядка от grad_fn по отношению к dy, который используется для генерации графов градиента по направлению вперёд из графов градиента по направлению назад, но не является тем же, что и градиент второго порядка от op по отношению к x.

Вместо этого оберните вложенные @tf.custom_gradients в другую функцию:

@tf.custom_gradient
def op_with_fused_backprop(x):
  y, x_grad = fused_op(x)
  def first_order_gradient(dy):
    @tf.custom_gradient
    def first_order_custom(unused_x):
      def second_order_and_transpose(ddy):
        return second_order_for_x(...), gradient_wrt_dy(...)
      return x_grad, second_order_and_transpose
    return dy * first_order_custom(x)
  return y, first_order_gradient

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

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

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