使用 tf2.0/keras 无法获取中间子模型层的 output

Question

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!

Answer 1

just change output to get_output_at(0)

inter_model_2 = tf.keras.Model(
    model.input, model.get_layer('A_3').get_output_at(0))