如何重新训练变分自编码器并重现相同的结果

Question

我正在使用自动编码器来学习 tensorflow 的表示。实验有不同的自动编码器（自动编码器、变分自动编码器 (VAE)、多模态自动编码器和多模态 VAE）。我想保存这些自动编码器模型，以供将来进行可重复性实验或使用这些学习表示。

• 我的第一个策略是使用函数“save_model”（来自 tensorflow），但是对于许多大型数据库和自动编码器上不同的参数，这会占用大量磁盘空间（例如，45G 用于 3 个数据集，我将使用 183)。
• 我的第二个想法是使用“tensorflow.random.set_seed()”，当我再次拟合模型时，它会产生相同的损失和结果。

这适用于自动编码器和多模式自动编码器，但在 VAE 上不起作用。我认为这是因为 VAE 上的随机抽样（重新参数化技巧）而发生的。

• 有什么方法可以在不保存模型的情况下在 VAE 中重现相同的损失和结果？
• 有没有办法在 VAE 中使用 set_seed()？

关注代码：

!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
========================================================================

Answer 1

正如我所怀疑的，采样函数行：

epsilon = tf.keras.backend.random_normal(shape=(batch, dim))

破坏了可重复性，所以使用函数的种子参数

epsilon = tf.keras.backend.random_normal(shape=(batch, dim), seed=1)

，可以使用

tf.random.set_seed(1)