Cannot save a Custom Variational Autoencoder Model built in Keras with tensor flow

Question

Keras 2.10.0

I am trying to save my Variational Autoencoder built in Keras and Tensorflow once it is trained. I am not being able to do so. What can I do to save the model?

Attached is the Colab link to my running base model.

https://colab.research.google.com/github/MaxIG1/var_auto_for_synth_data/blob/saving_branch/notebooks/var_auto_working.ipynb

The error message it creates is here

ValueError: Model <__main__.VAE object at 0x00000278F0A96E60> cannot be saved either because the input shape is not available or because the forward pass of the model is not defined.To define a forward pass, please override Model.call(). To specify an input shape, either call build(input_shape) directly, or call the model on actual data using Model(), Model.fit(), or Model.predict(). If you have a custom training step, please make sure to invoke the forward pass in train step through Model.__call__, i.e. model(inputs), as opposed to model.call().

Furthermore I add the main lines of code here if you do not wish to open Colab.

Any help would be greatly appreciated.

import numpy as np
from tensorflow import keras
from keras.layers import Input, Dense, Lambda
import tensorflow as tf
import pandas as pd
import matplotlib.pyplot as plt
from sklearn.datasets import make_classification
from keras import backend as K
import matplotlib.pyplot as plt
import numpy as np



# Creating fake data to replicate
# Define the desired number of samples for each class
class_samples = [2000, 3000, 2500, 1500]  # Adjust these numbers as needed

# Calculate weights based on the desired number of samples
class_weights = [num_samples / sum(class_samples) for num_samples in class_samples]

# Generate a synthetic dataset with different numbers of samples for each class
X, y = make_classification(
    n_samples=sum(class_samples),
    n_features=4,
    n_informative=4,
    n_redundant=0,
    n_classes=4,
    weights=class_weights,
    random_state=42,
)
columns = ["Feature_1", "Feature_2", "Feature_3", "Feature_4"]
synthetic_df = pd.DataFrame(data=X, columns=columns)
for column in synthetic_df:
    std = np.std(synthetic_df[column])
    mean = np.mean(synthetic_df[column])
    synthetic_df[column] = synthetic_df[column]-mean
    synthetic_df[column] = synthetic_df[column]/std
synthetic_df["target"] = y
synthetic_array =synthetic_df.values

# Defining the sampling layer that is also the call
class Sampling(keras.layers.Layer):

    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

# Encoder
latent_dim = 10
encoder_inputs = Input(shape=(5), name="input_layer")

n_x = 50
x = Dense(n_x, activation="relu", name="h1")(encoder_inputs)

# Split x into two halves
half_size = n_x // 2
x3_first_half = Lambda(lambda x: x[:, :half_size], name="select_z_mean")(x)
x3_second_half = Lambda(lambda x: x[:, half_size:], name="select_z_var")(x)

z_mean = Dense(latent_dim, name="z_mean")(x3_first_half)
z_log_var = Dense(latent_dim, name="z_log_var")(x3_second_half)
z = Sampling(name="Sampling")([z_mean, z_log_var])

encoder = keras.Model(encoder_inputs, [z_mean, z_log_var, z], name="encoder")
# Decoder
# Decoder
latent_inputs = keras.Input(shape=(latent_dim,))
n_x = 30

x = Dense(n_x, activation="relu", name="h4")(latent_inputs)

cont_decoder_outputs = Dense(4, activation="linear", name="cont_decoder_output")(x)
class_decoder_output = Dense(4, activation="softmax", name="classification_output")(x)

decoder = keras.Model(latent_inputs, [cont_decoder_outputs, class_decoder_output], name="decoder")
decoder.summary()
# Custom VAE
class VAE(keras.Model):
    def __init__(self, encoder, decoder, **kwargs):
        super().__init__(**kwargs)
        self.encoder = encoder
        self.decoder = decoder
        self.total_loss_tracker = keras.metrics.Mean(name="total_loss")

    @property
    def metrics(self):
        return [
            self.total_loss_tracker,
        ]

    def train_step(self, data):
        with tf.GradientTape() as tape:
            z_mean, z_log_var, z = self.encoder(data)
            reconstruction_cont, reconstruction_class = self.decoder(z)

            data_cont = data[
                :, :4
            ]  # Assuming the first 4 columns are for continuous variables
            data_class = data[:, 4:]  # Assuming the last column is for classification

            # Reconstruction loss for continuous outputs
            reconstruction_loss_cont = keras.losses.mean_squared_error(
                data_cont, reconstruction_cont
            )

            # Reconstruction loss for classification output
            reconstruction_loss_class = keras.losses.sparse_categorical_crossentropy(
                data_class, reconstruction_class
            )

            kl_loss = -0.5 * (1 + z_log_var - tf.square(z_mean) - tf.exp(z_log_var))
            kl_loss = tf.reduce_mean(kl_loss, axis=1)

            # Combine losses
            total_loss = (
                0.8*reconstruction_loss_cont + 0.2*reconstruction_loss_class + kl_loss
            )

        grads = tape.gradient(total_loss, self.trainable_weights)
        self.optimizer.apply_gradients(zip(grads, self.trainable_weights))
        self.total_loss_tracker.update_state(total_loss)

        return {
            "loss": self.total_loss_tracker.result(),
        }

    def get_config(self):
        config = super(VAE, self).get_config()
        config.update(
            {
                "encoder": self.encoder,
                "decoder": self.decoder,
            }
        )
        return config

# Compile and Train
vae = VAE(encoder, decoder)
vae.compile(optimizer=keras.optimizers.Adam())
vae.fit(
    synthetic_array,
    epochs=1,
    batch_size=16,
)

# Save in TensorFlow SavedModel format
vae.save("path_to_save_model_tf")

Answer 1

The error message tells you what to do, although it's a bit awkward in this case.

First you have to define a call function for your model. Add this to the VAE class:

def call(self, inputs):
    return self.decoder(self.encoder(inputs)[2])

This allows Keras to figure out what this model is actually supposed to do overall. It is not enough to call fit, because you overwrote train_step and it never actually calls the vae model itself, only the components (encoder and decoder).

Thus we still have to actually call the model once, so it has a defined input shape. Sprinkle this line anywhere after you defined the model, but before you try to save it:

vae(synthetic_array[:1])

This just calls the model on a single data point to define input/output shapes. After that, both saving and then loading the saved model works fine. Note that I tested this on Colab, which uses the more recent 2.15 version. I think they made some changes to saving models in the meantime, so this might not work properly on an old version.