How to retrain a variational autoencoder and reproduce the same results

I'm using Autoencoders to learn a representation with tensorflow. The experiments have different autoencoders (Autoencoder, Variational Autoencoder (VAE), Multimodal Autoencoder, and Multimodal VAE). I want to save these autoencoders models for future reproducibility experiments or the use of these learned representations.

• My first strategy is to use the function "save_model" (from tensorflow), but for many large databases and varying different parameters on the autoencoders, this uses a lot of disk space (for example, 45G for 3 datasets, I will use 183).
• My second idea is to use "tensorflow.random.set_seed()", and when I fit the model again, it produces the same losses and results.

This works with autoencoder and multimodal autoencoder, but on VAE doesn't work. I think this is happing because of the random sampling on VAE (reparametrization trick).

• Is there any way to reproduce the same losses and results in VAE, without saving the model?
• Is there a way to use set_seed() in VAE?

Follow the code:

!gdown --id 1-Q7VXXkdGGxIYmIDuotlN3sXNVy-RtYd 

import pandas as pd
df = pd.read_parquet('CSTR.parquet')
df_train = np.array(df[df['class'] == 'Theory'].DistilBERT.to_list())
df_test = np.array(df[df['class'] != 'Theory'].DistilBERT.to_list())

import tensorflow
from tensorflow import keras
import numpy as np
from tensorflow.keras import layers
from tensorflow.keras.layers import Activation,Dense, Input
from tensorflow.keras.models import Model
import tensorflow as tf

def Autoencoder(X_train, encoding_dim, epoca):
    encoder_inputs = Input(shape=(X_train.shape[1], ), name='Entrada_Encoder') 

    encoded = Dense(encoding_dim, activation="linear")(encoder_inputs) 

    decoder_output = Dense(X_train.shape[1], activation="linear")(encoded) 

    autoencoder = Model(encoder_inputs, decoder_output)

    autoencoder.compile(optimizer=tensorflow.keras.optimizers.Adam(), loss='mse')

    autoencoder.fit(X_train, X_train, epochs=epoca, batch_size=32)

    return autoencoder

class Sampling(layers.Layer):
"""Uses (z_mean, z_log_var) to sample z."""

    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 VAE(keras.Model):
    def __init__(self, encoder, decoder, factor_multiply, **kwargs):
        super(VAE, self).__init__(**kwargs)
        self.encoder = encoder
        self.decoder = decoder
        self.factor_multiply = factor_multiply

    def train_step(self, data):
        if isinstance(data, tuple):
            data = data[0]
        with tf.GradientTape() as tape:
            z_mean, z_log_var, z = self.encoder(data)
            reconstruction = self.decoder(z)
            reconstruction_loss = tf.reduce_mean(
                keras.losses.mean_squared_error(data, reconstruction)
            )
            reconstruction_loss *= self.factor_multiply
            kl_loss = 1 + z_log_var - tf.square(z_mean) - tf.exp(z_log_var)
            kl_loss = tf.reduce_mean(kl_loss)
            kl_loss *= -0.5
            total_loss = reconstruction_loss + kl_loss

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

            return {
                "loss": total_loss,
                "reconstruction_loss": reconstruction_loss,
                "kl_loss": kl_loss,
            }

def encoder_vae(input_dim):
    encoder_inputs = keras.Input(shape=(input_dim,), name='encoder_input')

    z_mean = layers.Dense(128, name="Z_mean")(encoder_inputs)
    z_log_var = layers.Dense(128, name="Z_log_var")(encoder_inputs)
    z = Sampling()([z_mean, z_log_var])

    encoder = keras.Model([encoder_inputs], [z_mean, z_log_var, z], name="Encoder")

    return encoder


def decoder_vae(output_dim):
    latent_inputs = keras.Input(shape=(128,), name='decoder_input')

    decoder_outputs = Dense(output_dim, activation="linear")(latent_inputs)

    decoder = keras.Model(latent_inputs, decoder_outputs, name="decoder")

   return decoder


def variationalautoencoder(input_dim):
    encoder = encoder_vae(input_dim)

    decoder = decoder_vae(input_dim)

    vae = VAE(encoder, decoder, input_dim)

    vae.compile(optimizer=keras.optimizers.Adam())

    return vae, encoder, decoder

=================================================
tf.random.set_seed(1)

ae = Autoencoder(df_train, 128, 10)

X_test = ae.predict(df_test) # GENERATE THE SAME X_test AND LOSSES EVER
=======================================================================

tf.random.set_seed(1)

vae,encoder, decoder = variationalautoencoder(len(df_train[0]))

vae.fit(df_train, df_train, epochs=10, batch_size=32, verbose=1)

x_test, _, _ = encoder.predict(np.array(df_test))  # GENERATE DIFFERENT x_test 
AND LOSSES EVER
========================================================================