Tensorflow tape.gradient is NONE when applying Grad-CAM on transfer learning model in model - keras

Question

I am using transfer learning to change my inner regression model with a outer model for classification.
The inner model is trained on the MNIST dataset where it had to predict ones from zeros. (Note that this is just an example to test the method before changing it to a more complex models and data)
I then freeze the inner model to prevent the regression output of the inner model from being changed while learning outer model.
The outer model classifies which instance the input is. So the first image of my batch should become the first class and with Tensorflow's Keras this works fine.
Then I remove the outer model's softmax activation function and then apply Grad-CAM.
Grad-CAM is applied on the final layer of the inner model (this is a 2D Conv layer).
But the grads = tape.gradient(class_channel, last_conv_layer_output) line in the code does produce only NONE.
Any help would be appreciated (I hope enough code is provided to reproduce the error ;-))
I have tried the steps from Grad-CAM Transfer learning error: Attempt to convert a value (None) with an unsupported type (<class 'NoneType'>) to a Tensor but to no avail. In addition allowing the inner model to be retrained also produces the same problem. I am aware of the import problems with Keras and TF with regards with tape.gradient.
Libaries
import numpy as np
import matplotlib.pyplot as plt
import tensorflow as tf
Preprocess dataset
(x_train, y_train), (x_test, y_test) = tf.keras.datasets.mnist.load_data()
x_train = x_train / 255
x_test = x_test / 255
# Select only ones and zeros.
x_train_zeros = x_train[y_train == 0]
x_train_ones = x_train[y_train == 1]
x_test_zeros = x_test[y_test == 0]
x_test_ones = x_test[y_test == 1]
# X and Y need to have same length
y_train_ones = x_train_ones[:len(x_train_zeros)]
y_test_ones = x_test_ones[:len(x_test_zeros)]
Model: "inner"
__________________________________________________________________________________________________
Layer (type) Output Shape Param # Connected to
==================================================================================================
input_1 (InputLayer) [(None, 28, 28, 1)] 0 []
conv2d (Conv2D) (None, 28, 28, 32) 320 ['input_1[0][0]']
add (Add) (None, 28, 28, 32) 0 ['input_1[0][0]',
'conv2d[0][0]']
conv2d_1 (Conv2D) (None, 28, 28, 32) 9248 ['add[0][0]']
add_1 (Add) (None, 28, 28, 32) 0 ['add[0][0]',
'conv2d_1[0][0]']
conv2d_2 (Conv2D) (None, 28, 28, 32) 9248 ['add_1[0][0]']
add_2 (Add) (None, 28, 28, 32) 0 ['add_1[0][0]',
'conv2d_2[0][0]']
conv2d_3 (Conv2D) (None, 28, 28, 32) 9248 ['add_2[0][0]']
add_3 (Add) (None, 28, 28, 32) 0 ['add_2[0][0]',
'conv2d_3[0][0]']
inner_final (Conv2D) (None, 28, 28, 1) 289 ['add_3[0][0]']
==================================================================================================
Total params: 28,353
Trainable params: 28,353
Non-trainable params: 0
__________________________________________________________________________________________________
history = inner_model.fit(x_train_zeros, y_train_ones, batch_size = batch_size, epochs = epochs, validation_data = (x_test_zeros, y_test_ones))
One hot encode the output for the classification
batch_size_classify = 40 # I want to "classify" the first 40 images.
X_train_class = x_train_zeros[: batch_size_classify]
# Produce "class labels"
Y_test = np.arange(0, batch_size_classify, 1, dtype=int)
from numpy import array
from numpy import argmax
from sklearn.preprocessing import LabelEncoder
from sklearn.preprocessing import OneHotEncoder
# define example
values = array(Y_test)
# integer encode
label_encoder = LabelEncoder()
integer_encoded = label_encoder.fit_transform(values)
# binary encode
onehot_encoder = OneHotEncoder(sparse=False)
integer_encoded = integer_encoded.reshape(len(integer_encoded), 1)
onehot_encoded = onehot_encoder.fit_transform(integer_encoded)
# invert first example
#inverted = label_encoder.inverse_transform([argmax(onehot_encoded[0, :])])
X_train_class = tf.convert_to_tensor(X_train_class)
onehot_encoded = tf.convert_to_tensor(onehot_encoded)
Build the outer model
# prevent changing the inner model
inner_model.trainable = False
inp = tf.keras.layers.Input(shape=(28, 28, 1))
x = inner_model(inp)
x = tf.keras.layers.Flatten()(x)
outputs = tf.keras.layers.Dense(batch_size_2, activation = 'softmax')(x)
Outer_model = tf.keras.Model(inp, outputs)
Outer_model.compile(loss="categorical_crossentropy", optimizer="adam", metrics=["accuracy"])
Model: "Outer"
__________________________________________________________________________________________________
Layer (type) Output Shape Param #
==================================================================================================
input_1 (InputLayer) [(None, 28, 28, 1)] 0
inner (Functional) (None, 28, 28, 1) 28353
|¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯|
| input_1 (InputLayer) [(None, 28, 28, 1)] 0 |
| |
| conv2d (Conv2D) (None, 28, 28, 32) 320 |
| |
| add (Add) (None, 28, 28, 32) 0 |
| |
| conv2d_1 (Conv2D) (None, 28, 28, 32) 9248 |
| |
| add_1 (Add) (None, 28, 28, 32) 0 |
| |
| conv2d_2 (Conv2D) (None, 28, 28, 32) 9248 |
| |
| add_2 (Add) (None, 28, 28, 32) 0 |
| |
| conv2d_3 (Conv2D) (None, 28, 28, 32) 9248 |
| |
| add_3 (Add) (None, 28, 28, 32) 0 |
| |
| inner_final (Conv2D) (None, 28, 28, 1) 289 |
¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯¯
flatten_2 (Flatten) (None, 784) 0
dense_2 (Dense) (None, 40) 31400
==================================================================================================
Total params: 59,753
Trainable params: 31,400
Non-trainable params: 28,353
__________________________________________________________________________________________________
epochs = 20
# Fit the model to the training data.
Outer_model.fit(X_train_class, onehot_encoded, batch_size = 1, epochs = epochs)
Grad-CAM where the grads = tape.gradient(class_channel, last_conv_layer_output) line produces NONE.
def make_gradcam_heatmap(img_array, model, last_conv_layer_name, pred_index=None):
# First, we create a model that maps the input image to the activations of the last conv layer as well as the output predictions
grad_model = tf.keras.models.Model([model.inputs], [model.get_layer(last_conv_layer_name).output, model.output])
# Then, we compute the gradient of the top predicted class for our input
# image with respect to the activations of the last conv layer
with tf.GradientTape() as tape:
last_conv_layer_output, preds = grad_model(img_array)
if pred_index is None:
pred_index = tf.argmax(preds[0])
class_channel = preds[:, pred_index]
# This is the gradient of the output neuron (top predicted or chosen)
# with regard to the output feature map of the last conv layer
# The line in Grad-CAM where the code produces NONE...
grads = tape.gradient(class_channel, last_conv_layer_output)
# This is a vector where each entry is the mean intensity
#of the gradient over a specific feature map channel
pooled_grads = tf.reduce_mean(grads, axis = (0, 1, 2))
# We multiply each channel in the feature map array
# by "how important this channel is" with regard to the top predicted class
# then sum all the channels to obtain the heatmap class activation
last_conv_layer_output = last_conv_layer_output[0]
heatmap = last_conv_layer_output # pooled_grads[..., tf.newaxis]
heatmap = tf.squeeze(heatmap)
# For visualization purpose, we will also normalize the heatmap between 0 & 1
heatmap = tf.maximum(heatmap, 0) / tf.math.reduce_max(heatmap)
return heatmap.numpy()
last_conv_layer_name = "inner"
# select first image to test Grad-CAM onto.
img_array = X_train_classs[:1]
# Remove softmax activation function.
Outer_model.layers[-1].activation = None
img_array = tf.expand_dims(img_array, axis = 3)
heatmap = make_gradcam_heatmap(img_array, Outer_model, last_conv_layer_name)
Cell output:
ValueError Traceback (most recent call last)
<ipython-input-44-4dea4cd32c74> in <module>
7 img_array = tf.expand_dims(img_array, axis = 3)
8
----> 9 heatmap = make_gradcam_heatmap(img_array, U_model, last_conv_layer_name)
10
11 # Display heatmap
2 frames
<ipython-input-19-f9b93747c0ae> in make_gradcam_heatmap(img_array, model, last_conv_layer_name, pred_index)
17
18 # This is a vector where each entry is the mean intensity of the gradient over a specific feature map channel
---> 19 pooled_grads = tf.reduce_mean(grads, axis = (0, 1, 2))
20
21 # We multiply each channel in the feature map array
/usr/local/lib/python3.8/dist-packages/tensorflow/python/util/traceback_utils.py in error_handler(*args, **kwargs)
151 except Exception as e:
152 filtered_tb = _process_traceback_frames(e.__traceback__)
--> 153 raise e.with_traceback(filtered_tb) from None
154 finally:
155 del filtered_tb
/usr/local/lib/python3.8/dist-packages/tensorflow/python/framework/constant_op.py in convert_to_eager_tensor(value, ctx, dtype)
100 dtype = dtypes.as_dtype(dtype).as_datatype_enum
101 ctx.ensure_initialized()
--> 102 return ops.EagerTensor(value, ctx.device_name, dtype)
103
104
ValueError: Attempt to convert a value (None) with an unsupported type (<class 'NoneType'>) to a Tensor.