Loading a Subclass Model doesnt load inputs and outputs and throws AttributeError Layer name has no inbound nodes.

[AshwinJay101]

System information.

• Have I written custom code (as opposed to using a stock example script provided in Keras): No
• OS Platform and Distribution (e.g., Linux Ubuntu 16.04): Ubuntu 18.04
• TensorFlow installed from (source or binary): Source
• TensorFlow version (use command below): 2.6.0
• Python version: 3.8

Describe the problem.

If you notice when I load from a subclass model, I am not getting the shape of the output which was present in the original model

Describe the current behavior.

I am getting the following error

File "/usr/local/lib/python3.8/site-packages/keras/engine/base_layer.py", line 2096, in output
    raise AttributeError('Layer ' + self.name + ' has no inbound nodes.')
AttributeError: Layer autoencoder has no inbound nodes.

Describe the expected behavior.

The expected behaviour would be that the model would show the output as seen in the original model

KerasTensor(type_spec=TensorSpec(shape=(None, 784), dtype=tf.float32, name=None), name='autoencoder/decoder/dense_4/Sigmoid:0', description="created by layer 'autoencoder'")

Contributing.

• Do you want to contribute a PR? (yes/no): no
• If yes, please read this page for instructions
• Briefly describe your candidate solution(if contributing):

Standalone code to reproduce the issue.

import tensorflow as tf
from tensorflow.keras import layers

class Sampling(layers.Layer):
    """Uses (z_mean, z_log_var) to sample z, the vector encoding a digit."""

    def call(self, inputs):
        z_mean, z_log_var = inputs
        batch = tf.shape(z_mean)[0]
        dim = tf.shape(z_mean)[1]
        epsilon = tf.keras.backend.random_normal(shape=(batch, dim))
        return z_mean + tf.exp(0.5 * z_log_var) * epsilon

class Encoder(layers.Layer):
    """Maps MNIST digits to a triplet (z_mean, z_log_var, z)."""

    def __init__(self, latent_dim=32, intermediate_dim=64, name="encoder", **kwargs):
        super(Encoder, self).__init__(name=name, **kwargs)
        self.dense_proj = layers.Dense(intermediate_dim, activation="relu")
        self.dense_mean = layers.Dense(latent_dim)
        self.dense_log_var = layers.Dense(latent_dim)
        self.sampling = Sampling()

    def call(self, inputs):
        x = self.dense_proj(inputs)
        z_mean = self.dense_mean(x)
        z_log_var = self.dense_log_var(x)
        z = self.sampling((z_mean, z_log_var))
        return z_mean, z_log_var, z

class Decoder(layers.Layer):
    """Converts z, the encoded digit vector, back into a readable digit."""

    def __init__(self, original_dim, intermediate_dim=64, name="decoder", **kwargs):
        super(Decoder, self).__init__(name=name, **kwargs)
        self.dense_proj = layers.Dense(intermediate_dim, activation="relu")
        self.dense_output = layers.Dense(original_dim, activation="sigmoid")

    def call(self, inputs):
        x = self.dense_proj(inputs)
        return self.dense_output(x)

class VariationalAutoEncoder(tf.keras.Model):
    """Combines the encoder and decoder into an end-to-end model for training."""

    def __init__(
        self,
        original_dim,
        intermediate_dim=64,
        latent_dim=32,
        name="autoencoder",
        **kwargs
    ):
        super(VariationalAutoEncoder, self).__init__(name=name, **kwargs)
        self.original_dim = original_dim
        self.encoder = Encoder(latent_dim=latent_dim, intermediate_dim=intermediate_dim)
        self.decoder = Decoder(original_dim, intermediate_dim=intermediate_dim)

    def call(self, inputs):
        z_mean, z_log_var, z = self.encoder(inputs)
        reconstructed = self.decoder(z)
        # Add KL divergence regularization loss.
        kl_loss = -0.5 * tf.reduce_mean(
            z_log_var - tf.square(z_mean) - tf.exp(z_log_var) + 1
        )
        self.add_loss(kl_loss)
        return reconstructed

vae = VariationalAutoEncoder(784, 64, 32)
inputs = tf.keras.Input(shape=(784))
vae(inputs)
optimizer = tf.keras.optimizers.Adam(learning_rate=1e-3)

(x_train, _), _ = tf.keras.datasets.mnist.load_data()
x_train = x_train.reshape(60000, 784).astype("float32") / 255
vae.compile(optimizer, loss=tf.keras.losses.MeanSquaredError())
vae.fit(x_train, x_train, epochs=1, batch_size=1000)
vae.save("VaeModel")
vae2 = tf.keras.models.load_model("VaeModel", compile=False)
print("-------------------------")
print(f"Original Model: {vae.output}")
print("-------------------------")
print(f"Loaded Model: {vae2.output}")

[tilakrayal]

@AshwinJay101,

When defining a network in Keras, the first layer added needs to have input_shape added and you should specify input shape for the first layer of the model.

Could you please refer the docs here: https://keras.io/getting-started/sequential-model-guide/#specifying-the-input-shape

Also for the MNIST, you should have something like input_shape=(28,28,1)

There is a example here for the reference: https://www.kaggle.com/adityaecdrid/mnist-with-keras-for-beginners-99457

[AshwinJay101]

This still does not solve the issue @tilakrayal

The subclass model actually has the output defined if you look at my code. Its after saving and loading that it does not retain those attributes

I tried adding input_shape and still get the same error

File "/usr/local/lib/python3.8/site-packages/keras/engine/base_layer.py", line 2096, in output
    raise AttributeError('Layer ' + self.name + ' has no inbound nodes.')
AttributeError: Layer autoencoder has no inbound nodes.

[AshwinJay101]

Also apparently after loading I need to add this for it to work

vae2(inputs)
print(f"Loaded Model: {vae2.output}")

This is not required for functional models but only for subclass models. Would still love for it to be supported by default for subclass models

[google-ml-butler[bot]]

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

[AshwinJay101]

Just commenting so that it doesnt get closed. I had provided the response

[tilakrayal]

@sachinprasadhs, I was able to reproduce the issue on tensorflow v2.11 and tf-nightly. Kindly find the gist of it here.