Tensorflow: ¿Cómo reemplazar o modificar el gradiente?

Me gustaría reemplazar o modificar el gradiente de una operación o parte de la gráfica en tensorflow. Sería ideal si puedo usar el gradiente existente en el cálculo.

De alguna manera, esto es lo opuesto a lo que hace tf.stop_gradient() : en lugar de agregar un cálculo que se ignora al calcular los gradientes, quiero un cálculo que solo se use al calcular los gradientes.

Un ejemplo simple sería algo que simplemente escala los gradientes multiplicándolos con una constante (pero no multiplica el cálculo directo por una constante). Otro ejemplo sería algo que recorte los gradientes a un rango dado.

Para tensorflow 1.7 o más nuevo, mire editar golpe.


Primero define tu gradiente personalizado:

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

Como no quiere que ocurra nada en el paso hacia adelante, anule el gradiente de una operación de identidad con su nuevo gradiente:

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

Aquí hay un ejemplo de trabajo con una capa que recorta gradientes en el paso hacia atrás y no hace nada en el paso hacia adelante, usando el mismo método:

 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]) 

Editar para TensorFlow 1.7

Desde 1.7 hay una nueva forma de redefinir el gradiente con una syntax más corta. (También permite redefinir el gradiente de múltiples operaciones al mismo tiempo, lo que no es necesario para esta pregunta). Aquí están los ejemplos de arriba, reescritos para TensorFlow 1.7:

Capa que escala gradientes en el paso hacia atrás:

 @tf.custom_gradient def scale_grad_layer(x): def grad(dy): return 5.0 * dy return tf.identity(x), grad 

Ejemplo con una capa que recorta gradientes en la pasada hacia atrás:

 import tensorflow as tf input = tf.Variable([3.0], dtype=tf.float32) @tf.custom_gradient def clip_grad_layer(x): def grad(dy): return tf.clip_by_value(dy, -0.1, 0.1) return tf.identity(x), grad output_clip = clip_grad_layer(input) 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]) 

use optimizer.compute_gradients o tf.gradient para obtener gradientes originales
entonces haz lo que quieras
por último, use optimizer.apply_gradients

Encontré un ejemplo de github.

Suponiendo que el cálculo hacia adelante es

 y = f(x) 

Y quieres que se propague como

 y = b(x) 

Un simple truco será:

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

La forma más general de hacerlo es mediante https://www.tensorflow.org/api_docs/python/tf/RegisterGradient

A continuación, implementé el recorte de gradiente matmul , que se puede usar con matmul , como se muestra aquí, o cualquier otra opción:

 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] 

Ejemplo de uso:

 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()) 

Salida:

 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]] 

Para TensorFlow r1.13 actual, use tf.custom_gradient .

La función decorada (los argumentos de entrada es una lista x ) debe devolver

  • el resultado del pase adelantado, y
  • una función que devuelve una lista de gradientes, una para cada elemento en x .

Aquí hay un ejemplo con una variable:

 @tf.custom_gradient def non_differentiable(x): f = tf.cast(x > 0, tf.float32) def grad(dy): return tf.math.maximum(0., 1 - tf.abs(x)) return f, grad 

Y uno con dos:

 @tf.custom_gradient def non_differentiable2(x0, x1): f = x0 * tf.cast(x1 > 0, tf.float32) def grad(dy): df_dx0 = tf.cast(x1 > 0, tf.float32) return dy*df_dx0, tf.zeros_like(dy) return f, grad