“Layer is not connected” issue while accessing intermediate layer from custom callback if model is built by sub-classing

[MasterJEET]

System information

• Have I written custom code (as opposed to using a stock example script provided in TensorFlow):
  Yes
• OS Platform and Distribution (e.g., Linux Ubuntu 16.04):
  Linux Ubuntu 18.04.2 LTS
• Mobile device (e.g. iPhone 8, Pixel 2, Samsung Galaxy) if the issue happens on mobile device:
  N/A
• TensorFlow installed from (source or binary):
  Binary
• TensorFlow version (use command below):
  v2.2.0-rc4-8-g2b96f3662b 2.2.0
• Python version:
  Python 3.7.3
• Bazel version (if compiling from source):
  N/A
• GCC/Compiler version (if compiling from source):
  N/A
• CUDA/cuDNN version:
  CUDA 10.2
• GPU model and memory:
  NVIDIA TITAN X (Pascal), ~12GB

Describe the current behavior
I've a simple model and need access of intermediate layers within a custom callback to get intermediate predictions. If I build the model by sub-classing, I get the error AttributeError: Layer dense is not connected.

Describe the expected behavior
It shouldn't cause any error and be able to get predictions using intermediate layers.

Standalone code to reproduce the issue

import tensorflow as tf
import numpy as np

X = np.ones((8,16))
y = np.sum(X, axis=1)

class CustomCallback(tf.keras.callbacks.Callback):
    def on_epoch_end(self, epoch, logs=None):
        get_output = tf.keras.backend.function(
            inputs = self.model.layers[0].input,
            outputs = self.model.layers[1].output
        )
        print("\nLayer output: ", get_output(X))

class Model(tf.keras.Model):
    def build(self, input_shape):
        self.dense1 = tf.keras.layers.Dense(units=32)
        self.dense2 = tf.keras.layers.Dense(units=1)
        
    def call(self, input_tensor):
        x = self.dense1(input_tensor)
        x = self.dense2(x)
        return x

model = Model()
model.compile(optimizer='adam',loss='mean_squared_error', metrics='accuracy')
model.fit(X,y, epochs=2, callbacks=[CustomCallback()])

Other info / logs
Traceback:

AttributeError                            Traceback (most recent call last)
<ipython-input-3-dd6e118e08d6> in <module>
     11 model = Model()
     12 model.compile(optimizer='adam',loss='mean_squared_error', metrics='accuracy')
---> 13 model.fit(X,y, epochs=2, callbacks=[CustomCallback()])

/opt/anaconda3/envs/brats/lib/python3.7/site-packages/tensorflow/python/keras/engine/training.py in _method_wrapper(self, *args, **kwargs)
     64   def _method_wrapper(self, *args, **kwargs):
     65     if not self._in_multi_worker_mode():  # pylint: disable=protected-access
---> 66       return method(self, *args, **kwargs)
     67 
     68     # Running inside `run_distribute_coordinator` already.

/opt/anaconda3/envs/brats/lib/python3.7/site-packages/tensorflow/python/keras/engine/training.py in fit(self, x, y, batch_size, epochs, verbose, callbacks, validation_split, validation_data, shuffle, class_weight, sample_weight, initial_epoch, steps_per_epoch, validation_steps, validation_batch_size, validation_freq, max_queue_size, workers, use_multiprocessing)
    874           epoch_logs.update(val_logs)
    875 
--> 876         callbacks.on_epoch_end(epoch, epoch_logs)
    877         if self.stop_training:
    878           break

/opt/anaconda3/envs/brats/lib/python3.7/site-packages/tensorflow/python/keras/callbacks.py in on_epoch_end(self, epoch, logs)
    363     logs = self._process_logs(logs)
    364     for callback in self.callbacks:
--> 365       callback.on_epoch_end(epoch, logs)
    366 
    367   def on_train_batch_begin(self, batch, logs=None):

<ipython-input-2-a1f33c1e2e52> in on_epoch_end(self, epoch, logs)
      8     def on_epoch_end(self, epoch, logs=None):
      9         get_output = tf.keras.backend.function(
---> 10             inputs = self.model.layers[0].input,
     11             outputs = self.model.layers[1].output
     12         )

/opt/anaconda3/envs/brats/lib/python3.7/site-packages/tensorflow/python/keras/engine/base_layer.py in input(self)
   1806     if not self._inbound_nodes:
   1807       raise AttributeError('Layer ' + self.name +
-> 1808                            ' is not connected, no input to return.')
   1809     return self._get_node_attribute_at_index(0, 'input_tensors', 'input')
   1810 

AttributeError: Layer dense is not connected, no input to return.

If I build the model using functional API as shown below, it works fine:

initial = tf.keras.layers.Input((16,))
x = tf.keras.layers.Dense(units=32)(initial)
final = tf.keras.layers.Dense(units=1)(x)

model = tf.keras.Model(initial, final)
model.compile(optimizer='adam',loss='mean_squared_error', metrics='accuracy')
model.fit(X,y, epochs=2, callbacks=[CustomCallback()])

Here's the stackoverflow question I created on the same issue.

[ravikyram]

I have tried in colab with TF versions 2.2, 2.3-rc0,nightly version(2.4.0-dev20200701) and was able to reproduce the issue.Please, find the gist here.Thanks!

[ujjwal-ai]

Is there any update on this ? This is an extremely important use case

[cecabert]

Run into similar issue trying to access intermediate layers in #46605, apparently you can't do it with subclassed model. Which is weird since it is just accessing subpart of the computational graph (build with submodel / functional api).

[saikumarchalla]

Was able to reproduce the issue using TF 2.11.0. Please find the gist here. Thanks!

[ZGJY95]

I got same error with my code. And I fixed it with adding following code before model.fit.

model.build([None, 32,32,3])
model.call(tf.keras.Input(shape=(32,32,3)))