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.