Unable to create functional model as slice from internal layer output if it is a Sequential model

[shkarupa-alex]

System information

• Have I written custom code (as opposed to using a stock example script provided in TensorFlow): No
• OS Platform and Distribution (e.g., Linux Ubuntu 16.04): Colab
• Mobile device (e.g. iPhone 8, Pixel 2, Samsung Galaxy) if the issue happens on mobile device: No
• TensorFlow installed from (source or binary): Colab default
• TensorFlow version (use command below): v2.4.1-0-g85c8b2a817f 2.4.1
• Python version: 3 (Colab default)
• Bazel version (if compiling from source): No
• GCC/Compiler version (if compiling from source): No
• CUDA/cuDNN version: No
• GPU model and memory: No

Describe the current behavior
I'm trying to create feature extractor for semantic segmentation model like here https://www.tensorflow.org/tutorials/images/segmentation
But my pretrained backbone built with combination of custom layers and sequential models (used as layer). Particularly it is BiT-S model from here https://github.com/google-research/big_transfer/blob/master/bit_tf2/models.py

During the feature extraction i got a Graph disconnected error every time i try to use Sequetial-layer output in new model.

Describe the expected behavior
There should be no error

Standalone code to reproduce the issue
https://colab.research.google.com/drive/1hg4RJtux1md5zJErzG7RDCQg0ye5cmFi?usp=sharing

Other info / logs Include any logs or source code that would be helpful to
diagnose the problem. If including tracebacks, please include the full
traceback. Large logs and files should be attached.

---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
<ipython-input-5-6cce9314b5b2> in <module>()
----> 1 ext3_model = tf.keras.Model(inputs=funct_model.inputs, outputs=funct_model.get_layer(name='block2').output)

4 frames
/usr/local/lib/python3.6/dist-packages/tensorflow/python/training/tracking/base.py in _method_wrapper(self, *args, **kwargs)
    515     self._self_setattr_tracking = False  # pylint: disable=protected-access
    516     try:
--> 517       result = method(self, *args, **kwargs)
    518     finally:
    519       self._self_setattr_tracking = previous_value  # pylint: disable=protected-access

/usr/local/lib/python3.6/dist-packages/tensorflow/python/keras/engine/functional.py in __init__(self, inputs, outputs, name, trainable, **kwargs)
    118     generic_utils.validate_kwargs(kwargs, {})
    119     super(Functional, self).__init__(name=name, trainable=trainable)
--> 120     self._init_graph_network(inputs, outputs)
    121 
    122   @trackable.no_automatic_dependency_tracking

/usr/local/lib/python3.6/dist-packages/tensorflow/python/training/tracking/base.py in _method_wrapper(self, *args, **kwargs)
    515     self._self_setattr_tracking = False  # pylint: disable=protected-access
    516     try:
--> 517       result = method(self, *args, **kwargs)
    518     finally:
    519       self._self_setattr_tracking = previous_value  # pylint: disable=protected-access

/usr/local/lib/python3.6/dist-packages/tensorflow/python/keras/engine/functional.py in _init_graph_network(self, inputs, outputs)
    202     # Keep track of the network's nodes and layers.
    203     nodes, nodes_by_depth, layers, _ = _map_graph_network(
--> 204         self.inputs, self.outputs)
    205     self._network_nodes = nodes
    206     self._nodes_by_depth = nodes_by_depth

/usr/local/lib/python3.6/dist-packages/tensorflow/python/keras/engine/functional.py in _map_graph_network(inputs, outputs)
    988                              'The following previous layers '
    989                              'were accessed without issue: ' +
--> 990                              str(layers_with_complete_input))
    991         for x in nest.flatten(node.outputs):
    992           computable_tensors.add(id(x))

ValueError: Graph disconnected: cannot obtain value for tensor KerasTensor(type_spec=TensorSpec(shape=(None, None, None, 10), dtype=tf.float32, name='conv2d_2_input'), name='conv2d_2_input', description="created by layer 'conv2d_2_input'") at layer "conv2d_2". The following previous layers were accessed without issue: []

[ravikyram]

Was able to reproduce the issue with TF v2.3, TF v2.4 and TF-nightly(2.5.0-dev20210204). Please find the gist of it here. Thanks!

[AdityaKane2001]

Just a side note.
ext3_model = tf.keras.Model(inputs=funct_model.inputs, outputs=funct_model.get_layer(name='proj').input)

This works as expected (gist).

Output is:

Model: "model_5"
_________________________________________________________________
Layer (type)                 Output Shape              Param #   
=================================================================
image (InputLayer)           [(None, None, None, 3)]   0         
_________________________________________________________________
block0 (Sequential)          (None, None, None, 10)    320       
_________________________________________________________________
block1 (Sequential)          (None, None, None, 10)    950       
_________________________________________________________________
block2 (Sequential)          (None, None, None, 10)    950       
=================================================================
Total params: 2,220
Trainable params: 2,160
Non-trainable params: 60
_________________________________________________________________

[google-ml-butler]

This issue has been automatically marked as stale because it has not had recent activity. It will be closed if no further activity occurs. Thank you.

[shkarupa-alex]

The issue is still here, "stalled" label please

[saikumarchalla]

@shkarupa-alex Seems the code is working as expected with the changes mentioned above by @AdityaKane2001. Please let us know if it is stil an issue?.Thanks!

[shkarupa-alex]

@saikumarchalla , changes mentioned above are not a solution but a hack. That hack can't be used always.
E.g. sequential models can be enclosed in another sequential many times.

So, i think this issue is still actual.

[google-ml-butler]

This issue has been automatically marked as stale because it has not had recent activity. It will be closed if no further activity occurs. Thank you.