How to edit functional model in keras?

Question

I am using the tf.keras.applications.efficientnet_v2.EfficientNetV2L model and I want to edit the last layers of the model to make the model a regression and classification layer. However, I am unsure of how to edit this model because it is not a linear sequential model, and thus I cannot do:

for layer in model.layers[:-2]:
    model.add(layer)

as certain layers of the model have multiple inputs. Is there a way of preserving the model except the last layer so the model will diverge before the last layer?

    efficentnet[:-2]
          |
          |
         / \
        /   \
       /     \
  output1   output2

Answer 1

To enable a functional model to have a classification layer and a regression layer, you can change the model as follows. Note, there are various ways to achieve this, and this is one of them.

import tensorflow as tf
from tensorflow import keras 

prev_model = keras.applications.EfficientNetV2B0(
    input_tensor=keras.Input(shape=(224, 224, 3)),
    include_top=False
)

Next, we will write our expected head layers, shown below.

neck_branch = keras.Sequential(
    [
        # we can add more layers i.e. batch norm, etc.
        keras.layers.GlobalAveragePooling2D()
    ],
    name='neck_head'
)

classification_head = keras.Sequential(
    [
        keras.layers.Dense(10, activation='softmax')
    ],
    name='classification_head'
)

regression_head = keras.Sequential(
    [
        keras.layers.Dense(1, activation=None)
    ],
    name='regression_head'
)

Now, we can build the desired model.

x = neck_branch(prev_model.output)
output_a = classification_head(x)
output_b = regression_head(x)
final_model = keras.Model(prev_model.inputs, [output_a, output_b])

Test

keras.utils.plot_model(final_model, expand_nested=True)
# OK

final_model(tf.ones(shape=(1, 224, 224, 3)))
# OK

---

Update

Based on your comments,

how you would tackle the problem if the previous model was imported from a h5 file since there I cannot declare the top layer not to be included?

If I understand your query, you have a saved model (in .h5 format) with top layers. In that case, you don't have include_top params to exclude the top branch. So, what you can do is remove the top branch of your saved model first. Here is how,

# a saved model with top layers
prev_model = keras.models.load_model('model.h5')

prev_model_with_top_remove = keras.Model(
     prev_model.input ,
     prev_model.layers[-4].output
)
prev_model_with_top_remove.summary()

This prev_model.layers[-4].output will remove the top branch. In the end, you will give similar output as we can get with include_top=True . Check the model summary to visually inspect.

Answer 2

Keras' functional API works by linking Keras tensors (hereby called KTensor) and not your everyday TF tensors.

Therefore, the first thing you need to do is feeding KTensors (created using tf.keras.Input ) of proper shapes to the original model. This will trigger the forward chain, prompting the model's layers to produce their own output KTensors that are properly linked to the input KTensors. After the forward pass,

• The layers will store their received/produced KTensors in their input and output attributes.
• The model itself will also store the KTensors you fed to it and the corresponding final output KTensors in its inputs and outputs attributes (note the s).

Like so,

>>> from tensorflow.keras import Input
>>> from tensorflow.keras.layers import Dense
>>> from tensorflow.keras.models import Sequential, Model
>>> seq_model = Sequential()
>>> seq_model.add(Dense(1))
>>> seq_model.add(Dense(2))
>>> seq_model.add(Dense(3))
>>> hasattr(seq_model.layers[0], 'output')
False
>>> seq_model.inputs is None
True
>>> _ = seq_model(Input(shape=(10,))) # <--- Feed input KTensor to the model
>>> seq_model.layers[0].output
<KerasTensor: shape=(None, 1) dtype=float32 (created by layer 'dense')>
>>> seq_model.inputs
[<KerasTensor: shape=(None, 10) dtype=float32 (created by layer 'dense_input')>]

Once you've obtained these internal KTensors, everything becomes trivial. To extract the KTensor right before the last two layers and forward it to two different branches to form a new functional model, do

>>> intermediate_ktensor = seq_model.layers[-3].output
>>> branch_1_output = Dense(20)(intermediate_ktensor)
>>> branch_2_output = Dense(30)(intermediate_ktensor)
>>> branched_model = Model(inputs=seq_model.inputs, outputs=[branch_1_output, branch_2_output])

Note that the shapes of the KTensors you fed at the very first step must conform to the shape requirements of the layers that receive them. In my example, the input KTensor would be fed to Dense(1) layer. As Dense requires the input shape to be defined in the last dimension, the input KTensor could be of shapes, eg, (10,) or (None,10) but not (None,) or (10, None) .