如何保存具有 TensorFlow 概率正则化器的模型?

数据挖掘 喀拉斯 张量流 自动编码器
2022-02-21 15:10:02

考虑以下最小 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')

最后一行抛出一个错误:AttributeError: 'Tensor' object has no attribute 'log_prob' 跟踪问题,似乎活动正则化tfpl.KLDivergenceRegularizer(prior)器没有被序列化更正。关于如何在正则化器完好无损的情况下保存此模型的任何想法?

1个回答

我发现的一种解决方案是以 HD5 格式保存模型。这似乎绕过了正则化器的序列化: vae.save("./VAE.h5")

您也可以尝试仅保存权重:vae.save_weights("./VAE-weights")