Тензорный поток: как заменить или изменить градиент?

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

В некотором смысле это на противоположном тому, что tf.stop_gradient() делает:. Вместо добавления расчета, который игнорируется при расчете градиентов, я хочу, расчет, который используется только при расчете градиентов

Простым примером может быть то, что просто масштабирует градиенты, умножая их на константу (но не умножая прямое вычисление на константу). Другим примером может быть то, что скопирует градиенты в заданный диапазон.

Ответ 1

Сначала определите свой собственный градиент:

@tf.RegisterGradient("CustomGrad")
def _const_mul_grad(unused_op, grad):
  return 5.0 * grad

Поскольку вы не хотите, чтобы в переднем проходе ничего не происходило, переопределите градиент операции идентификации с новым градиентом:

g = tf.get_default_graph()
with g.gradient_override_map({"Identity": "CustomGrad"}):
  output = tf.identity(input, name="Identity")

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

import tensorflow as tf

@tf.RegisterGradient("CustomClipGrad")
def _clip_grad(unused_op, grad):
  return tf.clip_by_value(grad, -0.1, 0.1)

input = tf.Variable([3.0], dtype=tf.float32)

g = tf.get_default_graph()
with g.gradient_override_map({"Identity": "CustomClipGrad"}):
  output_clip = tf.identity(input, name="Identity")
grad_clip = tf.gradients(output_clip, input)

# output without gradient clipping in the backwards pass for comparison:
output = tf.identity(input)
grad = tf.gradients(output, input)

with tf.Session() as sess:
  sess.run(tf.global_variables_initializer())
  print("with clipping:", sess.run(grad_clip)[0])
  print("without clipping:", sess.run(grad)[0])

Ответ 2

используйте optimizer.compute_gradients или tf.gradient, чтобы получить исходные градиенты
затем сделайте все, что вы хотите
наконец, используйте optimizer.apply_gradients

Я нашел пример из github

Ответ 3

Предполагая, что прямое вычисление

y = f(x)

И вы хотите, чтобы он backpropagate, как

y = b(x)

Простым взломом будет:

y = b(x) + tf.stop_gradients(f(x) - b(x))

Ответ 4

Самый общий способ сделать это - использовать https://www.tensorflow.org/api_docs/python/tf/RegisterGradient

Ниже я реализовал обратное сокращение градиента, которое можно использовать с matmul, как показано здесь, или любой другой op:

import tensorflow as tf
import numpy as np

# from https://gist.github.com/harpone/3453185b41d8d985356cbe5e57d67342
def py_func(func, inp, Tout, stateful=True, name=None, grad=None):

    # Need to generate a unique name to avoid duplicates:
    rnd_name = 'PyFuncGrad' + str(np.random.randint(0, 1E+8))

    tf.RegisterGradient(rnd_name)(grad)
    g = tf.get_default_graph()
    with g.gradient_override_map({"PyFunc": rnd_name}):
        return tf.py_func(func, inp, Tout, stateful=stateful, name=name)

def clip_grad(x, clip_value, name=None):
    """"
    scales backpropagated gradient so that
    its L2 norm is no more than `clip_value`
    """
    with tf.name_scope(name, "ClipGrad", [x]) as name:
        return py_func(lambda x : x,
                        [x],
                        [tf.float32],
                        name=name,
                        grad=lambda op, g : tf.clip_by_norm(g, clip_value))[0]

Пример использования:

with tf.Session() as sess:
    x = tf.constant([[1., 2.], [3., 4.]])
    y = tf.constant([[1., 2.], [3., 4.]])

    print('without clipping')
    z = tf.matmul(x, y)
    print(tf.gradients(tf.reduce_sum(z), x)[0].eval())

    print('with clipping')
    z = tf.matmul(clip_grad(x, 1.0), clip_grad(y, 0.5))
    print(tf.gradients(tf.reduce_sum(z), x)[0].eval())

    print('with clipping between matmuls')
    z = tf.matmul(clip_grad(tf.matmul(x, y), 1.0), y)
    print(tf.gradients(tf.reduce_sum(z), x)[0].eval())

Вывод:

without clipping
[[ 3.  7.]
 [ 3.  7.]]
with clipping
[[ 0.278543   0.6499337]
 [ 0.278543   0.6499337]]
with clipping between matmuls
[[ 1.57841039  3.43536377]
 [ 1.57841039  3.43536377]]