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.

Переменная dy определяется как градиент из предыдущего уровня. То есть градиент из всех слоёв или функций, исходящих из данного слоя.

По правилу цепи мы знаем, что dy/dx = dy/x_0 * dx_0/dx_1 * ... * dx_i/dx_i+1 * ... * dx_n/dx

В этом случае градиент нашей текущей функции определяется как dx_i/dx_i+1 = (1 - 1 / (1 + e)). Градиент из предыдущего уровня dy будет dx_i+1/dx_i+2 * dx_i+2/dx_i+3 * ... * dx_n/dx. Затем произведение градиента из предыдущего уровня и текущего градиента передаётся на следующий уровень.

В случае, если функция принимает несколько переменных в качестве входных данных, функция grad также должна возвращать то же количество переменных. Рассмотрим пример функции z = x * y.

@tf.custom_gradient
def bar(x, y):
  def grad(upstream):
    dz_dx = y
    dz_dy = x
    return upstream * dz_dx, upstream * dz_dy
  z = x * y
  return z, grad
x = tf.constant(2.0, dtype=tf.float32)
y = tf.constant(3.0, dtype=tf.float32)
with tf.GradientTape(persistent=True) as tape:
  tape.watch(x)
  tape.watch(y)
  z = bar(x, y)
z
<tf.Tensor: shape=(), dtype=float32, numpy=6.0>
tape.gradient(z, x)
<tf.Tensor: shape=(), dtype=float32, numpy=3.0>
tape.gradient(z, y)
<tf.Tensor: shape=(), dtype=float32, numpy=2.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ы.