帶有 TensorFlow Probability 的 Edward2 的簡單哈密頓蒙特卡羅示例

Question

愛德華的例子

由於Edward已棄用並且需要舊版本的 TensorFlow，因此可以為以下示例創建專用的虛擬環境

$ python3 --version
Python 3.6.8
$ python3 -m venv edward
$ source edward/bin/activate
(edward) $ pip3 install --upgrade pip setuptools wheel
(edward) $ cat edward.txt
tensorflow==1.7
edward~=1.3
scipy~=1.2
pandas~=0.24
matplotlib~=3.0
(edward) $ pip3 install -r edward.txt

我有一個非常簡單的最小工作示例，用於將哈密頓蒙特卡羅與愛德華一起使用，名為edward_old.py

#!/usr/bin/env python3

import numpy as np
import scipy.stats
import tensorflow as tf
import edward as ed
import pandas as pd
import matplotlib.pyplot as plt


def generate_samples(data, n_samples):
    # Pick initial point for MCMC chains based on the data
    low, med, high = np.percentile(data, (16, 50, 84))
    mu_init = np.float32(med)
    t_init = np.float32(np.log(0.5 * (high - low)))

    # Build a very simple model
    mu = ed.models.Uniform(-1.0, 1.0)
    t = ed.models.Uniform(*np.log((0.05, 1.0), dtype=np.float32))
    X = ed.models.Normal(
        loc=tf.fill(data.shape, mu), scale=tf.fill(data.shape, tf.exp(t))
    )

    # Emperical samples of a sclar
    q_mu = ed.models.Empirical(params=tf.Variable(tf.fill((n_samples,), mu_init)))
    q_t = ed.models.Empirical(params=tf.Variable(tf.fill((n_samples,), t_init)))

    # Run inference using HMC to generate samples.
    with tf.Session() as sess:
        inference = ed.HMC({mu: q_mu, t: q_t}, data={X: data})
        inference.run(step_size=0.01, n_steps=10)
        mu_samples, t_samples = sess.run([q_mu.params, q_t.params])
    return mu_samples, t_samples


def visualize(samples, mu_grid, sigma_grid):
    fig, ax = plt.subplots(1, 1, figsize=(6, 5))
    ax.scatter(samples['mu'], samples['sigma'], s=5, lw=0, c='black')
    ax.set_xlim(mu_grid[0], mu_grid[-1])
    ax.set_ylim(sigma_grid[0], sigma_grid[-1])
    ax.set_title('Edward')
    ax.set_xlabel('$\mu$')
    ax.set_ylabel('$\sigma$')
    plt.savefig('edward_old.pdf')


def main():
    np.random.seed(0)
    tf.set_random_seed(0)

    # Generate pseudodata from draws from a single normal distribution
    dist_mean = 0.0
    dist_std = 0.5
    n_events = 5000
    toy_data = scipy.stats.norm.rvs(dist_mean, dist_std, size=n_events)

    mu_samples, t_samples = generate_samples(toy_data, n_events)
    samples = pd.DataFrame({'mu': mu_samples, 'sigma': np.exp(t_samples)})

    n_grid = 50
    mu_grid = np.linspace(*np.percentile(mu_samples, (0.5, 99.5)), n_grid)
    sigma_grid = np.linspace(*np.exp(np.percentile(t_samples, (0.5, 99.5))), n_grid)
    visualize(samples, mu_grid, sigma_grid)


if __name__ == '__main__':
    main()

它通過產生下面的情節

(edward) $ python3 edward_old.py

Edward2 示例

但是，當我嘗試在以下環境中使用TensorFlow Probability和Edward2復制它時

$ python3 --version
Python 3.6.8
$ python3 -m venv tfp-edward2
$ source tfp-edward2/bin/activate
(tfp-edward2) $ pip3 install --upgrade pip setuptools wheel
(tfp-edward2) $ cat tfp-edward2.txt
tensorflow~=1.13
tensorflow-probability~=0.6
scipy~=1.2
pandas~=0.24
matplotlib~=3.0
(tfp-edward2) $ pip3 install -r tfp-edward2.txt

以及edward_old.py的generate_samples在一個名為edward2.py的文件中的以下更改

#!/usr/bin/env python3

import numpy as np
import scipy.stats
import tensorflow as tf
import tensorflow_probability as tfp
from tensorflow_probability import edward2 as ed
import pandas as pd
import matplotlib.pyplot as plt


def generate_samples(data, n_samples):
    # Pick initial point for MCMC chains based on the data
    low, med, high = np.percentile(data, (16, 50, 84))
    mu_init = np.float32(med)
    t_init = np.float32(np.log(0.5 * (high - low)))

    def model(data_shape):
        mu = ed.Uniform(
            low=tf.fill(data_shape, -1.0), high=tf.fill(data_shape, 1.0), name="mu"
        )
        t = ed.Uniform(
            low=tf.log(tf.fill(data_shape, 0.05)),
            high=tf.log(tf.fill(data_shape, 1.0)),
            name="t",
        )
        x = ed.Normal(loc=mu, scale=tf.exp(t), name="x")
        return x

    log_joint = ed.make_log_joint_fn(model)

    def target_log_prob_fn(mu, t):
        """Target log-probability as a function of states."""
        return log_joint(data.shape, mu=mu, t=t, x=data)

    step_size = tf.get_variable(
        name='step_size',
        initializer=0.01,
        use_resource=True,  # For TFE compatibility
        trainable=False,
    )

    num_burnin_steps = 1000

    hmc_kernel = tfp.mcmc.HamiltonianMonteCarlo(
        target_log_prob_fn=target_log_prob_fn,
        num_leapfrog_steps=5,
        step_size=step_size,
        step_size_update_fn=tfp.mcmc.make_simple_step_size_update_policy(
            num_adaptation_steps=int(num_burnin_steps * 0.8)
        ),
    )

    # How should these be done?
    q_mu = tf.random_normal(data.shape, mean=mu_init)
    q_t = tf.random_normal(data.shape, mean=t_init)

    states, kernel_results = tfp.mcmc.sample_chain(
        num_results=n_samples,
        current_state=[q_mu, q_t],
        kernel=hmc_kernel,
        num_burnin_steps=num_burnin_steps,
    )

    # Initialize all constructed variables.
    init_op = tf.global_variables_initializer()

    # Run the inference using HMC to generate samples
    with tf.Session() as sess:
        init_op.run()
        states_, results_ = sess.run([states, kernel_results])

    mu_samples, t_samples = states_[0][0], states_[1][0]
    return mu_samples, t_samples

運行

(tfp-edward2) $ python3 edward2.py

表明存在一些明顯的問題。 我不認為我正確地制定了ed.models.Empirical的等價物，因此，如果對此有任何想法或其他任何我做錯的事情，那就太好了。

我已經嘗試遵循“ 從 Edward 升級到 Edward2 ”示例，但我無法理解它們，無法從deep_exponential_family模型中使用的示例轉移到此示例。

Answer 1

我為自己創造的問題完全弄亂了我的分布的形狀。 起初我未能正確理解的是，我的tfp.mcmc.sample_chain的current_state應該是代表鏈初始位置的標量（ shape==() ）。 一旦我意識到這一點，那么很明顯，這些位置q_mu和q_t的形狀非常錯誤，應該是來自數據確定的位置的樣本的平均值

q_mu = tf.reduce_mean(tf.random_normal((1000,), mean=mu_init))
q_t = tf.reduce_mean(tf.random_normal((1000,), mean=t_init))

由於這些值是標量，因此我也一直在錯誤地創建模型的形狀。 我一直在創建與我的數據形狀相同的隨機變量樣本，錯誤地認為這只是將x的形狀移動到mu和t的形狀。 當然mu和t是來自它們各自均勻分布的標量隨機變量，然后是x的正態分布的參數， data.shape抽取data.shape樣本。

def model(data_shape):
    mu = ed.Uniform(low=-1.0, high=1.0, name="mu")
    t = ed.Uniform(low=tf.log(0.05), high=tf.log(1.0), name="t")
    x = ed.Normal(
        loc=tf.fill(data_shape, mu), scale=tf.fill(data_shape, tf.exp(t)), name="x"
    )
    return x

完成此操作后，唯一要做的就是現在正確訪問狀態

with tf.Session() as sess:
    init_op.run()
    states_, results_ = sess.run([states, kernel_results])
    mu_samples, t_samples = (states_[0], states_[1])

並產生下面的圖像

(tfp-edward2) $ python3 edward2.py

這與使用Edward的原始版本非常匹配。

完全更正的腳本如下

#!/usr/bin/env python3

import numpy as np
import scipy.stats
import tensorflow as tf
import tensorflow_probability as tfp
from tensorflow_probability import edward2 as ed
import pandas as pd
import matplotlib.pyplot as plt


def generate_samples(data, n_samples):
    # Pick initial point for MCMC chains based on the data
    low, med, high = np.percentile(data, (16, 50, 84))
    mu_init = np.float32(med)
    t_init = np.float32(np.log(0.5 * (high - low)))

    def model(data_shape):
        mu = ed.Uniform(low=-1.0, high=1.0, name="mu")
        t = ed.Uniform(low=tf.log(0.05), high=tf.log(1.0), name="t")
        x = ed.Normal(
            loc=tf.fill(data_shape, mu), scale=tf.fill(data_shape, tf.exp(t)), name="x"
        )
        return x

    log_joint = ed.make_log_joint_fn(model)

    def target_log_prob_fn(mu, t):
        """Target log-probability as a function of states."""
        return log_joint(data.shape, mu=mu, t=t, x=data)

    step_size = tf.get_variable(
        name='step_size',
        initializer=0.01,
        use_resource=True,  # For TFE compatibility
        trainable=False,
    )

    num_burnin_steps = 1000

    hmc_kernel = tfp.mcmc.HamiltonianMonteCarlo(
        target_log_prob_fn=target_log_prob_fn,
        num_leapfrog_steps=5,
        step_size=step_size,
        step_size_update_fn=tfp.mcmc.make_simple_step_size_update_policy(
            num_adaptation_steps=int(num_burnin_steps * 0.8)
        ),
    )

    # Initial states of chains
    q_mu = tf.reduce_mean(tf.random_normal((1000,), mean=mu_init))
    q_t = tf.reduce_mean(tf.random_normal((1000,), mean=t_init))

    states, kernel_results = tfp.mcmc.sample_chain(
        num_results=n_samples,
        current_state=[q_mu, q_t],
        kernel=hmc_kernel,
        num_burnin_steps=num_burnin_steps,
    )

    # Initialize all constructed variables.
    init_op = tf.global_variables_initializer()

    # Run the inference using HMC to generate samples
    with tf.Session() as sess:
        init_op.run()
        states_, results_ = sess.run([states, kernel_results])
        mu_samples, t_samples = (states_[0], states_[1])

    return mu_samples, t_samples


def visualize(samples, mu_grid, sigma_grid):
    fig, ax = plt.subplots(1, 1, figsize=(6, 5))
    ax.scatter(samples['mu'], samples['sigma'], s=5, lw=0, c='black')
    ax.set_xlim(mu_grid[0], mu_grid[-1])
    ax.set_ylim(sigma_grid[0], sigma_grid[-1])
    ax.set_title('tfp and Edward2')
    ax.set_xlabel('$\mu$')
    ax.set_ylabel('$\sigma$')
    plt.savefig('tfp-edward2.pdf')
    plt.savefig('tfp-edward2.png')


def main():
    np.random.seed(0)
    tf.set_random_seed(0)

    # Generate pseudodata from draws from a single normal distribution
    dist_mean = 0.0
    dist_std = 0.5
    n_events = 5000
    toy_data = scipy.stats.norm.rvs(dist_mean, dist_std, size=n_events)

    mu_samples, t_samples = generate_samples(toy_data, n_events)
    samples = pd.DataFrame({'mu': mu_samples, 'sigma': np.exp(t_samples)})

    n_grid = 50
    mu_grid = np.linspace(*np.percentile(mu_samples, (0.5, 99.5)), n_grid)
    sigma_grid = np.linspace(*np.exp(np.percentile(t_samples, (0.5, 99.5))), n_grid)
    visualize(samples, mu_grid, sigma_grid)


if __name__ == '__main__':
    main()