tensorflow: layer.output raises AttributeError because inbound nodes lost after call to activation function

@chahld , When defining a network in Keras, the first layer added needs to have input_shape added, please check the link for more information. Also refer this example.Thanks!

@chahld @zzj0402 I solved this issue with a simple trick

import tensorflow as tf

class MyModel(tf.keras.Model):
    def __init__(self, **kwargs):
        super().__init__(**kwargs)

        self.dense0 = tf.keras.layers.Dense(10, name='dense0')
        self.dense1 = tf.keras.layers.Dense(10, name='dense1')
        self.dense2 = tf.keras.layers.Dense(10, name='dense2')

    def build(self, input_shape):
        self.dense2(self.dense1(self.dense0(tf.keras.layers.Input(shape=input_shape[1:], name="input_x"))))

    def call(self, x):
        x = self.dense0(x)
        # if you use this line it works
        x = tf.keras.layers.ReLU()(x)
        x = self.dense1(x)

        print('correct:', self.dense1.inbound_nodes)

        # if you use this line it doesn't work
        relu = tf.keras.activations.get('relu')
        x = relu(x)
        x = self.dense2(x)

        print('incorrect:', self.dense2.inbound_nodes)
        return x


def main():
    my_model = MyModel()

    inp = tf.keras.Input(shape=(5, 5, 1))
    out = my_model(inp)
    my_model.summary()

    my_model.compile(optimizer='adam',
                     loss='sparse_categorical_crossentropy')

    for l in my_model.layers:
        try:
            print(l.output)
        except AttributeError:
            print('EXCEPTION: {}.output raises attribute error'.format(l.name))


if __name__=='__main__':
    main()

and this is the output:

correct: [<tensorflow.python.keras.engine.node.Node object at 0x7fc4e55bfb90>]
incorrect: [<tensorflow.python.keras.engine.node.Node object at 0x7fc4e4ea27d0>]
Model: "my_model_11"
_________________________________________________________________
Layer (type)                 Output Shape              Param #   
=================================================================
dense0 (Dense)               (None, 5, 5, 10)          20        
_________________________________________________________________
dense1 (Dense)               (None, 5, 5, 10)          110       
_________________________________________________________________
dense2 (Dense)               (None, 5, 5, 10)          110       
=================================================================
Total params: 240
Trainable params: 240
Non-trainable params: 0
_________________________________________________________________
KerasTensor(type_spec=TensorSpec(shape=(None, 5, 5, 10), dtype=tf.float32, name=None), name='dense0/BiasAdd:0', description="created by layer 'dense0'")
KerasTensor(type_spec=TensorSpec(shape=(None, 5, 5, 10), dtype=tf.float32, name=None), name='dense1/BiasAdd:0', description="created by layer 'dense1'")
KerasTensor(type_spec=TensorSpec(shape=(None, 5, 5, 10), dtype=tf.float32, name=None), name='dense2/BiasAdd:0', description="created by layer 'dense2'")

Basically, I did add a build method to your model and explicitly computed the output of each layer. You can now see the shapes of each layer in the summary. note that if you don’t pass tf.keras.Input(shape=(5, 5, 1)) to the model and pass tensor directly for example tf.random.uniform([5, 5, 1]) you need to get raide of the range selection in input_shape as follow:

    def build(self, input_shape):
        self.dense2(self.dense1(self.dense0(tf.keras.layers.Input(shape=input_shape, name="input_x"))))

    my_model(tf.random.uniform([5, 5, 1])

I hope that helps good luck This is the reproduced version in TF v2.5 but it should work on all TF v2.x versions, please find the gist here thanks

@oanush,

I am not using the Sequential Model. I discovered the problem while subclassing a layer where I needed to call tf.stop_gradient. [Note: Subclassing a layer does not require the input_shape parameter. In my example I’m calling the model on a tf.keras.Input which provides the appropriate input shape. Just to prove this to myself I’ve ammended the code in my original post, it still fails.]

TL;DNR: The inbound nodes problem happens whenever you are subclassing (either Model or Layer) and you use functions instead of classes for some of the steps in the call function. You can work around this problem by wrapping any tensorflow function inside a tf.keras.layers.Lambda.

Detailed comments:

• maybe tf.keras.activations.get should return the ReLU layer, not the relu function.

• so long as you use only Keras Layers in the model, the inbound nodes are updated correctly.

• If you use the functional interface outside of a keras.Layer, the error does not happen. So there is something about the context inside of Layer.call method that is not working properly. (see code below)

  the fact that the node structure is updated correctly in the functional interface makes me think that the layer.call problem is actually a bug.

• Workaround: My original problem happened because I needed to call tf.stop_gradient. There is no corresponding keras layer for this, but you can get around this by wrapping it inside a Lambda layer:

   StopGradient = tf.keras.layers.Lambda(tf.stop_gradient)

• or by defining a subclassed layer:

class StopGradient(KL.Layer):
    def call(self, x):
        x = tf.stop_gradient(x)
        return x

• I have no idea why this fixes the problem, because inside this layer’s call function the nodes are probably not being correctly assigned, but somehow in the layer that calls this function, it recovers. Maybe if the function call is the last or only step in a layer then the node construction works properly? I don’t know.
• The Sequential object is helpful because it throws an error if you try to add function to the sequence instead of a class:

    my_model.add(tf.keras.activations.get('relu')) 
    
    TypeError: The added layer must be an instance of class Layer. Found: <function relu at 0x7fee0be176a8>

• it would nice if the Layer subclassing system raised such a useful exception instead of silently building a bad node structure.

---

Addendum:

See my original post for code to reproduce the problem.

Here is an example using the functional interface that works:

  import tensorflow as tf

def main():
    inputs = tf.keras.Input(shape=(5, 5, 1))
    x = inputs
    x = tf.keras.layers.Dense(10, name='dense0')(x)
    x = tf.keras.layers.ReLU()(x)
    x = tf.stop_gradient(x)
    x = tf.keras.layers.Dense(10, name='dense1')(x)
    x = tf.keras.activations.relu(x)
    x = tf.keras.layers.Dense(10, name='dense2')(x)
    outputs = x
    my_model = tf.keras.Model(inputs, outputs)

    # inp = tf.keras.Input(shape=(5, 5, 1))

    my_model.compile(optimizer='adam',
                     loss='sparse_categorical_crossentropy')


    for l in my_model.layers:

        print(l.output)  # does not fail


if __name__=='__main__':
    main()