numpy 在Python中加速Metropolis-Hastings算法

mzillmmw  于 2023-10-19  发布在  Python
关注(0)|答案(2)|浏览(91)

我有一些使用MCMC对后验分布进行采样的代码,具体来说是Metropolis Hastings。我使用scipy生成随机样本:

import numpy as np
from scipy import stats

def get_samples(n):
    """
    Generate and return a randomly sampled posterior.

    For simplicity, Prior is fixed as Beta(a=2,b=5), Likelihood is fixed as Normal(0,2)

    :type n: int
    :param n: number of iterations

    :rtype: numpy.ndarray
    """
    x_t = stats.uniform(0,1).rvs() # initial value
    posterior = np.zeros((n,))
    for t in range(n):
        x_prime = stats.norm(loc=x_t).rvs() # candidate
        p1 = stats.beta(a=2,b=5).pdf(x_prime)*stats.norm(loc=0,scale=2).pdf(x_prime) # prior * likelihood 
        p2 = stats.beta(a=2,b=5).pdf(x_t)*stats.norm(loc=0,scale=2).pdf(x_t) # prior * likelihood 
        alpha = p1/p2 # ratio
        u = stats.uniform(0,1).rvs() # random uniform
        if u <= alpha:
            x_t = x_prime # accept
            posterior[t] = x_t
        elif u > alpha:
            x_t = x_t # reject
    posterior = posterior[np.where(posterior > 0)] # get rid of initial zeros that don't contribute to distribution
    return posterior

一般来说,我尽量避免在python中使用显式for循环-我会尝试使用纯numpy生成所有内容。然而,对于这种算法,带有if语句的for循环是不可避免的。因此,代码非常慢。当我分析我的代码时,它在for循环中花费了大部分时间(很明显),更具体地说,最慢的部分是生成随机数; stats.beta().pdf()stats.norm().pdf()
有时我使用numba来加速矩阵运算的代码。虽然numba与一些numpy操作兼容,但生成随机数不是其中之一。Numba有一个cuda rng,但这仅限于正态分布和均匀分布。
我的问题是,有没有一种方法可以使用与numba兼容的各种发行版的某种随机抽样来显着加速上面的代码?
我们不必把自己局限于numba,但它是我所知道的唯一易于使用的优化器。更一般地说,我正在寻找在python中的for循环**中加速各种分布(beta,gamma,poisson)**的随机采样的方法。

tgabmvqs

tgabmvqs1#

在你开始考虑numba et之前,你可以对这段代码进行很多优化。(我成功地在这个代码上获得了25倍的速度,只是因为我对算法的实现很聪明)
首先,你在执行大都会-黑斯廷斯算法时出错了.你需要保留方案的 * 每 * 次迭代,不管你的链是否移动。也就是说,您需要从代码中删除posterior = posterior[np.where(posterior > 0)],并在每个循环的末尾添加posterior[t] = x_t
其次,这个例子看起来很奇怪。通常,对于这类推理问题,我们希望在给定一些观察结果的情况下推断出分布的参数。然而,在这里,分布的参数是已知的,而不是你的抽样观察?不管怎样,不管这些,我很乐意用你的例子来告诉你如何加速它。

提速

首先,从主for循环中删除任何不依赖于t值的内容。首先从for循环中删除随机游走新息的生成:

x_t = stats.uniform(0,1).rvs()
    innov = stats.norm(loc=0).rvs(size=n)
    for t in range(n):
        x_prime = x_t + innov[t]

当然,也可以从for循环中移动u的随机生成:

x_t = stats.uniform(0,1).rvs()
    innov = stats.norm(loc=0).rvs(size=n)

    u = np.random.uniform(size=n)
    for t in range(n):
        x_prime = x_t + innov[t]
        ...
        if u[t] <= alpha:

另一个问题是,你在每个循环中计算当前的后验p2,这是不必要的。在每个循环中,你需要计算建议的后验p1,当建议被接受时,你可以更新p2等于p1

x_t = stats.uniform(0,1).rvs()
    innov = stats.norm(loc=0).rvs(size=n)

    u = np.random.uniform(size=n)

    p2 = stats.beta(a=2,b=5).pdf(x_t)*stats.norm(loc=0,scale=2).pdf(x_t)
    for t in range(n):
        x_prime = x_t + innov[t]

        p1 = stats.beta(a=2,b=5).pdf(x_prime)*stats.norm(loc=0,scale=2).pdf(x_prime)
        ...
        if u[t] <= alpha:
            x_t = x_prime # accept
            p2 = p1

        posterior[t] = x_t

一个非常小的改进是将scipy stats函数直接导入名称空间:

from scipy.stats import norm, beta

另一个非常小的改进是注意到代码中的elif语句没有做任何事情,因此可以删除。
把这些放在一起,使用我想出的更合理的变量名:

from scipy.stats import norm, beta
import numpy as np

def my_get_samples(n, sigma=1):

    x_cur = np.random.uniform()
    innov = norm.rvs(size=n, scale=sigma)
    u = np.random.uniform(size=n)

    post_cur = beta.pdf(x_cur, a=2, b=5) * norm.pdf(x_cur, loc=0, scale=2)

    posterior = np.zeros(n)
    for t in range(n):
        x_prop = x_cur + innov[t]

        post_prop = beta.pdf(x_prop, a=2, b=5) * norm.pdf(x_prop, loc=0, scale=2)
        alpha = post_prop / post_cur
        if u[t] <= alpha:
            x_cur = x_prop
            post_cur = post_prop

        posterior[t] = x_cur

    return posterior

现在,进行速度比较:

%timeit get_samples(1000)
3.19 s ± 5.28 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
%timeit my_get_samples(1000)
127 ms ± 484 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)

速度提升了25倍

ESS

值得注意的是,对于MCMC算法来说,蛮力速度并不是一切。实际上,你感兴趣的是每秒可以从后验中进行的独立(ish)绘制的数量。通常,这是用ESS (effective sample size)评估的。您可以通过调整随机游走来提高算法的效率(从而增加每秒提取的有效样本)。
为此,通常进行初始试运行,即,samples = my_get_samples(1000)。根据此输出计算sigma = 2.38**2 * np.var(samples)。然后,该值应用于将方案中的随机游走调整为innov = norm.rvs(size=n, scale=sigma)。2.38^2看似任意出现的原因是:
随机游走大都会算法的弱收敛和最优尺度(1997)。A. Gelman,W. R. Gilks和G. O. Roberts
为了说明调优可以带来的改进,让我们对我的算法进行两次运行,一次调优,另一次未调优,都使用了10000次迭代:

x = my_get_samples(10000)
y = my_get_samples(10000, sigma=0.12)

fig, ax = plt.subplots(1, 2)
ax[0].hist(x, density=True, bins=25, label='Untuned algorithm', color='C0')
ax[1].hist(y, density=True, bins=25, label='Tuned algorithm', color='C1')
ax[0].set_ylabel('density')
ax[0].set_xlabel('x'), ax[1].set_xlabel('x')
fig.legend()

您可以立即看到调优对算法效率的改进。请记住,两次运行的迭代次数相同。

最后的想法

如果你的算法需要很长时间才能收敛,或者你的样本有大量的自相关,我会考虑研究Cython来进一步优化速度。
我也建议你看看PyStan项目。它需要一点时间来适应,但它的NUTS HMC算法可能会优于任何大都会-黑斯廷斯算法,你可以手写。

eqfvzcg8

eqfvzcg82#

不幸的是,除了用numba兼容的python代码重写它们之外,我真的看不到任何加快随机分布的可能性。
但是,一个简单的方法可以加速代码的瓶颈,就是将stats函数的两次调用替换为:

p1, p2 = (
    stats.beta(a=2, b=5).pdf([x_prime, x_t])
    * stats.norm(loc=0, scale=2).pdf([x_prime, x_t]))

另一个轻微的调整可能是将u的生成外包到for循环之外:

x_t = stats.uniform(0, 1).rvs() # initial value
posterior = np.zeros((n,))
u = stats.uniform(0, 1).rvs(size=n) # random uniform
for t in range(n):  # and so on

然后在循环中索引u(当然循环中的行u = stats.uniform(0,1).rvs() # random uniform必须删除):

if u[t] <= alpha:
    x_t = x_prime # accept
    posterior[t] = x_t
elif u[t] > alpha:
    x_t = x_t # reject

微小的变化也可以通过省略elif语句来简化if条件,或者如果需要用于其他目的,则将其替换为else。但这只是一个小小的改进:

if u[t] <= alpha:
    x_t = x_prime # accept
    posterior[t] = x_t

编辑

另一个改进基于jwalton的回答:

def new_get_samples(n):
    """
    Generate and return a randomly sampled posterior.

    For simplicity, Prior is fixed as Beta(a=2,b=5), Likelihood is fixed as Normal(0,2)

    :type n: int
    :param n: number of iterations

    :rtype: numpy.ndarray
    """

    x_cur = np.random.uniform()
    innov = norm.rvs(size=n)
    x_prop = x_cur + innov
    u = np.random.uniform(size=n)

    post_cur = beta.pdf(x_cur, a=2,b=5) * norm.pdf(x_cur, loc=0,scale=2)
    post_prop = beta.pdf(x_prop, a=2,b=5) * norm.pdf(x_prop, loc=0,scale=2)

    posterior = np.zeros((n,))
    for t in range(n):        
        alpha = post_prop[t] / post_cur
        if u[t] <= alpha:
            x_cur = x_prop[t]
            post_cur = post_prop[t]
        posterior[t] = x_cur
    return posterior

随着时间的改进:

%timeit my_get_samples(1000)
# 187 ms ± 13 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)
%timeit my_get_samples2(1000)
# 1.55 ms ± 57.3 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)

这比jwalton的答案快了121倍。通过外包post_prop计算完成。
看了下直方图,似乎没什么问题。但是最好问问jwalton是否真的可以,他似乎对这个主题有更多的了解。:)

相关问题