Expanded model in Keras for visualization - python-3.x

Question

I want to use a pre-trained model in Keras, use it as a feature extractor (fine-tuning only the last layer) and then visualize the layers' outputs with Grad-CAM. However, when I add input/data augmentation/ preprocessing layers before my base_model, it compresses my base_model into a nested one and then I run into a problem during the visualization phase. Is there a way I can expand the nested model while having the input/data augmentation/ preprocessing layers before it? I know I can have it expanded without these layers and I put both codes below. The thing is when having all the three added layers before my base model I always have it nested.
Here is a code for when I get a nested model:
inputs = tf.keras.Input(shape=(255, 255, 3))
x = data_augmentation(inputs)
x = preprocess_input(x)
x = base_model(x, training=False)
x = global_average_layer(x)
outputs = prediction_layer(x)
model = tf.keras.Model(inputs, outputs)
here is the result:
Model: "model"
_________________________________________________________________
Layer (type) Output Shape Param #
=================================================================
input_2 (InputLayer) [(None, 255, 255, 3)] 0
sequential (Sequential) (None, 255, 255, 3) 0
tf.__operators__.getitem (S (None, 255, 255, 3) 0
licingOpLambda)
tf.nn.bias_add (TFOpLambda) (None, 255, 255, 3) 0
vgg16 (Functional) (None, 7, 7, 512) 14714688
global_average_pooling2d (G (None, 512) 0
lobalAveragePooling2D)
dropout (Dropout) (None, 512) 0
dense (Dense) (None, 38) 19494
=================================================================
Total params: 14,734,182
Trainable params: 19,494
Non-trainable params: 14,714,688
_________________________________________________________________
and the code when I do not get the input/data augmentation/preprocess layers but get the expanded version of the model:
base_model = tf.keras.applications.vgg16.VGG16(weights = 'imagenet',include_top = False)
x = base_model.output
x = GlobalAveragePooling2D()(x)
out = Dense(38, activation='softmax')(x)
model = Model(inputs=base_model.input, outputs=out)
here is the result:
_________________________________________________________________
Layer (type) Output Shape Param #
=================================================================
input_4 (InputLayer) [(None, None, None, 3)] 0
block1_conv1 (Conv2D) (None, None, None, 64) 1792
block1_conv2 (Conv2D) (None, None, None, 64) 36928
block1_pool (MaxPooling2D) (None, None, None, 64) 0
block2_conv1 (Conv2D) (None, None, None, 128) 73856
block2_conv2 (Conv2D) (None, None, None, 128) 147584
block2_pool (MaxPooling2D) (None, None, None, 128) 0
block3_conv1 (Conv2D) (None, None, None, 256) 295168
block3_conv2 (Conv2D) (None, None, None, 256) 590080
block3_conv3 (Conv2D) (None, None, None, 256) 590080
block3_pool (MaxPooling2D) (None, None, None, 256) 0
block4_conv1 (Conv2D) (None, None, None, 512) 1180160
block4_conv2 (Conv2D) (None, None, None, 512) 2359808
block4_conv3 (Conv2D) (None, None, None, 512) 2359808
block4_pool (MaxPooling2D) (None, None, None, 512) 0
block5_conv1 (Conv2D) (None, None, None, 512) 2359808
block5_conv2 (Conv2D) (None, None, None, 512) 2359808
block5_conv3 (Conv2D) (None, None, None, 512) 2359808
block5_pool (MaxPooling2D) (None, None, None, 512) 0
global_average_pooling2d_3 (None, 512) 0
(GlobalAveragePooling2D)
dense_3 (Dense) (None, 38) 19494
=================================================================
Total params: 14,734,182
Trainable params: 14,734,182
Non-trainable params: 0
_________________________________________________________________