During creating VAE model throws exception “you should implement a `call` method.”

I want to create VAE(variational autoencoder). During model creating it throws exception. When subclassing the Model class, you should implement a call method.

I am using Tensorflow 2.0

def vae():
    models ={}

    def apply_bn_and_dropout(x):
        return l.Dropout(dropout_rate)(l.BatchNormalization()(x))
    input_image = l.Input(batch_shape=(batch_size,28,28,1))
    x = l.Flatten()(input_image)
    x = l.Dense(256,activation="relu")(x)
    x = apply_bn_and_dropout(x)
    x = l.Dense(128,activation="relu")(x)
    x = apply_bn_and_dropout(x)

    z_mean = l.Dense(latent_dim)(x)
    z_log_var = l.Dense(latent_dim)(x)

    def sampling(args):
        z_mean, z_log_var = args
        epsilon = K.random_normal(shape=(batch_size,latent_dim),mean=0., stddev=1.0)
        return z_mean + K.exp(z_log_var/2) * epsilon

    lambda_layer = l.Lambda(sampling,output_shape=(latent_dim,))([z_mean,z_log_var])

    models["encoder"] = Model(input_image,lambda_layer,"Encoder")
    models["z_meaner"] = Model(input_image,z_mean,"Enc_z_mean")
    models["z_lvarer"] = Model(input_image, z_log_var,"Enc_z_log_var")

    z = l.Input(shape=(latent_dim,))
    x = l.Dense(128)(z)
    x = l.LeakyReLU()(x)
    x = apply_bn_and_dropout(x)
    x = l.Dense(256)(x)
    x = l.LeakyReLU()(x)
    x = apply_bn_and_dropout(x)
    x = l.Dense(28*28,activation="sigmoid")(x)
    decoded = l.Reshape((28,28,1))(x)

    models["decoder"] = Model(z,decoded,name="Decoder")
    models["vae"] = Model(input_image, models["decoder"](models["encoder"](input_image)), name="VAE")

    def vae_loss(x,decoded):
        x = K.reshape(x,shape=(batch_size,28*28))
        decoded = K.reshape(decoded,shape=(batch_size,28*28))
        xent_loss = 28*28*binary_crossentropy(x, decoded)
        kl_loss = -0.5 * K.sum(1 + z_log_var - K.square(z_mean) - K.exp(z_log_var), axis=-1)
        return (xent_loss + kl_loss)/2/28/28

    return models, vae_loss

---------------------------------------------------------------------------
NotImplementedError                       Traceback (most recent call last)
<ipython-input-34-186b31069dc3> in <module>
----> 1 models, vae_loss = vae()
      2 vae = models["vae"]

<ipython-input-33-0fa06b39e41c> in vae()
     36 
     37     models["decoder"] = Model(z,decoded,name="Decoder")
---> 38     models["vae"] = Model(input_image, models["decoder"](models["encoder"](input_image)), name="VAE")
     39 
     40     def vae_loss(x,decoded):

~AppDataLocalContinuumanaconda3libsite-packagestensorflowpythonkerasenginebase_layer.py in __call__(self, inputs, *args, **kwargs)
    610                       base_layer_utils.AutoAddUpdates(self,
    611                                                       inputs)) as auto_updater:
--> 612                 outputs = self.call(inputs, *args, **kwargs)
    613                 auto_updater.set_outputs(outputs)
    614 

~AppDataLocalContinuumanaconda3libsite-packagestensorflowpythonkerasenginenetwork.py in call(self, inputs, training, mask)
    865     """
    866     if not self._is_graph_network:
--> 867       raise NotImplementedError('When subclassing the `Model` class, you should'
    868                                 ' implement a `call` method.')
    869 

NotImplementedError: When subclassing the `Model` class, you should implement a `call` method.

Models with names

def create_dense_ae():
    encoding_dim = 64

    input_img = layers.Input(shape=(28, 28, 1))
    flat_img = layers.Flatten()(input_img)

    encoded = layers.Dense(encoding_dim, activation='relu')(flat_img)


    input_encoded = layers.Input(shape=(encoding_dim,))
    flat_decoded = layers.Dense(28*28, activation='sigmoid')(input_encoded)
    decoded = layers.Reshape((28, 28, 1))(flat_decoded)


    encoder = tf.keras.Model(input_img, encoded, name="encoder")
    decoder = tf.keras.Model(input_encoded, decoded, name="decoder")
    autoencoder = tf.keras.Model(input_img, decoder(encoder(input_img)), name="autoencoder")
    return encoder, decoder, autoencoder

I want to get model.

Answer

The problem is here:

models["encoder"] = Model(input_image,lambda_layer,"Encoder")
models["z_meaner"] = Model(input_image,z_mean,"Enc_z_mean")
models["z_lvarer"] = Model(input_image, z_log_var,"Enc_z_log_var")

You are passing three arguments to the construction, where only two are needed (inputs and outputs). Models do not have names. The problem is that three parameters will break the detection of network or sub-classed model as shown in the keras source code.

So just replace the code with:

models["encoder"] = Model(input_image,lambda_layer)
models["z_meaner"] = Model(input_image,z_mean)
models["z_lvarer"] = Model(input_image, z_log_var)