Stuck with the Implementation of Deep learning model

I intend to implement the deep learning model depicted in the image below.

enter image description here

The image depicts two models in the Training phase. The output of 1ST Model (Blue rectangle) should be the input of 2ND Model (Red rectangle) (yellow rectangle). Rest, each model will train on its own gradients relative to its own losses. Note that the gradients of the second model will not pass through to the first model.

And in Testing, I want to test samples only on the second model (yellow rectangle).

How can this be implemented in TensorFlow Keras? Any examples, suggestions, or ideas would be greatly appreciated.


You can use the functional API to define the two models and specify the inputs and outputs.

Here is an example code snippet:

from tensorflow import keras
from tensorflow.keras.layers import Input, Dense

# Define the first model
inputs1 = Input(shape=(input_shape,))
x1 = Dense(units=hidden_units, activation='relu')(inputs1)
outputs1 = Dense(units=output_units, activation='softmax')(x1)
model1 = keras.Model(inputs=inputs1, outputs=outputs1)

# Define the second model
inputs2 = Input(shape=(input_shape,))
x2 = Dense(units=hidden_units, activation='relu')(inputs2)
outputs2 = Dense(units=output_units, activation='softmax')(x2)
model2 = keras.Model(inputs=inputs2, outputs=outputs2)

# Connect the two models
x1 = model1(inputs2)
outputs = model2(x1)
3rd_model = keras.Model(inputs=inputs2, outputs=outputs)

# Compile the models
model1.compile(optimizer='adam', loss='categorical_crossentropy')
model2.compile(optimizer='adam', loss='categorical_crossentropy')
3rd_model.compile(optimizer='adam', loss='categorical_crossentropy')

# Train the models, y_train, epochs=10, validation_data=(x_val, y_val)), y_train, epochs=10, validation_data=(x_val, y_val)), y_train, epochs=10, validation_data=(x_val, y_val))

# Evaluate the second model only
model2.evaluate(x_test, y_test)

We first define the two models (model1 and model2) separately. We then connect the two models by passing the output of the first model (x1) as input to the second model, and then create a third model (3rd_model) that combines both models and takes the same input as the second model.

We compile all three models with the same optimizer and loss function, and train them on the training data. Finally, to evaluate only the second model, we use the evaluate() method on the model2 object and pass in the test data.

Note that because the gradients of the second model do not pass through to the first model during training, we can train each model separately and then combine them for testing.

Please let me know if it helps you to resolve your problem.