I want to create a combination of multiple tf.keras.Sequential models such that only one of the sub-models is evaluated at any given point in time. To better explain, I created the following model (the code for the model is at the end of this post):

In this graph, the models sequential, sequential_1, sequential_2 and sequential_3 are LSTM-based sub-models and label_0 is a simple fifth sub-model. The final layer arbiter decides based on a value in the input data (extracted from in_arb) which one of the five parallel paths is going to actually give the network output. The other four values are discarded.
Of course, the computation in the other four parallel layers (that do not contribute to the result) is wasted. So my question is: Is there any way to solve this inside TensorFlow, e.g. with some kind of conditional graph routing instead of parallel execution?
Example code for the model:
import tensorflow as tf
tf.compat.v1.enable_eager_execution()
config = tf.compat.v1.ConfigProto()
config.gpu_options.allow_growth = True
sess = tf.compat.v1.Session(config=config)
tf.compat.v1.keras.backend.set_session(sess)
batch_size = 1
def gen_lstm(base_label, num_features, num_units):
return tf.keras.Sequential([
tf.keras.layers.InputLayer(input_shape=(1, num_features), batch_size=batch_size,
name="input_{}".format(base_label)),
tf.keras.layers.LSTM(num_units,
batch_input_shape=(batch_size, 1, num_features),
return_sequences=False, stateful=True,
name="lstm_{}".format(base_label)),
tf.keras.layers.Dense(1, name="dense_{}".format(base_label)), # binary
tf.keras.layers.Activation('sigmoid', name="activ_{}".format(base_label)), # binary
])
models = {}
for l in [1, 3, 4, 5]:
global m
m = gen_lstm(l, 130, 88)
models[l] = m
in_all = tf.keras.layers.InputLayer(input_shape=(1, 132), batch_size=batch_size, name="input_all")
in_lstm = tf.keras.layers.Lambda(lambda x: tf.slice(x, [0, 0, 2], [-1, -1, -1]), name="in_lstm")(in_all.output)
def model_0(x):
return x[:, :, 1]
out_model_0 = tf.keras.layers.Lambda(model_0, name="label_0")(in_all.output)
out_concat = tf.keras.layers.Concatenate(axis=1, name="concat_infer")([out_model_0] + [m(in_lstm) for m in models.values()])
in_arb = tf.keras.layers.Lambda(lambda x: tf.reshape(tf.slice(x, [0, 0, 0], [-1, -1, 1]), (batch_size, 1)), name="in_arb")(in_all.output)
out_merged = tf.keras.layers.Concatenate(axis=1, name="concat_arb")([in_arb, out_concat])
def arbiter(x):
return tf.where(tf.equal(x[:, 0], tf.constant(0.0, dtype=tf.float32)), x[:, 1], tf.where(
tf.equal(x[:, 0], tf.constant(2.0, dtype=tf.float32)), x[:, 2] + tf.constant(2.0), tf.where(
tf.equal(x[:, 0], tf.constant(3.0, dtype=tf.float32)), x[:, 3] + tf.constant(3.0), tf.where(
tf.equal(x[:, 0], tf.constant(4.0, dtype=tf.float32)), x[:, 4] + tf.constant(4.0),
x[:, 5] + tf.constant(5.0)))))
out_merged = tf.keras.layers.Lambda(arbiter, name="arbiter")(out_merged)
lstm_model = tf.keras.Model([in_all.input], out_merged)
print(lstm_model.summary())
tf.keras.utils.plot_model(lstm_model, to_file="./temp.png", show_shapes=True)