Home > Software design >  How to stack two models to create a new model in TensorFlow?
How to stack two models to create a new model in TensorFlow?

Time:11-09

I am preparing two models (one is a subset of another).

input_shape = (480, 800, 3)
conv_layer = Conv2D(filters=3, kernel_size=5, strides=(1, 1), padding="same")

# The outputs of the conv layer match the inputs of model_1
model_1 = tf.keras.applications.MobileNetV3Small(input_shape)
model_2 = conv_layer   model_1  # I need help with this line

I would like to create the model_2 such that it has an additional Conv2D layer. The parameters of the conv_layer ensure that the input size is the same as its output. Therefore, it should be possible to stack conv_layer and model_1 to generate model_2.

I was thinking of creating a model with only the conv_layer, like:

conv_model = Sequential([
               Input(shape=input_shape),
               Conv2D(filters=3, kernel_size=5, strides=(1, 1), padding="same")
             ])

In this case, I will have to stack two models conv_model and model_1 to create model_2.

I am not sure how to do either of the two, to achieve what I need.

Edit: Fixed the padding in conv_layer to ensure the input size are the same.

CodePudding user response:

You could simply add model_1 to Sequential. Also please note that the outputs of the conv layer don't match the inputs of model_1 - you need to use padding="same" for that:

import tensorflow as tf
import numpy as np

input_shape = (480, 800, 3)
conv_layer = tf.keras.layers.Conv2D(filters=3, kernel_size=5, strides=(1, 1), padding="same")
model_1 = tf.keras.applications.MobileNetV3Small(input_shape)
model_2 = tf.keras.Sequential([
    tf.keras.layers.InputLayer(input_shape=input_shape),
    conv_layer,
    model_1
])
# let's test that everything works on random data
model_2.predict(np.random.random((1, 480, 800, 3)))
  • Related