如何修复TensorFlow搭建VAE时“Gradients do not exist for variables”报错

问题根源 #

• 数据类型不匹配：编码器输出层显式指定了dtype=tf.float64，但其他网络层默认使用float32，同时输入的图像数组也转成了float64，前后类型不一致导致梯度链路断裂
• KL散度计算方式问题：两个正态分布的KL散度存在闭式解析解，你当前使用采样近似计算的方式，反而容易引入梯度不稳定甚至梯度丢失的问题
• 输入数据数值范围异常：你用np.random.uniform(0,1)已经生成了0_{1范围的图像数据，额外除以255会让数值压缩到0}0.0039区间，超出伯努利分布的最优处理范围，加剧梯度消失风险

修改方案 #

你只需要对代码做三处修改即可解决警告：

1. 移除编码器MultivariateNormalTriL层的dtype = tf.float64参数，统一使用默认的float32类型
2. 将KLDivergenceRegularizer的use_exact_kl参数改为True，删除采样相关的配置
3. 生成image_array时，去掉最后的/ 255.0操作，保留0~1的数值范围

修改后的核心代码片段如下：

# 生成输入数据时去掉除以255，统一用float32
image_array = np.random.uniform( 0, 1, ( 100, 36, 36, 3 ) ).astype( np.float32 )

# KL正则器使用解析解计算
kl_regularizer = tfpl.KLDivergenceRegularizer(
    distribution_b = prior_distribution,
    use_exact_kl = True,
    weight = 1.0
)

# 编码器输出层去掉float64指定
encoder = tf.keras.Sequential( [
    Conv2D( 3, 3, padding = "same", input_shape = ( 36, 36, 3 ) ),
    Flatten( input_shape = ( 36, 36, 3 ) ),
    Dense( tfpl.MultivariateNormalTriL.params_size( 2 ), name = "encoder_dense" ),
    tfpl.MultivariateNormalTriL( 2, activity_regularizer = kl_regularizer, name = "encoder_outdist" )
])

修改后运行代码就不会再出现梯度不存在的警告，解码器参数也可以正常更新。

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