Say we use TensorFlow r2.0 and say we want to use sub-models in Keras, so we have a model A like this:
def create_model_A(in_num_units, name):
x = tf.keras.Input(shape=(in_num_units))
y = tf.keras.layers.Dense(in_num_units, activation='relu')(x)
y = tf.keras.layers.BatchNormalization()(y)
y = tf.keras.layers.Dense(in_num_units, activation='relu')(y)
y = tf.keras.layers.BatchNormalization()(y)
y = x + y
return tf.keras.Model(x, y, name=name)
and model B that makes use of model A:
def create_model_B(in_num_units):
x = tf.keras.Input(shape=(in_num_units))
y = create_model_A(x.shape[-1], name='A_1')(x)
y = tf.keras.layers.Dense(in_num_units // 2, name='downsize_1')(y)
y = tf.keras.layers.BatchNormalization()(y)
y = create_model_A(y.shape[-1], name='A_2')(y)
y = tf.keras.layers.Dense(in_num_units // 2, name='downsize_2')(y)
y = tf.keras.layers.BatchNormalization()(y)
y = create_model_A(y.shape[-1], name='A_3')(y)
y = tf.keras.layers.Dense(in_num_units // 2, name='downsize_3')(y)
y = tf.keras.layers.BatchNormalization()(y)
y = create_model_A(y.shape[-1], name='A_4')(y)
return tf.keras.Model(x, y)
This works like a charm. We can instantiate a model B like this:
num_in_units = 500
model = create_model_B(num_in_units) # Works!
And benefit from all the advantages of a tf.keras.Model. But, the problem arises when we want to obtain the result of an intermediate layer that is of sub-model A. If the layer is part of model B, everything works:
inter_model_1 = tf.keras.Model(
model.input, model.get_layer('downsize_1').output) # Works!
But if the layer is of sub-model A, it crashes with a ValueError. This command:
inter_model_2 = tf.keras.Model(
model.input, model.get_layer('A_3').output) # Does not work!
Gives:
ValueError: Graph disconnected: cannot obtain value for tensor Tensor("input_4:0", shape=(None, 250), dtype=float32) at layer "input_4". The following previous layers were accessed without issue: []
I'm not sure I understand all the inside mechanics of keras. But what I understand when I dove in the source code, is that sub-models used this way have two input tensors object created. They can be printed like this:
print([n.input_tensors.name for n in model.get_layer('A_3').inbound_nodes])
['input_4:0', 'batch_normalization_5/Identity:0']
One is the sub-model tf.keras.Input and the other one is the input linked to the top-model.
When building a new model from a top-model B input tensor to a top-model B output tensor, the path in the graph seems to correctly pass by the 'batch_normalization_5' input and both tensors are correctly connected in the graph.
However, when trying to link a top-model B input tensor to a sub-model A output tensor, the output tensors seems to be connected to the sub-model tf.keras.Input and both tensors are disconnected.
A solution I found at the moment is to use the top-model version of the tensor model.get_layer('A_3').output:
model.get_layer('A_3')._outbound_nodes[0].input_tensors
But this seems overcomplicated and not clean... In addition, it does not let us make use of the layers inside model A.
I wonder if someone could give me some precision on this particular tf.keras behavior. Am I right to do so? Is this the intended behavior? Is this a bug? Thanks a lot!