GradCAM with Transfer-Learning and augmentation layers at imput

Question

I'm trying to implement GradCAM to a transfer-learning model. For that reason I need an additional output from the last Convolutional Layer of the base model. My model consists of preprocessing/augmentation layers, pretrained MobileNet and a custom head. When MobileNet is implemented one functional layer I always get a disconnected graph error. And because of augmentation layers at the beginning I didn't manage to implement MobileNet as single layers, as other solutions proposed. Thanks a lot for any help!

# transfer-learning model

base_model = MobileNetV2(input_shape=(224, 224, 3), include_top=False, weights='imagenet')
          

inputs = Input(shape=(224, 224, 3))
augmented = RandomFlip("horizontal")(inputs)
augmented = RandomRotation(0.1)(augmented)
augmented = RandomZoom(height_factor=(0.0, 0.3), width_factor=(0.0, 0.3), 
                                 fill_mode='constant')(augmented)
mobilenet = base_model(augmented)
pooling = GlobalAveragePooling2D()(mobilenet)
dropout = Dropout(0.5)(pooling)
outputs = Dense(len(classes), activation="softmax")(dropout)
model = Model(inputs=inputs, outputs=outputs)
model.summary()

And here's my model for GradCAM:

gradModel = Model(inputs=[model.inputs],
                  outputs=[model.get_layer('mobilenetv2_1.00_224').get_layer('Conv_1').output,
                  model.output])

Answer 1

I had a similar problem and ended up implementing the augmentation at the dataset level rather than in the model layers.

  train_ds = tf.keras.utils.image_dataset_from_directory(
      train_dir,
      validation_split=0.3,
      label_mode='categorical',
      subset="training",
      seed=s,
      color_mode="rgb",
      image_size=image_size,
      batch_size=batch_size,
  )

 data_augmentation = tf.keras.Sequential([
    tf.keras.layers.RandomFlip("horizontal"),
    tf.keras.layers.RandomRotation(0.1),
    tf.keras.layers.RandomZoom(height_factor=(0.0, 0.3), width_factor=(0.0, 0.3), fill_mode='constant')
  ])
  
  train_ds = train_ds.map(
      lambda x, y: (data_augmentation(x, training=True), y)
  )
 

I would then feed this data into the model and it had the desired effect.

model.fit(train_ds, EPOCHS)