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_dx = tf.gradients(y, x) # Will be NaN when evaluated.

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

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

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

Переменная upstream определяется как градиент от предшествующих слоёв или функций. В приведённом примере нет предшествующих функций, поэтому upstream = dy/dy = 1.0.

Предположим, что x_i равно log1pexp в прямом проходе x_1 = x_1(x_0), x_2 = x_2(x_1), ..., x_i = x_i(x_i-1), ..., x_n = x_n(x_n-1). По правилу цепи мы знаем, что dx_n/dx_0 = dx_n/dx_n-1 * dx_n-1/dx_n-2 * ... * dx_i/dx_i-1 * ... * dx_1/dx_0.

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

Если функция принимает несколько переменных в качестве входных данных, функция 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, управляют ожидаемыми возвращаемыми значениями самой внутренней функции.

Примеры выше иллюстрируют, как указать пользовательские градиенты для функций, которые не считывают значения из переменных. Следующий пример использует переменные, которые требуют специального обращения, поскольку они фактически являются входными данными для функции прямого прохода.

weights = tf.Variable(tf.ones([2]))  # Trainable variable weights
@tf.custom_gradient
def linear_poly(x):
  # Creating polynomial
  poly = weights[1] * x + weights[0]

  def grad_fn(dpoly, variables):
    # dy/dx = weights[1] and we need to left multiply dpoly
    grad_xs = dpoly * weights[1]  # Scalar gradient

    grad_vars = []  # To store gradients of passed variables
    assert variables is not None
    assert len(variables) == 1
    assert variables[0] is weights
    # Manually computing dy/dweights
    dy_dw = dpoly * tf.stack([x ** 1, x ** 0])
    grad_vars.append(
        tf.reduce_sum(tf.reshape(dy_dw, [2, -1]), axis=1)
    )
    return grad_xs, grad_vars
  return poly, grad_fn
x = tf.constant([1., 2., 3.])
with tf.GradientTape(persistent=True) as tape:
  tape.watch(x)
  poly = linear_poly(x)
poly # poly = x + 1
<tf.Tensor: shape=(3,),
  dtype=float32,
  numpy=array([2., 3., 4.], dtype=float32)>
tape.gradient(poly, x)  # conventional scalar gradient dy/dx
<tf.Tensor: shape=(3,),
  dtype=float32,
  numpy=array([1., 1., 1.], dtype=float32)>
tape.gradient(poly, weights)
<tf.Tensor: shape=(2,), dtype=float32, numpy=array([6., 3.], dtype=float32)>

Приведённый выше пример демонстрирует использование обучаемой переменной weights. В примере внутренняя функция grad_fn принимает дополнительный аргумент variables и также возвращает дополнительный выход grad_vars. Этот дополнительный аргумент передаётся, если функция прямого прохода считывает какие-либо переменные. Вам необходимо вычислить градиент по каждой из этих variables и вывести его в виде списка grad_vars. Обратите внимание, что значение по умолчанию для variables установлено в None при отсутствии переменных в функции прямого прохода.

Следует отметить, что tf.GradientTape по-прежнему отслеживает прямой проход функции tf.custom_gradient и будет использовать отслеживаемые операции. В результате вызов tf.function в то время, когда лента ещё отслеживает, приводит к построению графа градиента. Если операция используется в tf.function без зарегистрированного градиента, будет выброшено исключение LookupError.

Пользователи могут вставить tf.stop_gradient для настройки этого поведения. Это показано в примере ниже. tf.random.shuffle не имеет зарегистрированного градиента. В результате tf.stop_gradient используется для предотвращения LookupError.

x = tf.constant([0.3, 0.5], dtype=tf.float32)

@tf.custom_gradient
def test_func_with_stop_grad(x):
  @tf.function
  def _inner_func():
    # Avoid exception during the forward pass
    return tf.stop_gradient(tf.random.shuffle(x))
    # return tf.random.shuffle(x)  # This will raise

  res = _inner_func()
  def grad(upstream):
    return upstream  # Arbitrarily defined custom gradient
  return res, grad

with tf.GradientTape() as g:
  g.watch(x)
  res = test_func_with_stop_grad(x)

g.gradient(res, x)

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

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