Layer wise propagation(LRP) in keras neural network

user_7090776 0
Viewed 1039

I have been following LRP implementation using pyTorch and wanted to test it out using Tensorflow and Keras. I am using the same model with weights(VGG16) in Keras and was able to successfully execute the forward pass and element wise division using

# keras-tensorflow implementation
z = incr(clasifierLayers[l](A[l])) # forward pass step(1)
s = (R[l+1]/z) # Element wise division step(2)

But i am facing trouble in recreating the backward pass. In the original code(LRP), which uses pyTorch, the backward pass is computed using

# pyTorch implementation
(z*s).sum().backward(); c = A[l].grad

and when i tried to find the replicate the backward pass using tensorflow, my gradient returns None. Here is my code trying to compute the backward pass.

def getGradients(product,layer,l):
    with tf.GradientTape() as tape:
        tape.watch(product)
        a=layers[l](A[l])
        gradient = tape.gradient(product, a) 
    return gradient

c = getGradients((z*s).numpy().sum(),layers[l],l) # backward pass step(3)

Can someone tell me whats wrong with this implementation?

Thanks in Advance

conv-neural-network keras tensorflow heatmap
1 Answers

I tried to replicate the issue with one layer and performing an LRP backward step, here is the code:

import tensorflow as tf

x = tf.ones((1,10))
layer=tf.keras.layers.Dense(10)
y=layer(x)

with tf.GradientTape() as tape:
     tape.watch(x)
     z = tf.keras.layers.Dense(10)(x)+1e-9
     s = y/z
     s = tf.reshape(s, z.shape)
     c = tape.gradient(tf.reduce_sum(z*s), x)
     y*c

This code works, in the sense that it returns the gradients to c. I did not test it with a dataset, so do not know if it works as it should. Nonetheless, I think the problem with your code is that you should have the first block:

# keras-tensorflow implementation
z = incr(clasifierLayers[l](A[l])) # forward pass step(1)
s = (R[l+1]/z) # Element wise division step(2)

inside the TapeGradient scope and ask for the gradients with respect to the A[l].

Edit: I forgot to avoid gradients being propagated through s. The gradient computation should be done as follows:

     c = tape.gradient(tf.reduce_sum(z*s.numpy()), x)
user_62144790
Related