为大型数据集的 HMC 创建自定义梯度函数

Question

我正在尝试使用张量流概率中的 HMC 推断高斯过程的参数。

我有多个独立的数据序列，它们是从同一个底层进程生成的，我想推断出它们共享的内核参数。

为了计算可能性，我正在使用急切模式并循环遍历每个独立序列。我能够计算出可能性，但在尝试计算梯度时我运行陷入资源耗尽错误。

我知道这会很慢，但我希望能够将 HMC 与任何大小的数据集一起使用而不会运行内存不足。

我使用以下代码创建合成数据，这会从具有 p 个数据点的 GP 中创建 N 个样本。

L = 5
variance=2
m_noise=0.05

kernel=psd_kernels.ExponentiatedQuadratic(np.float64(variance), np.float64(L))


def gram_matrix(xs):
    return kernel.matrix(xs,xs).numpy() + m_noise*np.identity(xs.shape[0])


observation_index_points = []
observations = []
N=200
p = 2000
for i in range(0, N):

    xs = np.sort(np.random.uniform(0,100,p))[...,None]
    mean = [0 for x in xs]
    gram = gram_matrix(xs)

    ys = np.random.multivariate_normal(mean, gram)

    observation_index_points.append(xs)

    observations.append(ys)

for i in range(0, N):

    plt.plot(observation_index_points[i],observations[i])
plt.show()

以下计算对数似然的代码将运行采样器用于较小的 N 值，但在较大的 N 值（资源耗尽）时失败。尝试计算可能性的梯度时发生错误。

@tf.function()
def gp_log_prob(amplitude, length_scale, seg_index_points, noise_variance, seg_observations):

    kernel = psd_kernels.ExponentiatedQuadratic(amplitude, length_scale)
    gp = tfd.GaussianProcess(kernel=kernel,
                                index_points=seg_index_points,
                                observation_noise_variance=noise_variance)


    return gp.log_prob(seg_observations)

rv_amplitude = tfd.LogNormal(np.float64(0.), np.float64(1))
rv_length_scale = tfd.LogNormal(np.float64(0.), np.float64(1))
rv_noise_variance = tfd.LogNormal(np.float64(0.), np.float64(1))

def joint_log_prob_no_grad(amplitude, length_scale, noise_variance):
    ret_val = rv_amplitude.log_prob(amplitude) \
                + rv_length_scale.log_prob(length_scale) \
                + rv_noise_variance.log_prob(noise_variance)

    for i in range(N):
        ret_val = ret_val + gp_log_prob(amplitude, 
                                        length_scale, 
                                        observation_index_points[i], 
                                        noise_variance, 
                                        observations[i])

    return ret_val

但是我可以使用循环内的梯度带计算大 N 的梯度。此代码运行s 对于任何 N 和 returns 正确的可能性和梯度：

def joint_log_prob(amplitude, length_scale, noise_variance):

    with tf.GradientTape() as tape:
        tape.watch(amplitude)
        tape.watch(length_scale)
        tape.watch(noise_variance)

        ret_val = rv_amplitude.log_prob(amplitude) \
                    + rv_length_scale.log_prob(length_scale) \
                    + rv_noise_variance.log_prob(noise_variance)

    grads = tape.gradient(ret_val, [amplitude, length_scale, noise_variance])

    for i in range(N):
        with tf.GradientTape() as tape:
            tape.watch([amplitude, length_scale, noise_variance])
            gp_prob = gp_log_prob(amplitude, length_scale, 
                                  observation_index_points[i], noise_variance, observations[i])

        gp_grads = tape.gradient(gp_prob, [amplitude, length_scale, noise_variance])


        grads = [a+b for a,b in zip(grads,gp_grads)]
        ret_val = ret_val + gp_prob

    return ret_val, grads

x = tf.convert_to_tensor(np.float64(1.0))
y = tf.convert_to_tensor(np.float64(1.0))
z = tf.convert_to_tensor(np.float64(0.1))

joint_log_prob(x,y,z) # correct output even for large N

如果我将其转换为自定义渐变，它会再次失败：

@tf.custom_gradient
def joint_log_prob_cg(amplitude, length_scale, noise_variance):

    with tf.GradientTape() as tape:
        tape.watch(amplitude)
        tape.watch(length_scale)
        tape.watch(noise_variance)

        ret_val = rv_amplitude.log_prob(amplitude) \
                    + rv_length_scale.log_prob(length_scale) \
                    + rv_noise_variance.log_prob(noise_variance)

    grads = tape.gradient(ret_val, [amplitude, length_scale, noise_variance])

    for i in range(N):
        with tf.GradientTape() as tape:
            tape.watch([amplitude, length_scale, noise_variance])
            gp_prob = gp_log_prob(amplitude, length_scale, 
                                  observation_index_points[i], noise_variance, observations[i])

        gp_grads = tape.gradient(gp_prob, [amplitude, length_scale, noise_variance])


        grads = [a+b for a,b in zip(grads,gp_grads)]
        ret_val = ret_val + gp_prob

    def grad(dy):
        return grads
    return ret_val, grad

with tf.GradientTape() as t:
    t.watch([x,y,z])
    lp = joint_log_prob_cg(x,y,z)

t.gradient(lp, [x,y,z]) # fails for large N

我的问题是如何从上面的 joint_log_prob 函数（我知道可以为任何大型数据集计算）中获取 grads 到 HMC 采样器中？似乎如果整个函数都包含在 gradienttape 调用中，那么 for 循环就会展开并且它运行内存不足 - 但是有没有办法解决这个问题？

Answer 1

如果有人感兴趣，我可以通过使用自定义渐变并在 for 循环中停止录音来解决这个问题。我需要导入磁带实用程序：

from tensorflow.python.eager import tape

然后在 for 循环周围停止录制

with tape.stop_recording():
    for i in range(N):
        ...

这将停止跟踪，然后我必须在图形模式下计算梯度并手动添加它们以停止 OOM 错误。

为大型数据集的 HMC 创建自定义梯度函数

Creating a custom gradient function for HMC for large datasets

python-3.x

tensorflow

tensorflow-probability