tf.custom_gradient

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

tf.custom_gradient(
    f=None
)

Используется в ноутбуках

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

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

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

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

with tf.GradientTape() as tape:
  tape.watch(x)
  y=log1pexp(x)
dy_dx = tape.gradient(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 = (exp(x_i) / (1 + exp(x_i))) = (1 - 1 / (1 + exp(x_i))). Градиент на входе 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, то окружающий контекст переменных должен использовать ResourceVariables.