Modify Tensorflow model at runtime - tensorflow

Question

I am trying to modify data flow of a tensorflow model at runtime. Consider a 3 layers FC neural network. Let's say I want to define 2 different layers for the middle position.
Let's say,
1st option: 64 neuron layer
2nd option: 128 neuran layer.
Then during predict function, I want to give an input alongside the input data like;
model.predict([x_test, decider])
Then if decider is 0, I want my model to execute 64 neuron layer as middle layer. Otherwise, I want my model to execute 128 neuron layer as middle layer.
If I choose one of them, I don't want the other option to be executed for performance reasons.
Note: I do not care for training.
Is there a way to do that? So far, I have been trying to use tf.cond() but could not make it work.

Answer 1

I think you could achieve same thing by recombine the independent models:
import tensorflow as tf
from tensorflow.keras import layers, models
import numpy as np
# First model
input_shape = (16, )
inputs_0 = layers.Input(shape=input_shape)
outputs_0 = layers.Dense(256, 'relu')(inputs_0)
fc_0 = models.Model(inputs_0, outputs_0)
# Middel model 0
inputs_1_0 = layers.Input(shape=(256, ))
outputs_1_0 = layers.Dense(64, 'relu')(inputs_1_0)
outputs_1_0 = layers.Dense(128, 'relu')(outputs_1_0)
fc_1_0 = models.Model(inputs_1_0, outputs_1_0)
# Middel model 1
inputs_1_1 = layers.Input(shape=(256, ))
outputs_1_1 = layers.Dense(128, 'relu')(inputs_1_1)
outputs_1_1 = layers.Dense(128, 'relu')(outputs_1_1)
fc_1_1 = models.Model(inputs_1_1, outputs_1_1)
# Last model
inputs_2 = layers.Input(shape=(128, ))
outputs_2 = layers.Dense(1, 'sigmoid')(inputs_2)
fc_2 = models.Model(inputs_2, outputs_2)
def custom_model(x, d):
h = fc_0(x)
if d == 1:
h = fc_1_0(h)
else:
h = fc_1_1(h)
return fc_2(h)
x = np.random.rand(1, input_shape[0])
decider = 0 # Middel model 0 or 1
y = custom_model(x, decider)