python - 使用 TF2 最小化 KL 散度

Question

我正在尝试更熟悉 TensorFlow 2，因此我正在尝试一些练习。我正在做一个以最小化两个高斯之间的 KL 散度并更新一个相对于另一个的参数。

下面是运行良好的代码，但即使 KL_divergence 具有有限值，梯度也被计算为 NaN 的均值和方差：

kl_divergence:  1947.9133
grads: [<tf.Tensor: shape=(1,), dtype=float32, numpy=array([nan], dtype=float32)>, <tf.Tensor: shape=(1,), dtype=float32, numpy=array([nan], dtype=float32)>]

你知道这个错误可能来自哪里吗？我虽然tf.where无法传播渐变，但看起来它实际上可以。

import numpy as np
from scipy.special import kl_div
from scipy.stats import norm
import matplotlib.pyplot as plt
import tensorflow as tf
import tensorflow.keras as Keras

x = np.arange(-10, 10, 0.001)
p = norm.pdf(x, 0, 5)

mean = tf.Variable(initial_value=tf.zeros((1,)), trainable=True, dtype=tf.float32, name='mean')
var = tf.Variable(initial_value=tf.ones((1,)), trainable=True, dtype=tf.float32, name='var')
optimizer = Keras.optimizers.SGD(0.1)
EPS = 1e-5

for _ in range(50):
    with tf.GradientTape() as tape:
        normal = tf.exp(-tf.square(x - mean / (2 * var)))
        q = normal / tf.reduce_sum(normal)
        kl_divergence = tf.reduce_sum(
            tf.where(q < EPS, tf.zeros(p.shape, tf.float32), p * tf.math.log(p / q))
        )

    print('kl_divergence: ', kl_divergence.numpy())
    grads = tape.gradient(kl_divergence, [mean, var])
    print(grads)
    # optimizer.minimize(kl_divergence, [mean, var], tape=tape)
    # print('mean: ', mean.numpy())
    # print('var: ', var.numpy())

标签： pythontensorflowprobability