How to remove layer from pre-trained TensorFlow model?

Question

I trained the following model with three hidden layers in TensorFlow:

inputs = tf.keras.Input(shape=(timesteps, feature_size))

l = LSTM(units=state_size,
               return_sequences=True)(inputs)
l = LSTM(units=state_size,
               return_sequences=True)(l)
l = LSTM(units=state_size,
               return_sequences=True)(l)

output = TimeDistributed(Dense(output_size, activation='softmax'))(l)

model = tf.keras.Model(inputs=inputs, outputs=output)

Now, I would like to use the model but skip the second hidden layer, ie directly pass the output from the first layer to the third layer without going through the second layer. I understand that I can get a hold of the output from the first layer by:

output = model.layers[idx].Output

But how could I feed this output to the third layer now? Many thanks for any help!

Answer 1

One approach is to use layer names to create a new model.

The example below uses specified names. You can also use the default names given by Keras.

inputs = tf.keras.Input(shape=(timesteps, feature_size))

l = LSTM(units=state_size, return_sequences=True, name="lstm1")(inputs)
l = LSTM(units=state_size, return_sequences=True, name="lstm2")(l)
l = LSTM(units=state_size, return_sequences=True, name="lstm3")(l)

output = TimeDistributed(Dense(output_size, activation='softmax'))(l)

model = tf.keras.Model(inputs=inputs, outputs=output)

# Now create the second model using specific layers from the first model
reuse_layers = ["lstm1", "lstm3"]

inputs = tf.keras.Input(shape=(timesteps, feature_size))
l = inputs
for layer_name in reuse_layers:
    l = model.get_layer(layer_name)(l)

output = TimeDistributed(Dense(output_size, activation='softmax'))(l)
new_model = Model(inputs=inputs, outputs=output)