如何在 TF2.0 中用 GradientTape 替换 Keras 的 gradients() function？

Question

With the "old" Keras library I created heatmaps for my CNNs using the keras.backend.gradients() function, like this:

# load model and image, then predict the class this image belongs to
model = load_model(os.path.join(model_folder, "custom_model.h5"))
image = image.load_img(image_path)
img_tensor = image.img_to_array(image)
img_tensor = np.expand_dims(img_tensor, axis=0)
img_tensor = preprocess_input(img_tensor)

preds = model.predict(img_tensor)
model_prediction = model.output[:, np.argmax(preds[0])]

# Calculate pooled grads for heatmap
conv_layer = model.get_layer("block5_conv3")  # last conv. layer
grads = K.gradients(model_prediction, conv_layer.output)[0]
pooled_grads = K.mean(grads, axis=(0, 1, 2))

# Get values of pooled grads and model conv. layer output as Numpy arrays
input_layer = model.get_layer("model_input")
iterate = K.function([input_layer], [pooled_grads, conv_layer.output[0]])
pooled_grads_value, conv_layer_output_value = iterate([img_tensor])

# Continue with heatmap generation ...

Now I switched to TF2.0 and it's built-in Keras implementation. Everything works fine, however, using that code I get the following error when calling K.gradients() :

tf.gradients is not supported when eager execution is enabled. Use tf.GradientTape instead.

I did some research and tried to understand how I can make use of GradientTape , but unfortunately I don't know much about TF nor TF2.0 - I always worked with Keras. Can you guys guide me how I can make this gradient calculation work again with my setup?

Answer 1

Here is a solution to your problem. It implies creating a model that outputs both conv_output and predictions, so we can apply the GradientTape properly.

Didn't have your model/data so I took a ResNet50 and random values.

import numpy as np
import tensorflow as tf

model = tf.keras.applications.resnet50.ResNet50()
img_tensor = np.random.random((1, 224, 224, 3))

conv_layer = model.get_layer('conv5_block3_1_conv')

heatmap_model = tf.keras.models.Model(
    [model.inputs], [model.get_layer('conv5_block3_1_conv').output, model.output]
)

with tf.GradientTape() as tape:
    conv_output, predictions = heatmap_model(img_tensor)
    loss = predictions[:, np.argmax(predictions[0])]

grads = tape.gradient(loss, conv_output)

Answer 2

You cannot use K.function as it won't work when eager execution is enabled, which is done by default. K.function from my limited understanding, can be used on a static graph, which occurs when eager execution is disabled.

tensorflow.compat.v1.disable_eager_execution()

Above block can alleviate some issues but create new ones. Better workaround is to use a guided gradients approach, where you only consider positive gradients and positive convolution outputs to obtain guided gradients.

castConvOutputs = tensorflow.cast(conv_output > 0, "float32")
castGrads = tensorflow.cast(grads > 0, "float32")
guidedGrads = castConvOutputs * castGrads * grads

# My goal was to create the class activation map for single image, so we are skipping axis=0 which is meant to have the batch_size of images at axis=0
convOutputs = conv_output[0]
guidedGrads = guidedGrads[0]

Finally I took the mean along the width and height of the image (pooled guided gradients) and summed the mathematically weighted activations across all channels (depth layer) to create the class activation map.

weights = tensorflow.reduce_mean(guidedGrads, axis=(0, 1))
cam = tensorflow.reduce_sum(tensorflow.multiply(weights, convOutputs), axis=-1)

This can later be modified with a colormap of choice with a cv2.applycolormap after normalizing and scaling to 0-255 integer range.

Credits: Adrian Rosebrock blog