TensorFlow概率库VAE保存报错及解码器单独加载方案咨询

你在编码器最后一层MultivariateNormalTriL中配置的KLDivergenceRegularizer依赖外部定义的先验分布prior，模型序列化过程中无法正确捕获该分布对象，导致其被误转换为Tensor类型，触发log_prob属性不存在的报错。

• 整体保存加载思路是正确的，只要解决TFP自定义分布层的序列化依赖问题即可正常使用。
• 单独保存解码器的方案完全合理，解码器本身不依赖KL正则项和先验分布，序列化逻辑更简单，加载后可直接用于隐向量生成样本的场景。

方案1：SavedModel直接保存（推荐） #

训练完成后先移除推理阶段不需要的KL正则项，再执行保存，加载时指定TFP自定义层映射即可：

保存代码 #

# 移除编码器的KL正则（仅推理时不需要，不影响模型权重）
encoder.layers[-1].activity_regularizer = None

# 保存完整VAE
vae.save('saved_vae')
# 单独保存解码器
decoder.save('saved_decoder')

加载代码 #

import tensorflow as tf
import tensorflow_probability as tfp
tfd = tfp.distributions
tfpl = tfp.layers

# 声明自定义层映射，加载时必须传入
custom_objects = {
    'MultivariateNormalTriL': tfpl.MultivariateNormalTriL,
    'IndependentBernoulli': tfpl.IndependentBernoulli
}

# 加载完整VAE验证重构
vae_rec = tf.keras.models.load_model('saved_vae', custom_objects=custom_objects)
x = next(iter(eval_dataset))[0][:10]
xhat = vae_rec(x)
assert isinstance(xhat, tfd.Distribution)

# 加载解码器验证生成
decoder_rec = tf.keras.models.load_model('saved_decoder', custom_objects=custom_objects)
# 先定义和训练时完全一致的先验分布
prior = tfd.Independent(tfd.Normal(loc=tf.zeros(encoded_size), scale=1), reinterpreted_batch_ndims=1)
z = prior.sample(10)
xtilde = decoder_rec(z)
assert isinstance(xtilde, tfd.Distribution)

方案2：权重分离保存（兼容性最优） #

如果不想修改原模型配置，可直接保存权重，加载时先初始化完全相同的模型结构再载入权重，完全规避序列化问题：

# 保存权重
vae.save_weights('vae_weights.h5')
decoder.save_weights('decoder_weights.h5')

# 加载权重：先复制训练阶段的encoder、decoder、vae初始化代码，再执行
vae.load_weights('vae_weights.h5')
decoder.load_weights('decoder_weights.h5')

内容的提问来源于stack exchange，提问作者kiriloff