[Python] GradientTape won't calculate gradients after restoring a model from ModelCheckpoint...

I'm training a CNN on Tensorflow for binary classification and executing my code in Google Colab.

CNN_model = tf.keras.Sequential([
tf.keras.layers.Input(shape=(IMAGE_SIZE, IMAGE_SIZE, 3)),
tf.keras.layers.Conv2D(filters=32, kernel_size=3, activation="relu", name='conv1'),
tf.keras.layers.Conv2D(filters=64, kernel_size=3, activation="relu", name='conv2'),
tf.keras.layers.MaxPooling2D(pool_size=2, name='maxpool1'),
tf.keras.layers.Conv2D(filters=128, kernel_size=3, activation="relu", name='conv3'),

tf.keras.layers.Flatten(),
tf.keras.layers.Dense(128, activation='relu', name='dense1'),
tf.keras.layers.Dense(1, activation='sigmoid', name='prediction')

], name='CNN_model')


Since it takes a while for the model to train, I figured I should save it (through the ModelCheckpoint callback) and restore it when needed.

callbacks = [
tf.keras.callbacks.ModelCheckpoint("best_CNN.keras",
save_best_only=True),
]


CNN_model = tf.keras.models.load_model("best_CNN.keras")


It went all right until I implemented Grad-CAM (for decision transparency). I'll paste the first section of make_gradcam_heatmap:

def make_gradcam_heatmap(img_array, model, last_conv_layer_name):

grad_model = tf.keras.models.Model(
model.inputs,
[model.get_layer(last_conv_layer_name).output, model.layers[-1].input]
)


with tf.GradientTape() as tape:
tape.watch(img_array)

last_conv_layer_output, last_layer_input = grad_model(img_array)

logits = tf.matmul(last_layer_input, model.layers[-1].kernel) + model.layers[-1].bias


grads = tape.gradient(logits, last_conv_layer_output)


When training the model and creating heatmaps in the same session, there's no problem, but when I pass make_gradcam_heatmap a restored model, tape.gradient(logits, last_conv_layer_output) will return None.

I thought it could be because by restoring a model it would turn to non-trainable, but CNN_model.trainable is True.

I tried turning each layer's trainability to True, with no result

for layer in CNN_model.layers:
layer.trainable = True


In case it might help, I'm using dataset 'malaria' from TensorFlow datasets

Continue reading...