changing the gradient map estimation - tensorflow

Question

I am trying to estimate the forward pass and the backword gradient of the function below:
def func(img-batch, X1,X2):
L=1
A1 = X1*L**2
A2 = X2*L**2
AA1 = A1*A1
AA2 = A2*A2
A11A2 = A1*A2
v = tf.nn.conv2d(img-batch, A1A2, strides=[1, 1, 1, 1], padding='SAME')
v = v+ AA1+AA2
return v
When I add this function to the network, the gradient will be performed on each instruction of the function by default.
How can I use this function and calculate it in the forward pass, in the meantime ignoring the gradient of each instruction in the function and provide other gradient estimation and add it to the main gradient of the model?

Answer 1

You can use py_func to ignore the gradients in this function, and use gradient_override_map to provide customized gradients. Here is an example:
import tensorflow as tf
def myfunc(X1, X2):
L = 1
A1 = X1 * L**2
A2 = X2 * L**2
AA1 = A1 * A1
AA2 = A2 * A2
A11A2 = A1 * A2
...
v = AA1 + AA2 + A11A2
return v
#tf.RegisterGradient("GradMyfunc")
def grad_myfunc(op, grad):
X1 = op.inputs[0]
X2 = op.inputs[1]
return [grad * X2, grad * X1]
X1 = tf.Variable(tf.constant(1.1, dtype=tf.float64))
X2 = tf.Variable(tf.constant(2.2, dtype=tf.float64))
g = tf.get_default_graph()
with g.gradient_override_map({"PyFunc": "GradMyfunc"}):
y = tf.py_func(myfunc, [X1, X2], [tf.float64])
with tf.Session() as sess:
grad = tf.gradients(y, [X1, X2])
sess.run(tf.global_variables_initializer())
print(sess.run(y))
print(sess.run(grad))