Implement Guided BackProp in TensorFlow 2.0?

Question

I am trying to apply Guided Backpropagation (https://arxiv.org/abs/1412.6806) using TensorFlow 2.0. To apply guided backpropagation, we need to modify relu gradients. I read the conversation in How to apply Guided BackProp in Tensorflow 2.0? and tried to adapt the code from https://gist.github.com/falcondai/561d5eec7fed9ebf48751d124a77b087, however the results are not as I expected. I am not sure what I am missing.

Here is what I have (combining code from sources above):

import tensorflow as tf

@tf.RegisterGradient("GuidedRelu")
def _GuidedReluGrad(op, grad):
    dtype = op.inputs[0].dtype
    gate_f = tf.cast(op.outputs[0] > 0, dtype) #for f^l > 0
    gate_R = tf.cast(grad > 0, dtype) #for R^l+1 > 0
    return gate_f * gate_R * grad

with tf.compat.v1.get_default_graph().gradient_override_map({'Relu': 'GuidedRelu'}):
    with tf.GradientTape() as tape:
        x = tf.constant([10., 2.])
        tape.watch(x)
        y = tf.nn.relu(x)
        z = tf.reduce_sum(-y ** 2)
        print(x.numpy())
        print(y.numpy())
        print(z.numpy())
        print(tape.gradient(z, x).numpy())

The output is

[10.  2.]
[10.  2.]
-103.99999
[-20.  -4.]

Instead of

[10.  2.]
[10.  2.]
-103.99999
[0.  0.]

Answer 1

It does not seem to have a clean way of doing it in tf2.0/2.1. The workaround I've used is to modify my model by changing the ReLU with a custom ReLU that makes use of @custom_gradient. I've been inspired by this thread. It is a bit slow, but at least it works. TF will surely be updated to support a remap of gradient again. Hope it will help you meanwhile.

Edit: Issue discussing that issue here.