Cómo registrar un gradiente personalizado para una operación compuesta de operaciones tf

Más específicamente, tengo un fprop simple que es una composición de operaciones tf. Quiero anular el cálculo del gradiente de tensorflow con mi propio método de gradiente utilizando RegisterGradient.

¿Qué hay de malo con este código?

import tensorflow as tf from tensorflow.python.framework import ops @ops.RegisterGradient("MyopGrad") def frop_grad(op, grad): x = op.inputs[0] return 0 * x # zero out to see the difference: def fprop(x): x = tf.sqrt(x) out = tf.maximum(x, .2) return out a = tf.Variable(tf.constant([5., 4., 3., 2., 1.], dtype=tf.float32)) h = fprop(a) h = tf.identity(h, name="Myop") grad = tf.gradients(h, a) g = tf.get_default_graph() with g.gradient_override_map({'Myop': 'MyopGrad'}): with tf.Session() as sess: sess.run(tf.initialize_all_variables()) result = sess.run(grad) print(result[0]) 

Quiero ver todos los ceros en la impresión, pero en cambio estoy obteniendo:

 [ 0.2236068 0.25000003 0.28867513 0.35355341 0.5 ] 

with g.gradient_override_map({'Myop': 'MyopGrad'}) definir la with g.gradient_override_map({'Myop': 'MyopGrad'}) dentro del scope de with g.gradient_override_map({'Myop': 'MyopGrad'})

Además, debe asignar Identity lugar del nombre Myop a su nuevo gradiente.

Aquí está el código completo:

 import tensorflow as tf from tensorflow.python.framework import ops @ops.RegisterGradient("MyopGrad") def frop_grad(op, grad): x = op.inputs[0] return 0 * x # zero out to see the difference: def fprop(x): x = tf.sqrt(x) out = tf.maximum(x, .2) return out a = tf.Variable(tf.constant([5., 4., 3., 2., 1.], dtype=tf.float32)) h = fprop(a) g = tf.get_default_graph() with g.gradient_override_map({'Identity': 'MyopGrad'}): h = tf.identity(h, name="Myop") grad = tf.gradients(h, a) with tf.Session() as sess: sess.run(tf.initialize_all_variables()) result = sess.run(grad) print(result[0]) 

Salida:

 [ 0. 0. 0. 0. 0.] 

Si desea utilizar tf.RegisterGradient() para este propósito, no estoy seguro de que sea una solución adecuada. Porque en los documentos oficiales https://www.tensorflow.org/api_docs/python/tf/RegisterGradient , dice:

Este decorador solo se usa cuando se define un nuevo tipo de operación.

lo que significa que necesita definir una nueva operación escrita en C ++ o envuelta en py_func . No estoy totalmente seguro de si puede aplicarse en el grupo de “tf op” que dijiste.


Sin embargo, también puede consultar los métodos de “truco” mencionados en este hilo:

¿Cómo puedo definir solo el gradiente para un subgrafo de Tensorflow?

donde podría combinar tf.stop_gradient() y tfgradient_override_map() juntos para redefinir los gradientes para grupos de operaciones

Vea esta respuesta (tenga en cuenta que las diferentes preguntas pueden responderse satisfactoriamente con la misma respuesta).