0
$\begingroup$

Consider the following minimal VAE:

import tensorflow as tf
import tensorflow_probability as tfp

tfk = tf.keras
tfkl = tf.keras.layers
tfpl = tfp.layers
tfd = tfp.distributions

#Fake dataset
cim = np.random.randint(2, size=(10,10))

#Parameters

vector_size = 10
input_shape = (vector_size,)
encoded_size = 3
latent_dim = 5

#Model

prior = tfd.Independent(tfd.Normal(loc=tf.zeros(encoded_size), scale=1), reinterpreted_batch_ndims=1)
vae = tfk.Sequential([
    #Encoder
    tfkl.InputLayer(input_shape=input_shape),    
    tfkl.Dense(
        tfpl.MultivariateNormalTriL.params_size(encoded_size),
        activation=None,
        use_bias = False
    ),
    tfpl.MultivariateNormalTriL(
        encoded_size,
        activity_regularizer=tfpl.KLDivergenceRegularizer(prior)
    ),
    #Decoder
    tfkl.Dense(
        units = vector_size, 
        activation = tf.nn.leaky_relu,
        use_bias = False
    )
])


vae.compile(optimizer=tf.optimizers.Adam(learning_rate=1e-3), loss='mse')
history = vae.fit(cim, cim, epochs = 1)

tf.saved_model.save(vae, './VAE')

The final line throws an error: AttributeError: 'Tensor' object has no attribute 'log_prob' Tracing the issue, it seems that the activity regularizer tfpl.KLDivergenceRegularizer(prior) is not being serialized corrected. Any ideas for how to save this Model with the regularizer intact?

$\endgroup$
0
$\begingroup$

One solution I found was saving the model in HD5 format. That seems to bypass the serialization of the regularizer: vae.save("./VAE.h5")

You can also try saving the weights only: vae.save_weights("./VAE-weights")

$\endgroup$
1
  • $\begingroup$ I cannot get the model to reload and continue training with HD5 , sadly. $\endgroup$ – Parker Wieck Oct 19 '20 at 16:50

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service, privacy policy and cookie policy

Not the answer you're looking for? Browse other questions tagged or ask your own question.